# Copyright © 2023-2024 ValidMind Inc. All rights reserved.
# See the LICENSE file in the root of this repository for details.
# SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial

import warnings
from dataclasses import dataclass

import matplotlib.pyplot as plt
import numpy as np
import shap

from validmind.errors import UnsupportedModelForSHAPError
from validmind.logging import get_logger
from validmind.models import CatBoostModel, SKlearnModel, StatsModelsModel
from validmind.vm_models import Figure, Metric

logger = get_logger(__name__)


@dataclass
class SHAPGlobalImportance(Metric):
    """
    Evaluates and visualizes global feature importance using SHAP values for model explanation and risk identification.

    ### Purpose

    The SHAP (SHapley Additive exPlanations) Global Importance metric aims to elucidate model outcomes by attributing
    them to the contributing features. It assigns a quantifiable global importance to each feature via their respective
    absolute Shapley values, thereby making it suitable for tasks like classification (both binary and multiclass).
    This metric forms an essential part of model risk management.

    ### Test Mechanism

    The exam begins with the selection of a suitable explainer which aligns with the model's type. For tree-based
    models like XGBClassifier, RandomForestClassifier, CatBoostClassifier, TreeExplainer is used whereas for linear
    models like LogisticRegression, XGBRegressor, LinearRegression, it is the LinearExplainer. Once the explainer
    calculates the Shapley values, these values are visualized using two specific graphical representations:

    1. Mean Importance Plot: This graph portrays the significance of individual features based on their absolute
    Shapley values. It calculates the average of these absolute Shapley values across all instances to highlight the
    global importance of features.

    2. Summary Plot: This visual tool combines the feature importance with their effects. Every dot on this chart
    represents a Shapley value for a certain feature in a specific case. The vertical axis is denoted by the feature
    whereas the horizontal one corresponds to the Shapley value. A color gradient indicates the value of the feature,
    gradually changing from low to high. Features are systematically organized in accordance with their importance.
    These plots are generated by the function `_generate_shap_plot()`.

    ### Signs of High Risk

    - Overemphasis on certain features in SHAP importance plots, thus hinting at the possibility of model overfitting
    - Anomalies such as unexpected or illogical features showing high importance, which might suggest that the model's
    decisions are rooted in incorrect or undesirable reasoning
    - A SHAP summary plot filled with high variability or scattered data points, indicating a cause for concern

    ### Strengths

    - SHAP does more than just illustrating global feature significance, it offers a detailed perspective on how
    different features shape the model's decision-making logic for each instance.
    - It provides clear insights into model behavior.

    ### Limitations

    - High-dimensional data can convolute interpretations.
    - Associating importance with tangible real-world impact still involves a certain degree of subjectivity.
    """

    name = "shap"
    required_inputs = ["model", "dataset"]
    tasks = ["classification", "text_classification"]
    tags = [
        "sklearn",
        "binary_classification",
        "multiclass_classification",
        "feature_importance",
        "visualization",
    ]
    default_params = {
        "kernel_explainer_samples": 10,
        "tree_or_linear_explainer_samples": 200,
        "class_of_interest": None,
    }

    def _generate_shap_plot(self, type_, shap_values, x_test):
        """
        Plots two types of SHAP global importance (SHAP).
        :params type: mean, summary
        :params shap_values: a matrix
        :params x_test:
        """
        plt.close("all")

        # preserve styles
        # mpl.rcParams["grid.color"] = "#CCC"
        ax = plt.axes()
        ax.set_facecolor("white")

        summary_plot_extra_args = {}
        if type_ == "mean":
            # Calculate the mean absolute SHAP value for each feature
            mean_abs_shap = np.abs(shap_values).mean(axis=0)

            # Find the maximum mean absolute SHAP value
            max_shap_value = np.max(mean_abs_shap)

            # Normalize all SHAP values based on the top feature
            shap_values = (
                shap_values / max_shap_value * 100
            )  # scaling factor to make the top feature 100%
            summary_plot_extra_args = {"plot_type": "bar"}

            shap.summary_plot(
                shap_values, x_test, show=False, **summary_plot_extra_args
            )

            # Customize the plot using matplotlib
            plt.xlabel("Normalized SHAP Value (Percentage)", fontsize=13)
            plt.ylabel("Features", fontsize=13)
            plt.title("Normalized Feature Importance", fontsize=13)
        else:
            shap.summary_plot(
                shap_values, x_test, show=False, **summary_plot_extra_args
            )

        figure = plt.gcf()
        # avoid displaying on notebooks and clears the canvas for the next plot
        plt.close()

        return Figure(
            for_object=self,
            figure=figure,
            key=f"shap:{type_}",
            metadata={"type": type_},
        )

    def run(self):
        if not isinstance(self.inputs.model, SKlearnModel) or isinstance(
            self.inputs.model, (CatBoostModel, StatsModelsModel)
        ):
            logger.info(f"Skiping SHAP for {self.inputs.model.library} models")
            return

        trained_model = self.inputs.model.model
        model_class = self.inputs.model.class_

        # the shap library generates a bunch of annoying warnings that we don't care about
        warnings.filterwarnings("ignore", category=UserWarning)

        # Any tree based model can go here
        if (
            model_class == "XGBClassifier"
            or model_class == "RandomForestClassifier"
            or model_class == "CatBoostClassifier"
            or model_class == "DecisionTreeClassifier"
            or model_class == "RandomForestRegressor"
            or model_class == "GradientBoostingRegressor"
        ):
            explainer = shap.TreeExplainer(trained_model)
        elif (
            model_class == "LogisticRegression"
            or model_class == "XGBRegressor"
            or model_class == "LinearRegression"
            or model_class == "LinearSVC"
        ):
            explainer = shap.LinearExplainer(trained_model, self.inputs.dataset.x)
        elif model_class == "SVC":
            # KernelExplainer is slow so we use shap.sample to speed it up
            explainer = shap.KernelExplainer(
                trained_model.predict,
                shap.sample(
                    self.inputs.dataset.x,
                    self.params["kernel_explainer_samples"],
                ),
            )
        else:
            model_class = "<ExternalModel>" if model_class is None else model_class
            raise UnsupportedModelForSHAPError(
                f"Model {model_class} not supported for SHAP importance."
            )

        # KernelExplainer is slow so we use shap.sample to speed it up
        if isinstance(explainer, shap.KernelExplainer):
            shap_sample = shap.sample(
                self.inputs.dataset.x_df(),
                self.params["kernel_explainer_samples"],
            )
        else:
            shap_sample = self.inputs.dataset.x_df().sample(
                min(
                    self.params["tree_or_linear_explainer_samples"],
                    self.inputs.dataset.x_df().shape[0],
                )
            )

        shap_values = explainer.shap_values(shap_sample)

        # Select the SHAP values for the specified class (classification) or for the regression output.
        class_of_interest = self.params["class_of_interest"]
        shap_values = _select_shap_values(shap_values, class_of_interest)

        figures = [
            self._generate_shap_plot("mean", shap_values, shap_sample),
            self._generate_shap_plot("summary", shap_values, shap_sample),
        ]

        # restore warnings
        warnings.filterwarnings("default", category=UserWarning)

        return self.cache_results(figures=figures)

    def test(self):
        """Unit Test for SHAP Global Importance Metric"""
        # Verify that the result object is not None
        assert self.result is not None

        # Verify that there are exactly two figures in the figures list
        assert len(self.result.figures) == 2

        # Verify that each figure is an instance of Figure and has the correct metadata type
        for fig_num, type_ in enumerate(["mean", "summary"], start=1):
            assert isinstance(self.result.figures[fig_num - 1], Figure)
            assert self.result.figures[fig_num - 1].metadata["type"] == type_


