from core.explanations.analyzers.BaseExplanationAnalyzer import BaseExplanationAnalyzer

import numpy as np
import shap


class ShapExplanationAnalyzer(BaseExplanationAnalyzer):

    _ranking_cache = {}

    @classmethod
    def calc_feature_ranking(cls, sample, model, dataset):
        """
        Compute feature ranking using SHAP values.
        Args:
            sample: np.ndarray, shape (n_features,)
            model: model object (should be compatible with SHAP Explainer)
            dataset: dataset object (not used directly, but may be useful for feature names)
            features: list of feature names or indices
        Returns:
            List of feature indices sorted by SHAP importance (descending)
        """
        # Create a unique cache key
        sample_key = tuple(sample.tolist()) if hasattr(sample, 'tolist') else tuple(sample)
        dataset_key = getattr(dataset, 'name', id(dataset))
        cache_key = (sample_key, dataset_key)

        if cache_key in cls._ranking_cache:
            return cls._ranking_cache[cache_key]

        # Wrap the model's predict method in a callable for SHAP
        # Use a background dataset for SHAP explanation
        background_data = dataset.get_data()[0][0][:100]  # Use first 100 training samples as background

        explainer = shap.Explainer(model.classify, background_data)
        shap_values = explainer(np.array([sample]))

        if hasattr(shap_values, 'values'):
            values = shap_values.values
        else:
            values = shap_values

        if values.ndim == 3:
            mean_abs_shap = np.mean(np.abs(values), axis=(0,1))
        else:
            mean_abs_shap = np.mean(np.abs(values), axis=0)

        permutation = np.argsort(-mean_abs_shap)

        # Convert permutation to ranking: ranking[i] = rank of feature i
        ranking = np.empty_like(permutation)
        ranking[permutation] = np.arange(len(permutation))
        ranking = ranking.tolist()

        # Cache the ranking
        cls._ranking_cache[cache_key] = ranking

        return ranking
