"""LIME explainer implementation."""

import lime
import lime.lime_tabular

from xai_suite.explainers.base import BaseExplainer
from xai_suite.utils.model_wrapper import ModelWrapper


class LimeExplainer(BaseExplainer):
    """LIME implementation of the BaseExplainer."""

    def _initialize_explainer(self, **kwargs):
        """Initializes the LIME TabularExplainer."""
        print("Initializing LIME TabularExplainer...")
        self.mode = kwargs.get("mode", "classification")  # Store the mode

        return lime.lime_tabular.LimeTabularExplainer(
            training_data=self.data.values,
            feature_names=self.data.columns.tolist(),
            class_names=kwargs.get("class_names", ["class_0", "class_1"]),
            mode=self.mode,
        )

    def explain_local(self, instance):
        """
        Generates LIME explanation for a single instance.
        Dynamically selects the prediction function based on the mode.
        """
        model_wrapper = ModelWrapper(self.model)
        instance_np = instance.values[0]

        # Choose the appropriate prediction function based on the mode
        if self.mode == "classification":
            predict_fn = model_wrapper.predict_proba
        elif self.mode == "regression":
            predict_fn = model_wrapper.predict_score
        else:
            raise ValueError(f"Unsupported LIME mode: {self.mode}")

        explanation = self.explainer.explain_instance(
            instance_np, predict_fn, num_features=len(self.data.columns)
        )

        explanation_list = explanation.as_list()

        return {
            "method": "LIME",
            "feature_names": [item[0] for item in explanation_list],
            "importance_scores": [item[1] for item in explanation_list],
        }
