
import pandas as pd
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import log_loss
from .config import Config

class FeatureSelector:
    def __init__(self):
        self.selected_features = []
        self.dropped_features = {}
        
    def filter_by_iv(self, iv_df, threshold=0.02):
        """
        Filter features based on Information Value.
        iv_df: DataFrame with 'variable' and 'iv' columns.
        """
        selected = iv_df[iv_df['iv'] >= threshold]['variable'].tolist()
        dropped = iv_df[iv_df['iv'] < threshold]['variable'].tolist()
        
        self.dropped_features['iv_filter'] = dropped
        return selected
        
    def filter_by_corr(self, df, features=None, threshold=0.7):
        """
        Filter features based on correlation.
        Keeps the one with higher IV if iv_df provided (not implemented here for simplicity, just first one).
        """
        if features is None:
            features = df.columns.tolist()
            
        df_sub = df[features]
        corr_matrix = df_sub.corr().abs()
        upper = corr_matrix.where(np.triu(np.ones(corr_matrix.shape), k=1).astype(bool))
        
        to_drop = [column for column in upper.columns if any(upper[column] > threshold)]
        
        selected = [f for f in features if f not in to_drop]
        self.dropped_features['corr_filter'] = to_drop
        return selected

    def stepwise_selection(self, X, y, initial_features=None, method='forward', criterion='aic'):
        """
        Stepwise selection using Logistic Regression and AIC/BIC.
        """
        if initial_features is None:
            remaining = list(X.columns)
            selected = []
        else:
            remaining = [f for f in X.columns if f not in initial_features]
            selected = list(initial_features)
            
        current_score = np.inf
        best_new_score = np.inf
        
        # Initial score
        if selected:
            current_score = self._calculate_score(X[selected], y, criterion)
        else:
            # Null model score (intercept only)
            # We simulate intercept only by passing a column of ones if using fit_intercept=False
            # But sklearn LR adds intercept by default.
            # So passing empty X is tricky.
            # We calculate score for null model manually or using DummyClassifier
            # Or just start with inf and let the first feature be added.
            null_lr = LogisticRegression(penalty=None, solver='lbfgs')
            # Fake feature of 1s
            null_lr.fit(np.ones((len(y), 1)), y)
            probs = null_lr.predict_proba(np.ones((len(y), 1)))
            ll = log_loss(y, probs)
            k = 1 # intercept
            n = len(y)
            if criterion == 'aic':
                current_score = 2*k + 2*n*ll
            else:
                current_score = np.log(n)*k + 2*n*ll
        
        while remaining:
            scores_with_candidates = []
            
            for candidate in remaining:
                features_to_test = selected + [candidate]
                score = self._calculate_score(X[features_to_test], y, criterion)
                scores_with_candidates.append((score, candidate))
                
            scores_with_candidates.sort()
            best_new_score, best_candidate = scores_with_candidates[0]
            
            if best_new_score < current_score:
                remaining.remove(best_candidate)
                selected.append(best_candidate)
                current_score = best_new_score
                # print(f"Added {best_candidate}, score: {current_score}")
            else:
                break
                
        self.selected_features = selected
        return selected

    def _calculate_score(self, X, y, criterion='aic'):
        """
        Calculate AIC/BIC for Logistic Regression.
        """
        # Use simple model without penalty for stepwise to be faster and standard
        # However, if features are correlated, it might fail.
        # Adding a small L2 regularization helps convergence.
        lr = LogisticRegression(penalty='l2', C=1e9, solver='lbfgs', max_iter=1000)
        try:
            lr.fit(X, y)
            probs = lr.predict_proba(X)
            ll = log_loss(y, probs)
            
            k = X.shape[1] + 1 # +1 for intercept
            n = len(y)
            
            if criterion == 'aic':
                return 2*k + 2*n*ll
            elif criterion == 'bic':
                return np.log(n)*k + 2*n*ll
        except:
            return np.inf
