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_missing_rate(self, df, threshold=0.95):
        """
        Filter features based on missing rate.
        df: DataFrame with features
        threshold: maximum missing rate allowed (default 0.95, i.e., 95%)
        """
        missing_rates = df.isnull().sum() / len(df)
        selected = missing_rates[missing_rates <= threshold].index.tolist()
        dropped = missing_rates[missing_rates > threshold].index.tolist()
        
        self.dropped_features['missing_rate_filter'] = dropped
        
        # 记录特征选择过程
        self.selection_process.append({
            'step': 'Missing Rate Filter',
            'description': f'Filter features with missing rate > {threshold:.2%}',
            'initial_count': len(df.columns),
            'selected_count': len(selected),
            'dropped_features': dropped
        })
        
        return selected
        
    def filter_by_single_value_rate(self, df, threshold=0.99):
        """
        Filter features based on single value rate.
        df: DataFrame with features
        threshold: maximum single value rate allowed (default 0.99, i.e., 99%)
        """
        single_value_rates = {}
        for col in df.columns:
            nunique = df[col].nunique()
            if nunique == 1:
                single_value_rates[col] = 1.0
            else:
                single_value_rates[col] = df[col].value_counts().max() / len(df)
        
        selected = [col for col, rate in single_value_rates.items() if rate <= threshold]
        dropped = [col for col, rate in single_value_rates.items() if rate > threshold]
        
        self.dropped_features['single_value_rate_filter'] = dropped
        
        # 记录特征选择过程
        self.selection_process.append({
            'step': 'Single Value Rate Filter',
            'description': f'Filter features with single value rate > {threshold:.2%}',
            'initial_count': len(df.columns),
            'selected_count': len(selected),
            'dropped_features': dropped
        })
        
        return selected
        
    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()
            
        # Filter numeric features only for correlation calculation
        numeric_features = [f for f in features if pd.api.types.is_numeric_dtype(df[f])]
        
        if len(numeric_features) == 0:
            return features
        
        df_sub = df[numeric_features]
        corr_matrix = df_sub.corr().abs()
        upper = corr_matrix.where(np.triu(np.ones(corr_matrix.shape), k=1).astype(bool))
        
        to_drop = []
        for column in upper.columns:
            if any(upper[column] > threshold):
                to_drop.append(column)
        
        selected = [f for f in features if f not in to_drop]
        self.dropped_features['corr_filter'] = to_drop
        return selected
        
    def filter_by_vif(self, X, features=None, threshold=10):
        """
        Filter features based on VIF (Variance Inflation Factor).
        X: DataFrame with features
        features: list of features to consider (default: all columns)
        threshold: maximum VIF allowed (default 10)
        """
        if features is None:
            features = X.columns.tolist()
        
        # Filter numeric features only
        numeric_features = [f for f in features if pd.api.types.is_numeric_dtype(X[f])]
        
        if len(numeric_features) == 0:
            return features
        
        # Iteratively remove features with high VIF
        selected = numeric_features.copy()
        dropped = []
        
        while True:
            # Calculate VIF for each feature
            vif_dict = {}
            for feat in selected:
                # VIF = 1 / (1 - R²)
                # Use linear regression to get R²
                other_feats = [f for f in selected if f != feat]
                if len(other_feats) == 0:
                    vif_dict[feat] = 1.0
                else:
                    try:
                        from sklearn.linear_model import LinearRegression
                        lr = LinearRegression()
                        lr.fit(X[other_feats], X[feat])
                        r2 = lr.score(X[other_feats], X[feat])
                        if r2 >= 1:
                            vif_dict[feat] = np.inf
                        else:
                            vif_dict[feat] = 1 / (1 - r2)
                    except:
                        vif_dict[feat] = np.inf
            
            # Find feature with highest VIF
            max_vif_feat = max(vif_dict, key=vif_dict.get)
            max_vif = vif_dict[max_vif_feat]
            
            if max_vif > threshold:
                selected.remove(max_vif_feat)
                dropped.append(max_vif_feat)
            else:
                break
        
        # Add back non-numeric features
        non_numeric_features = [f for f in features if f not in numeric_features]
        selected = selected + non_numeric_features
        
        self.dropped_features['vif_filter'] = dropped
        
        # 记录特征选择过程
        self.selection_process.append({
            'step': 'VIF Filter',
            'description': f'Filter features with VIF > {threshold}',
            'initial_count': len(features),
            'selected_count': len(selected),
            'dropped_features': dropped
        })
        
        return selected
        
    def stepwise_selection(self, X, y, initial_features=None, method='forward', criterion='aic', sl=0.05):
        """
        Stepwise selection using Logistic Regression.
        Based on p-value and AIC/BIC.
        
        Parameters:
        -----------
        X : DataFrame
            Independent variables
        y : Series
            Dependent variable
        initial_features : list, optional
            Initial features to start with
        method : str, default='forward'
            Selection method ('forward' or 'backward')
        criterion : str, default='aic'
            Selection criterion ('aic', 'bic', 'r2', 'adjr2')
        sl : float, default=0.05
            Significance level for p-value test
        """
        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)
            from sklearn.linear_model import LogisticRegression
            lr_null = LogisticRegression(
                penalty='l2', C=1.0, solver='lbfgs', 
                max_iter=1000, random_state=42
            )
            # Fit with dummy feature
            lr_null.fit(np.ones((len(y), 1)), y)
            probs = lr_null.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
            
            # Step 1: Calculate p-values for remaining features
            pvals = []
            for candidate in remaining:
                features_to_test = selected + [candidate]
                pval = self._calculate_pvalue(X[features_to_test], y, candidate)
                pvals.append((pval, candidate))
            
            # Step 2: Filter by significance level
            significant_candidates = [(p, f) for p, f in pvals if p <= sl]
            
            if not significant_candidates:
                # No significant candidates, stop
                self.selection_process.append({
                    'step': f'Stepwise Selection - Iteration {iteration}',
                    'description': f'No significant candidates (p-value <= {sl}), stopped. Best score: {current_score:.4f}',
                    'initial_count': len(X.columns),
                    'selected_count': len(selected),
                    'selected_features': selected.copy(),
                    'stopped': True
                })
                break
            
            # Step 3: For significant candidates, find the one that improves score the most
            scores_with_candidates = []
            for pval, candidate in significant_candidates:
                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_pvalue(self, X, y, feature):
        """
        Calculate p-value for a feature using Logistic Regression.
        """
        from sklearn.linear_model import LogisticRegression
        from scipy import stats
        
        lr = LogisticRegression(
            penalty='l2', C=1.0, solver='lbfgs',
            max_iter=1000, random_state=42
        )
        
        try:
            lr.fit(X, y)
            
            # Get coefficients
            coef = lr.coef_[0]
            
            # Calculate standard errors using Fisher information
            probs = lr.predict_proba(X)
            W = np.diag(probs[:, 0] * probs[:, 1])
            
            # Add intercept column to X
            X_with_intercept = np.column_stack([np.ones(X.shape[0]), X])
            
            # Calculate covariance matrix
            try:
                cov = np.linalg.inv(X_with_intercept.T @ np.diag(W) @ X_with_intercept)
                std_err = np.sqrt(np.diag(cov))
                
                # Get p-value for last feature
                feature_idx = X.shape[1] - 1
                z_score = coef[feature_idx] / std_err[feature_idx + 1]  # +1 for intercept
                p_value = 2 * (1 - stats.norm.cdf(abs(z_score)))
                
                return p_value
            except np.linalg.LinAlgError:
                # If matrix is singular, return high p-value
                return 1.0
        except:
            return 1.0  # Return high p-value if calculation fails
    
    def _calculate_score(self, X, y, criterion='aic'):
        """
        Calculate AIC/BIC for Logistic Regression using sklearn.
        """
        from sklearn.linear_model import LogisticRegression
        from sklearn.metrics import log_loss
        
        # Use Logistic Regression with L2 regularization
        lr = LogisticRegression(
            penalty='l2',
            C=1.0,
            solver='lbfgs',
            max_iter=1000,
            random_state=42,
            fit_intercept=True
        )
        
        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
            else:
                return -ll  # For minimization, use negative log-likelihood
        except:
            return np.inf
    
    def backward_selection_by_iv(self, X, y, iv_df, auc_threshold=0.005):
        """
        Backward selection based on IV order.
        1. Sort features by IV (ascending).
        2. Iteratively remove features.
        3. Check model performance (AUC) change.
        4. Retain feature if removal causes significant drop, otherwise drop.
        """
        from sklearn.linear_model import LogisticRegression
        from sklearn.metrics import roc_auc_score
        
        # Sort features by IV ascending (remove low IV first)
        # Assuming iv_df has 'variable' and 'iv' columns
        sorted_features = iv_df.sort_values('iv', ascending=True)['variable'].tolist()
        
        # Keep only features present in X
        sorted_features = [f for f in sorted_features if f in X.columns]
        
        current_features = sorted_features.copy()
        
        # Initial AUC
        lr = LogisticRegression(solver='lbfgs', max_iter=1000, random_state=42)
        lr.fit(X[current_features], y)
        probs = lr.predict_proba(X[current_features])[:, 1]
        current_auc = roc_auc_score(y, probs)
        
        print(f"Initial AUC with {len(current_features)} features: {current_auc:.4f}")
        
        dropped = []
        
        # Iterate through features (from low IV to high IV)
        for feature in sorted_features:
            # Try removing feature
            test_features = [f for f in current_features if f != feature]
            
            if len(test_features) == 0:
                break
                
            lr.fit(X[test_features], y)
            probs = lr.predict_proba(X[test_features])[:, 1]
            new_auc = roc_auc_score(y, probs)
            
            # Check AUC drop
            auc_drop = current_auc - new_auc
            
            if auc_drop < auc_threshold:
                # AUC didn't drop significantly (or increased), remove feature
                current_features = test_features
                current_auc = new_auc
                dropped.append(feature)
                print(f"Dropped {feature} (IV={iv_df[iv_df['variable']==feature]['iv'].values[0]:.4f}), AUC change: {auc_drop:.4f}, New AUC: {new_auc:.4f}")
            else:
                # AUC dropped significantly, keep feature
                print(f"Kept {feature} (IV={iv_df[iv_df['variable']==feature]['iv'].values[0]:.4f}), AUC drop: {auc_drop:.4f}")
        
        self.dropped_features['backward_iv_selection'] = dropped
        self.selected_features = current_features
        return current_features

    def comprehensive_selection(self, X, y, iv_df=None,
                         iv_threshold=0.02,
                         missing_threshold=0.95,
                         single_value_threshold=0.99,
                         vif_threshold=10,
                         corr_threshold=0.8,
                         selection_method='backward_lr',
                         auc_threshold_decrease=0.001,
                         protected_features=None,
                         woe_transformer=None):
        """
        Comprehensive feature selection pipeline with multiple filtering steps.

        Parameters:
        -----------
        X : DataFrame
            Raw feature data (NOT WOE-transformed data)
        y : Series
            Target variable
        iv_df : DataFrame, optional
            DataFrame with 'variable' and 'iv' columns from binning process
        iv_threshold : float, default=0.02
            IV filtering threshold
        missing_threshold : float, default=0.95
            Missing rate threshold for feature filtering
        single_value_threshold : float, default=0.99
            Single value rate threshold for feature filtering
        vif_threshold : float, default=10
            VIF threshold for feature filtering
        corr_threshold : float, default=0.8
            Correlation threshold for feature filtering
        selection_method : str, default='backward_lr'
            Feature selection method ('stepwise', 'backward_lr', or 'backward_iv')
        auc_threshold_decrease : float, default=0.001
            AUC threshold decrease for backward selection
        protected_features : list, optional
            List of feature names that must be preserved and not filtered out
        woe_transformer : BinningProcess, optional
            WOE transformer for converting raw data to WOE values for model-based selection

        Returns:
        --------
        list
            Selected feature names
        """
        # Initialize protected features set
        protected_features = set(protected_features or [])
        print(f"Protected features (will not be filtered): {list(protected_features)}")

        # Step 1: Missing rate filter
        features_after_missing = self.filter_by_missing_rate(X, threshold=missing_threshold)
        # Restore protected features that may have been filtered
        features_after_missing = list(set(features_after_missing) | protected_features)
        print(f"Features after missing rate filter: {len(features_after_missing)}")

        # Step 2: Single value rate filter
        X_after_missing = X[features_after_missing]
        features_after_single_value = self.filter_by_single_value_rate(X_after_missing, threshold=single_value_threshold)
        # Restore protected features
        features_after_single_value = list(set(features_after_single_value) | protected_features)
        print(f"Features after single value rate filter: {len(features_after_single_value)}")

        # Step 3: IV filter
        # Fix: Only calculate IV if iv_df is NOT provided (was incorrectly checking is not None)
        if iv_df is None:
            # Calculate IV if not provided
            from .binning import BinningProcess
            binning = BinningProcess()
            binning.fit(X[features_after_single_value], y)
            iv_df = binning.iv_df

        features_after_iv = self.filter_by_iv(iv_df, threshold=iv_threshold)
        # Restore protected features
        features_after_iv = list(set(features_after_iv) | protected_features)
        print(f"Features after IV filter: {len(features_after_iv)}")

        # Step 4: Correlation filter
        X_after_iv = X[features_after_iv]
        features_after_corr = self.filter_by_corr(X_after_iv, features=features_after_iv, threshold=corr_threshold)
        # Restore protected features
        features_after_corr = list(set(features_after_corr) | protected_features)
        print(f"Features after correlation filter: {len(features_after_corr)}")

        # Step 5: VIF filter
        X_after_corr = X[features_after_corr]
        features_after_vif = self.filter_by_vif(X_after_corr, features=features_after_corr, threshold=vif_threshold)
        # Restore protected features
        features_after_vif = list(set(features_after_vif) | protected_features)
        print(f"Features after VIF filter: {len(features_after_vif)}")

        # Step 6: Final selection (stepwise or IV ranking or backward IV)
        # NOTE: For model-based selection (stepwise, backward_lr, backward_iv),
        # we need WOE-transformed data since the model expects numeric inputs
        X_final = X[features_after_vif]

        if selection_method in ['stepwise', 'backward_lr', 'backward_iv']:
            if woe_transformer is not None:
                # Transform to WOE for model-based selection
                print(f"Transforming to WOE for {selection_method} selection...")
                X_final_woe = woe_transformer.transform(X_final)
                # Only keep features that exist in WOE transformer
                available_features = [f for f in features_after_vif if f in woe_transformer.rules]
                X_final_woe = X_final_woe[available_features]
                # Fill NaN values with 0 for model training
                X_final_woe = X_final_woe.fillna(0)
            else:
                # Fallback: use only numeric features
                numeric_features = [f for f in features_after_vif if pd.api.types.is_numeric_dtype(X[f])]
                X_final_woe = X[numeric_features].fillna(0)
                print(f"Warning: No WOE transformer provided, using {len(numeric_features)} numeric features for selection")
        else:
            X_final_woe = X_final.fillna(0)

        if selection_method == 'stepwise':
            selected_features = self.stepwise_selection(X_final_woe, y)
        elif selection_method == 'backward_iv':
            selected_features = self.backward_selection_by_iv(X_final_woe, y, iv_df, auc_threshold=auc_threshold_decrease)
        else:  # backward_lr
            selected_features = self.backward_selection_by_iv(X_final_woe, y, iv_df, auc_threshold=auc_threshold_decrease)

        # Map back to original feature names (in case WOE transformation changed column names)
        # and ensure protected features are always included
        selected_features = list(set(selected_features) | protected_features)

        print(f"Final selected features: {len(selected_features)}")
        self.selected_features = selected_features

        return selected_features