import shap
import numpy as np
import pandas as pd

def compute_shap_stability(model, X: pd.DataFrame, y: pd.Series, n_splits=5, classification=True, random_state=42):
    from sklearn.model_selection import KFold, StratifiedKFold

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

    fold_importances = []

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

        model_clone = clone_model(model)
        model_clone.fit(X_tr, y_tr)

        explainer = shap.TreeExplainer(model_clone)

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

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

    df = pd.DataFrame({
        "feature": X.columns,
        "mean_shap": mean_importance,
        "std_shap": std_importance,
        "coef_var": coef_var,
        "stability_score": stability_score
    }).sort_values("stability_score", ascending=False).reset_index(drop=True)

    return df

def clone_model(model):
    from sklearn.base import clone
    return clone(model)