"""
Advanced Winnow
Advanced selector extending BaseSelector with:
- Feature Importance support
- PCA support
"""

import numpy as np
import pandas as pd
from typing import List, Optional, Dict, Union, Tuple, Callable
from dataclasses import dataclass
from sklearn.decomposition import PCA
from sklearn.feature_selection import mutual_info_classif, mutual_info_regression
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from .base_selector import BaseSelector, SelectionResult
from .logger import logger
try:
    import shap
except ImportError:
    shap = None


@dataclass
class AdvancedSelectionResult(SelectionResult):
    """Class to hold advanced feature selection results"""
    pca_components: Optional[np.ndarray] = None
    pca_explained_variance_ratio: Optional[np.ndarray] = None
    feature_importance_ranking: Optional[pd.DataFrame] = None


class AdvancedSelector(BaseSelector):
    """
    Advanced Winnow class.
    
    Includes all features of BaseSelector plus:
    1. Feature Importance support (externally provided)
    2. PCA support
    
    Parameters
    ----------
    correlation_threshold : float, default=0.8
        Threshold for inter-feature correlation.
    
    target_correlation_threshold : float, default=0.05
        Threshold for feature-target correlation.
    
    variance_threshold : float, default=0.01
        Threshold for variance.
    
    use_feature_importance : bool, default=False
        Whether to use feature importance provided externally.
    
    feature_importance_percentile : float, default=0.5
        Percentage of top features to keep based on importance.
        Example: 0.5 = keep top 50%.
    
    use_pca : bool, default=False
        Whether to use PCA for feature filtering.
    
    pca_variance_threshold : float, default=0.95
        Cumulative explained variance threshold for PCA.
    
    n_pca_components : int, optional
        Number of PCA components (if specified, overrides pca_variance_threshold).
    
    use_borutashap : bool, default=False
        Whether to use BorutaSHAP for feature selection.
    
    n_boruta_trials : int, default=20
        Number of iterations for BorutaSHAP.
    
    use_mrmr : bool, default=False
        Whether to use mRMR (minimum Redundancy Maximum Relevance).
    
    n_mrmr_features : int, optional
        Number of features to select via mRMR. If None, selects until score is non-positive.
    
    Examples
    --------
    >>> import pandas as pd
    >>> from winnow import AdvancedSelector
    >>> 
    >>> X = pd.DataFrame({
    ...     'f1': [1, 2, 3, 4, 5],
    ...     'f2': [2, 4, 6, 8, 10],
    ...     'f3': [5, 4, 3, 2, 1]
    ... })
    >>> y = pd.Series([0, 0, 1, 1, 1])
    >>> 
    >>> # Usage with feature importance
    >>> feature_importance = {'f1': 0.5, 'f2': 0.3, 'f3': 0.2}
    >>> selector = AdvancedSelector(
    ...     use_feature_importance=True,
    ...     feature_importance_percentile=0.5
    ... )
    >>> result = selector.fit_transform(X, y, feature_importance=feature_importance)
    """
    
    def __init__(
        self,
        correlation_threshold: float = 0.8,
        target_correlation_threshold: float = 0.05,
        variance_threshold: float = 0.01,
        use_feature_importance: bool = False,
        feature_importance_percentile: float = 0.5,
        use_pca: bool = False,
        pca_variance_threshold: float = 0.95,
        n_pca_components: Optional[int] = None,
        use_borutashap: bool = False,
        n_boruta_trials: int = 20,
        use_mrmr: bool = False,
        n_mrmr_features: Optional[int] = None
    ):
        super().__init__(
            correlation_threshold=correlation_threshold,
            target_correlation_threshold=target_correlation_threshold,
            variance_threshold=variance_threshold
        )
        
        self.use_feature_importance = use_feature_importance
        self.feature_importance_percentile = feature_importance_percentile
        self.use_pca = use_pca
        self.pca_variance_threshold = pca_variance_threshold
        self.n_pca_components = n_pca_components
        self.use_borutashap = use_borutashap
        self.n_boruta_trials = n_boruta_trials
        self.use_mrmr = use_mrmr
        self.n_mrmr_features = n_mrmr_features
        
        # Results storage
        self._feature_importance: Optional[Dict[str, float]] = None
        self._pca: Optional[PCA] = None
        self._pca_components: Optional[np.ndarray] = None
        self._pca_feature_importance: Optional[pd.Series] = None
        
        # Custom methods storage: name -> callable
        self._custom_methods: Dict[str, Callable] = {}
    
    def fit(
        self, 
        X: pd.DataFrame, 
        y: Optional[Union[pd.Series, str]] = None,
        target_column: Optional[str] = None,
        feature_importance: Optional[Dict[str, float]] = None
    ) -> 'AdvancedSelector':
        """
        Calculates selection criteria.
        
        Parameters
        ----------
        X : pd.DataFrame
            Feature matrix or DataFrame containing all data
        y : pd.Series or str, optional
            Target variable or column name
        target_column : str, optional
            Name of the target column
        feature_importance : Dict[str, float], optional
            Externally provided feature importance scores
            
        Returns
        -------
        self : AdvancedSelector
            Fitted selector instance
        """
        logger.info("Initializing AdvancedSelector fit process...")
        
        # Base selector fit
        super().fit(X, y, target_column)
        
        # Store importance
        if feature_importance is not None:
            self._feature_importance = feature_importance
            logger.debug(f"Stored external importance for {len(feature_importance)} features.")
        
        # PCA fit if enabled
        if self.use_pca:
            logger.debug("Fitting PCA...")
            X_features, _ = self._parse_target(X, y, target_column)
            self._fit_pca(X_features)
        
        return self

    def add_custom_method(self, name: str, method: Callable[[pd.DataFrame, pd.Series, List[str]], List[str]]) -> 'AdvancedSelector':
        """
        Registers a custom feature selection method.
        
        Parameters
        ----------
        name : str
            Unique name for the method.
        method : callable
            A function that takes (X, y, candidate_features) and returns 
             a list of features to REMOVE.
            
        Returns
        -------
        self : AdvancedSelector
        """
        self._custom_methods[name] = method
        logger.info(f"Registered custom selection method: '{name}'")
        return self

    def remove_custom_method(self, name: str) -> 'AdvancedSelector':
        """
        Removes a registered custom selection method.
        """
        if name in self._custom_methods:
            del self._custom_methods[name]
            logger.info(f"Removed custom selection method: '{name}'")
        return self
    
    def transform(
        self, 
        X: pd.DataFrame, 
        y: Optional[pd.Series] = None,
        feature_importance: Optional[Dict[str, float]] = None
    ) -> AdvancedSelectionResult:
        """
        Performs advanced feature selection based on calculated criteria.
        
        Parameters
        ----------
        X : pd.DataFrame
            Feature matrix
        y : pd.Series, optional
            Target variable
        feature_importance : Dict[str, float], optional
            Feature importance scores (if not provided during fit)
            
        Returns
        -------
        AdvancedSelectionResult
            Selected features, PCA results, and statistics
        """
        self._check_is_fitted()
        logger.info("Starting advanced transformation...")
        
        # Update importance
        if feature_importance is not None:
            self._feature_importance = feature_importance
        
        # Run base selection
        base_result = super().transform(X, y)
        
        candidate_features = set(base_result.selected_features)
        removed_features = dict(base_result.removed_features)
        removed_features['low_importance'] = []
        removed_features['pca_eliminated'] = []
        removed_features['borutashap_rejected'] = []
        removed_features['mrmr_excluded'] = []
        
        # Initialize placeholders for custom methods in the result dict
        for name in self._custom_methods.keys():
            removed_features[name] = []
        
        feature_importance_ranking = None
        
        # Feature importance filtering
        if self.use_feature_importance and self._feature_importance is not None:
            low_imp_features, feature_importance_ranking = \
                self._filter_by_importance(list(candidate_features))
            removed_features['low_importance'] = low_imp_features
            candidate_features -= set(low_imp_features)
            logger.debug(f"Removed {len(low_imp_features)} features due to low external importance.")
        
        # PCA analysis
        pca_components = None
        pca_explained_variance = None
        
        if self.use_pca and self._pca is not None:
            pca_eliminated = self._filter_by_pca(list(candidate_features))
            removed_features['pca_eliminated'] = pca_eliminated
            candidate_features -= set(pca_eliminated)
            logger.debug(f"Removed {len(pca_eliminated)} features due to low PCA contribution.")
            
            pca_components = self._pca.components_
            pca_explained_variance = self._pca.explained_variance_ratio_
            
        # BorutaSHAP analysis
        if self.use_borutashap:
            if shap is None:
                logger.warning("SHAP is not installed. Skipping BorutaSHAP.")
            else:
                boruta_rejected = self._filter_by_borutashap(X, y, list(candidate_features))
                removed_features['borutashap_rejected'] = boruta_rejected
                candidate_features -= set(boruta_rejected)
                logger.info(f"BorutaSHAP rejected {len(boruta_rejected)} features.")

        # mRMR analysis
        if self.use_mrmr:
            mrmr_excluded = self._filter_by_mrmr(X, y, list(candidate_features))
            removed_features['mrmr_excluded'] = mrmr_excluded
            candidate_features -= set(mrmr_excluded)
            logger.info(f"mRMR excluded {len(mrmr_excluded)} features.")
            
        # Run custom methods
        for name, method in self._custom_methods.items():
            logger.debug(f"Running custom method: '{name}'...")
            try:
                to_remove = method(X, y, list(candidate_features))
                to_remove_valid = [f for f in to_remove if f in candidate_features]
                
                removed_features[name] = to_remove_valid
                candidate_features -= set(to_remove_valid)
                
                logger.info(f"Custom method '{name}' removed {len(to_remove_valid)} features.")
            except Exception as e:
                logger.error(f"Error in custom method '{name}': {str(e)}")
                raise
        
        # Results
        selected_features = list(candidate_features)
        self._selected_features = selected_features
        
        # Update statistics
        statistics = base_result.statistics.copy()
        if self._pca_feature_importance is not None:
            statistics['pca_feature_importance'] = self._pca_feature_importance.to_frame('pca_importance')
        
        logger.info(f"Advanced transformation complete. Selected {len(selected_features)} features.")
        
        return AdvancedSelectionResult(
            selected_features=selected_features,
            removed_features=removed_features,
            statistics=statistics,
            pca_components=pca_components,
            pca_explained_variance_ratio=pca_explained_variance,
            feature_importance_ranking=feature_importance_ranking
        )
    
    def fit_transform(
        self, 
        X: pd.DataFrame, 
        y: Optional[Union[pd.Series, str]] = None,
        target_column: Optional[str] = None,
        feature_importance: Optional[Dict[str, float]] = None
    ) -> AdvancedSelectionResult:
        """
        Performs fit and transform in one step.
        """
        self.fit(X, y, target_column, feature_importance)
        X_features, y_target = self._parse_target(X, y, target_column)
        return self.transform(X_features, y=y_target, feature_importance=feature_importance)
    
    def _fit_pca(self, X: pd.DataFrame) -> None:
        """Fits PCA and calculates feature contributions."""
        # Normalize for PCA
        X_normalized = (X - X.mean()) / X.std()
        X_normalized = X_normalized.fillna(0)
        
        # Determine number of components
        if self.n_pca_components is not None:
            n_components = min(self.n_pca_components, min(X.shape))
        else:
            temp_pca = PCA()
            temp_pca.fit(X_normalized)
            cum_variance = np.cumsum(temp_pca.explained_variance_ratio_)
            n_components = np.argmax(cum_variance >= self.pca_variance_threshold) + 1
            n_components = max(1, min(n_components, len(X.columns)))
        
        # PCA fit
        self._pca = PCA(n_components=n_components)
        self._pca.fit(X_normalized)
        
        # Calculate feature contributions to components
        # weighted by explained variance ratio
        comp_abs = np.abs(self._pca.components_)
        weighted_imp = comp_abs.T @ self._pca.explained_variance_ratio_
        self._pca_feature_importance = pd.Series(
            weighted_imp, 
            index=X.columns
        )
        logger.debug(f"PCA fitted with {n_components} components.")
    
    def _filter_by_importance(self, candidate_features: List[str]) -> Tuple[List[str], pd.DataFrame]:
        """Filters features by external importance scores."""
        imp_values = {f: self._feature_importance.get(f, 0) for f in candidate_features}
        imp_df = pd.DataFrame([{'feature': f, 'importance': v} for f, v in imp_values.items()])
        imp_df = imp_df.sort_values('importance', ascending=False)
        
        n_keep = max(1, int(len(candidate_features) * self.feature_importance_percentile))
        keep = imp_df.head(n_keep)['feature'].tolist()
        remove = [f for f in candidate_features if f not in keep]
        
        return remove, imp_df
    
    def _filter_by_pca(self, candidate_features: List[str]) -> List[str]:
        """Filters features by PCA contribution."""
        if self._pca_feature_importance is None:
            return []
        
        pca_imp = self._pca_feature_importance[candidate_features]
        threshold = pca_imp.median() * 0.5
        return pca_imp[pca_imp < threshold].index.tolist()
    
    def get_pca_transformed(self, X: pd.DataFrame) -> pd.DataFrame:
        """Returns PCA transformed data."""
        if self._pca is None:
            raise ValueError("PCA was not used or is not fitted yet.")
        
        X_normalized = (X - X.mean()) / X.std()
        X_normalized = X_normalized.fillna(0)
        
        pca_data = self._pca.transform(X_normalized)
        cols = [f'PC{i+1}' for i in range(pca_data.shape[1])]
        
        return pd.DataFrame(pca_data, index=X.index, columns=cols)
    
    def get_combined_importance(self, X: pd.DataFrame) -> pd.DataFrame:
        """Returns a DataFrame combining all importance metrics."""
        self._check_is_fitted()
        
        metrics = {
            'target_correlation': self._target_correlations,
            'variance': self._variances
        }
        
        if self._feature_importance is not None:
            metrics['external_importance'] = pd.Series(self._feature_importance)
        
        if self._pca_feature_importance is not None:
            metrics['pca_importance'] = self._pca_feature_importance
        
        df = pd.DataFrame(metrics)
        
        # Normalized scores
        for col in df.columns:
            m = df[col].max()
            if m > 0:
                df[f'{col}_normalized'] = df[col] / m
        
        norm_cols = [c for c in df.columns if c.endswith('_normalized')]
        if norm_cols:
            df['combined_score'] = df[norm_cols].mean(axis=1)
            df = df.sort_values('combined_score', ascending=False)
        
        return df
    
    def _filter_by_borutashap(self, X: pd.DataFrame, y: pd.Series, candidate_features: List[str]) -> List[str]:
        """Simplified BorutaSHAP implementation."""
        logger.info(f"Running BorutaSHAP with {self.n_boruta_trials} trials...")
        X_curr = X[candidate_features]
        hits = np.zeros(len(candidate_features))
        
        # Determine model based on target
        if y.dtype.kind in 'fc' or y.nunique() > 20: # Heuristic for reg
            model_class = RandomForestRegressor
        else:
            model_class = RandomForestClassifier
            
        for _ in range(self.n_boruta_trials):
            # Create shadow features
            shadow_X = X_curr.copy().apply(np.random.permutation)
            shadow_X.columns = [f"shadow_{c}" for c in shadow_X.columns]
            
            combined_X = pd.concat([X_curr, shadow_X], axis=1)
            
            model = model_class(n_estimators=50, max_depth=5, random_state=42)
            model.fit(combined_X, y)
            
            explainer = shap.TreeExplainer(model)
            shap_values = explainer.shap_values(combined_X)
            
            # SHAP values handling (ensure 2D: samples x features)
            if isinstance(shap_values, list):
                # For multiclass RF, shap_values is a list of [samples, features]
                # We take mean across all classes for simplicity in importance
                shap_values = np.mean([np.abs(v) for v in shap_values], axis=0)
            elif len(shap_values.shape) == 3:
                # Some newer SHAP versions return (samples, features, classes)
                shap_values = np.abs(shap_values).mean(axis=2)
            else:
                shap_values = np.abs(shap_values)
                
            # Collapse to feature-wise importance (1D array: features)
            global_importance = shap_values.mean(axis=0)
            
            real_importance = global_importance[:len(candidate_features)]
            shadow_importance = global_importance[len(candidate_features):]
            
            max_shadow = np.max(shadow_importance)
            hits += (real_importance > max_shadow)
            
        # Statistical decision (simple threshold for now)
        # Keep features that performed better than shadow features in at least 50% of trials
        threshold = self.n_boruta_trials * 0.5
        rejected = [f for i, f in enumerate(candidate_features) if hits[i] < threshold]
        
        return rejected

    def _filter_by_mrmr(self, X: pd.DataFrame, y: pd.Series, candidate_features: List[str]) -> List[str]:
        """mRMR implementation (Minimum Redundancy Maximum Relevance)."""
        logger.info("Running mRMR selection...")
        X_curr = X[candidate_features]
        
        # Relevance: MI (X_i, y)
        if y.dtype.kind in 'fc' or y.nunique() > 20:
            relevance = mutual_info_regression(X_curr, y)
        else:
            relevance = mutual_info_classif(X_curr, y)
            
        relevance_series = pd.Series(relevance, index=candidate_features)
        
        n_to_select = self.n_mrmr_features or len(candidate_features)
        selected = []
        remaining = list(candidate_features)
        
        # Step 1: Select feature with max relevance
        first_feature = relevance_series.idxmax()
        selected.append(first_feature)
        remaining.remove(first_feature)
        
        # Step 2: Iteratively select features
        while len(selected) < n_to_select and remaining:
            best_score = -np.inf
            best_feature = None
            
            for f in remaining:
                # Max Relevance
                rel = relevance_series[f]
                
                # Min Redundancy: Average MI (f, s) for s in selected
                # For performance, we sample or just use correlation as a proxy if MI is slow
                # Here we use mutual info for consistency
                redundancy = 0
                for s in selected:
                    # Approximation: MI between features
                    # Using correlation as proxy for redundancy to keep it faster
                    redundancy += abs(X_curr[f].corr(X_curr[s]))
                
                redundancy /= len(selected)
                
                score = rel - redundancy
                if score > best_score:
                    best_score = score
                    best_feature = f
            
            if best_score <= 0 and self.n_mrmr_features is None:
                break
                
            selected.append(best_feature)
            remaining.remove(best_feature)
            
        return remaining

    def __repr__(self) -> str:
        base_repr = super().__repr__()
        return (
            f"AdvancedSelector("
            f"use_feature_importance={self.use_feature_importance}, "
            f"use_borutashap={self.use_borutashap}, "
            f"use_mrmr={self.use_mrmr}, "
            f"{base_repr})"
        )

