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

class FeatureSelector:
    def __init__(self):
        self.selected_features = []
        self.dropped_features = {}
        self.selection_process = []  # 记录特征选择过程
        
    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
        
        # 记录特征选择过程
        self.selection_process.append({
            'step': 'IV Filter',
            'description': f'Filter features with IV < {threshold}',
            'initial_count': len(iv_df),
            'selected_count': len(selected),
            'dropped_features': 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 SGD Classifier (logistic regression equivalent) 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
        
        # 记录初始状态
        self.selection_process.append({
            'step': 'Initial State',
            'description': 'Initial feature set before stepwise selection',
            'initial_count': len(X.columns),
            'selected_count': len(selected),
            'selected_features': selected.copy()
        })
        
        # 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_sgd = SGDClassifier(loss='log_loss', penalty='l2', alpha=1e-4, max_iter=1000)
            # Fake feature of 1s
            null_sgd.fit(np.ones((len(y), 1)), y)
            probs = null_sgd.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
        
        iteration = 0
        while remaining:
            iteration += 1
            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
                
                # 记录每一步的选择过程
                self.selection_process.append({
                    'step': f'Stepwise Selection - Iteration {iteration}',
                    'description': f'Added feature: {best_candidate}, Score: {current_score:.4f}',
                    'initial_count': len(X.columns),
                    'selected_count': len(selected),
                    'selected_features': selected.copy(),
                    'added_feature': best_candidate,
                    'score': current_score
                })
            else:
                # 记录终止原因
                self.selection_process.append({
                    'step': f'Stepwise Selection - Iteration {iteration}',
                    'description': f'No improvement, stopped. Best score: {current_score:.4f}',
                    'initial_count': len(X.columns),
                    'selected_count': len(selected),
                    'selected_features': selected.copy(),
                    'stopped': True
                })
                break
                
        self.selected_features = selected
        return selected
        
    def _calculate_score(self, X, y, criterion='aic'):
        """
        Calculate AIC/BIC for Logistic Regression using SGD Classifier.
        """
        # Use SGD Classifier with log_loss (equivalent to Logistic Regression)
        # Using L2 regularization with small alpha to avoid convergence issues
        sgd = SGDClassifier(
            loss='log_loss',
            penalty='l2',
            alpha=1e-4,
            max_iter=1000,
            tol=1e-3,
            random_state=42
        )
        try:
            sgd.fit(X, y)
            probs = sgd.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