def _select_shap_values(shap_values, class_of_interest=None):
    """
    Selects SHAP values for binary or multiclass classification. For regression models,
    returns the SHAP values directly as there are no classes.

    Parameters:
    -----------
    shap_values : list or numpy.ndarray
        The SHAP values returned by the SHAP explainer. For multiclass classification,
        this will be a list where each element corresponds to a class. For regression,
        this will be a single array of SHAP values.

    class_of_interest : int, optional
        The class index for which to retrieve SHAP values. If None (default), the function
        will assume binary classification and use class 1 by default.

    Returns:
    --------
    numpy.ndarray
        The SHAP values for the specified class (classification) or for the regression output.

    Raises:
    -------
    ValueError
        If class_of_interest is specified and is out of bounds for the number of classes.
    """
    # Check if we are dealing with a multiclass classification
    if isinstance(shap_values, list):
        num_classes = len(shap_values)

        # Default to class 1 for binary classification
        if num_classes == 2 and class_of_interest is None:
            logger.info(
                "Binary classification detected: using SHAP values for class 1 (positive class)."
            )
            return shap_values[1]
        else:
            # Multiclass classification: use the specified class_of_interest
            if class_of_interest is not None and 0 <= class_of_interest < num_classes:
                logger.info(
                    f"Multiclass classification: using SHAP values for class {class_of_interest}."
                )
                return shap_values[class_of_interest]
            else:
                raise ValueError(
                    f"Invalid class_of_interest: {class_of_interest}. Must be between 0 and {num_classes - 1}."
                )
    else:
        # For regression, return the SHAP values as they are
        logger.info("Regression model detected: returning SHAP values as-is.")
        return shap_values
