import pandas as pd
import numpy as np
from sklearn.model_selection import KFold, StratifiedKFold
from ..core.component import AivoraComponent
from ..generators.oof import OOFModelGenerator

class SHAPStabilitySelector(AivoraComponent):

    def __init__(self, model=None, n_splits=5, threshold=0.05, classification=True, random_state=42):
        self.model = model
        self.n_splits = n_splits
        self.threshold = threshold
        self.classification = classification
        self.random_state = random_state
        self.selected_features_ = None
        self.feature_scores_ = None

    def fit(self, X: pd.DataFrame, y: pd.Series):

        X_clean = X.apply(pd.to_numeric, errors='coerce').fillna(0)
        X_clean.replace([np.inf, -np.inf], 0, inplace=True)
        
        n_samples, n_features = X_clean.shape
        all_importances = []

        if self.classification:
            splitter = StratifiedKFold(n_splits=self.n_splits, shuffle=True, random_state=self.random_state)
        else:
            splitter = KFold(n_splits=self.n_splits, shuffle=True, random_state=self.random_state)

        import shap
        
        for train_idx, val_idx in splitter.split(X_clean, y):
            X_tr, X_val = X_clean.iloc[train_idx], X_clean.iloc[val_idx]
            y_tr = y.iloc[train_idx]

            model_clone = self._clone_model(self.model)
            model_clone.fit(X_tr, y_tr)

            explainer = shap.TreeExplainer(model_clone)

            if self.classification:
                shap_values = explainer.shap_values(X_val)
                shap_values = np.abs(shap_values).mean(axis=0)
            else:
                shap_values = np.abs(explainer(X_val).values).mean(axis=0)

            all_importances.append(shap_values)

        all_importances = np.array(all_importances)
        mean_importance = all_importances.mean(axis=0)
        std_importance = all_importances.std(axis=0)
        coef_var = std_importance / (mean_importance + 1e-8)
        stability_score = mean_importance / (1 + coef_var)

        self.feature_scores_ = pd.DataFrame({
            "feature": X_clean.columns,
            "stability_score": stability_score
        }).sort_values("stability_score", ascending=False).reset_index(drop=True)

        self.selected_features_ = self.feature_scores_[self.feature_scores_["stability_score"] > self.threshold]["feature"].tolist()
        return self

    def transform(self, X: pd.DataFrame) -> pd.DataFrame:
        
        if self.selected_features_ is None:
            raise ValueError("Selector not fitted yet.")
            
        return X[self.selected_features_]

    def fit_transform(self, X, y):
        return self.fit(X, y).transform(X)

    def _clone_model(self, model):
        from sklearn.base import clone
        return clone(model)
