"""SHAP explainer implementation."""

import shap

from xai_suite.explainers.base import BaseExplainer


class ShapExplainer(BaseExplainer):
    """
    SHAP implementation that automatically selects the best explainer
    (e.g., TreeExplainer for tree models, KernelExplainer for others).
    """

    def _initialize_explainer(self, **kwargs):
        """Initializes the SHAP Explainer."""
        print("Initializing SHAP Explainer (auto-detecting model type)...")

        # We pass the model and the background data.
        # SHAP's API will automatically wrap the model and choose the best explainer.
        return shap.Explainer(self.model, self.data)

    def explain_local(self, instance):
        """
        Generates SHAP values for a single instance.
        Handles both classification and regression/scoring.

        Args:
            instance (pd.DataFrame): A single row DataFrame to explain.

        Returns:
            dict: A dictionary with SHAP explanation values.
        """
        shap_explanation = self.explainer(instance)

        # The output structure of shap_explanation.values depends on the model output.
        # - For classifiers (2+ outputs): shape is (num_samples, num_features, num_classes)
        # - For regression/scoring (1 output): shape is (num_samples, num_features)

        if len(shap_explanation.values.shape) == 3:  # It's a classifier
            # Default to explaining the positive class (class 1)
            class_index = 1
            shap_values = shap_explanation.values[0, :, class_index]
            base_value = shap_explanation.base_values[0, class_index]
        else:  # It's a regression or scoring model (like IsolationForest)
            shap_values = shap_explanation.values[0, :]
            base_value = shap_explanation.base_values[0]

        return {
            "method": "SHAP",
            "feature_names": list(instance.columns),
            "feature_values": instance.values[0].tolist(),
            "importance_scores": shap_values.tolist(),
            "base_value": float(base_value),
            "output_value": float(base_value + shap_values.sum()),
        }

    def explain_global(self):
        """Generates a SHAP summary plot."""
        print("Generating SHAP global summary plot...")
        sample_data = self.data.sample(min(100, len(self.data)), random_state=42)

        shap_values_obj = self.explainer(sample_data)

        # For plot, we may need to specify the output if multi-class
        plot_values = shap_values_obj
        if len(shap_values_obj.values.shape) == 3:
            plot_values = shap_values_obj[:, :, 1]  # Summary for positive class

        shap.summary_plot(plot_values, sample_data, show=False)
        return "SHAP summary plot generated. Use matplotlib.pyplot.show() to display."
