from core.explanations.analyzers.BaseExplanationAnalyzer import BaseExplanationAnalyzer

import lime.lime_tabular
import numpy as np


class LimeExplanationAnalyzer(BaseExplanationAnalyzer):
    _ranking_cache = {}

    @classmethod
    def calc_feature_ranking(cls, sample, model, dataset):
        """
        Compute feature ranking using LIME values.
        Args:
            sample: np.ndarray, shape (n_features,)
            model: model object (should be compatible with LIME explainer)
            dataset: dataset object (should provide get_data())
        Returns:
            List where the i-th entry is the rank of the i-th feature (0 = most important)
        """
        # 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]

        # Use first 100 training samples as background
        background_data = dataset.get_data()[0][0][:100]

        explainer = lime.lime_tabular.LimeTabularExplainer(
            training_data=background_data,
            mode="classification",
            feature_names=dataset.feature_names,
            discretize_continuous=True
        )

        exp = explainer.explain_instance(sample, model.predict_proba, num_features=dataset.num_features)

        # Get feature importances as a dict {feature_index: importance}
        feature_importance = {i: 0.0 for i in range(dataset.num_features)}
        for idx, importance in exp.as_map()[1]:
            feature_importance[idx] = abs(importance)

        permutation = np.argsort([-feature_importance[i] for i in range(dataset.num_features)])

        # Convert to ranking
        ranking = np.empty_like(permutation)
        ranking[permutation] = np.arange(len(permutation))
        ranking = ranking.tolist()

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

        return ranking
