from typing import Dict, List, Optional

import matplotlib.pyplot as plt
import numpy as np
import torch
from captum._utils.models.linear_model import SkLearnLinearRegression
from captum.attr import (
    DeepLift,
    FeatureAblation,
    FeaturePermutation,
    IntegratedGradients,
    KernelShap,
    Lime,
    ShapleyValueSampling,
    GradientShap,
)
from captum.attr._core.lime import get_exp_kernel_similarity_function
from torch import nn


DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class Explainer:
    """
    Explainer instance, consisting of several explainability methods.
    """

    def __init__(
        self,
        model: nn.Module,
        feature_names: List,
        explainer_list: List = [
            "feature_ablation",
            "feature_permutation",
            "integrated_gradients",
            "deeplift",
            "shapley_value_sampling",
            "lime",
        ],
        n_steps: int = 500,
        perturbations_per_eval: int = 10,
        n_samples: int = 1000,
        kernel_width: float = 1.0,
        baseline: Optional[torch.Tensor] = None,
    ) -> None:
        self.baseline = baseline
        self.explainer_list = explainer_list
        self.feature_names = feature_names

        # Feature ablation
        feature_ablation_model = FeatureAblation(model)

        def feature_ablation_cbk(X_test: torch.Tensor) -> torch.Tensor:
            out = feature_ablation_model.attribute(
                X_test, n_steps=n_steps, perturbations_per_eval=perturbations_per_eval
            )

            return out

        # Integrated gradients
        integrated_gradients_model = IntegratedGradients(model)

        def integrated_gradients_cbk(X_test: torch.Tensor) -> torch.Tensor:
            return integrated_gradients_model.attribute(X_test, n_steps=n_steps)

        # DeepLift
        deeplift_model = DeepLift(model)

        def deeplift_cbk(X_test: torch.Tensor) -> torch.Tensor:
            return deeplift_model.attribute(X_test)

        # Feature permutation
        feature_permutation_model = FeaturePermutation(model)

        def feature_permutation_cbk(X_test: torch.Tensor) -> torch.Tensor:
            return feature_permutation_model.attribute(
                X_test, n_steps=n_steps, perturbations_per_eval=perturbations_per_eval
            )

        # LIME
        exp_eucl_distance = get_exp_kernel_similarity_function(
            kernel_width=kernel_width
        )
        lime_model = Lime(
            model,
            interpretable_model=SkLearnLinearRegression(),
            similarity_func=exp_eucl_distance,
        )

        def lime_cbk(X_test: torch.Tensor) -> torch.Tensor:
            return lime_model.attribute(
                X_test,
                n_samples=n_samples,
                perturbations_per_eval=perturbations_per_eval,
            )

        # Shapley value sampling
        shapley_value_sampling_model = ShapleyValueSampling(model)

        def shapley_value_sampling_cbk(X_test: torch.Tensor) -> torch.Tensor:
            return shapley_value_sampling_model.attribute(
                X_test,
                n_samples=n_samples,
                perturbations_per_eval=perturbations_per_eval,
            )

        # Kernel SHAP
        kernel_shap_model = KernelShap(model)

        def kernel_shap_cbk(X_test: torch.Tensor) -> torch.Tensor:
            return kernel_shap_model.attribute(
                X_test,
                n_samples=n_samples,
                perturbations_per_eval=perturbations_per_eval,
            )

        # Gradient SHAP
        gradient_shap_model = GradientShap(model)

        def gradient_shap_cbk(X_test: torch.Tensor) -> torch.Tensor:
            return gradient_shap_model.attribute(X_test, baselines=self.baseline)

        self.explainers = {
            "feature_ablation": feature_ablation_cbk,
            "integrated_gradients": integrated_gradients_cbk,
            "deeplift": deeplift_cbk,
            "feature_permutation": feature_permutation_cbk,
            "lime": lime_cbk,
            "shapley_value_sampling": shapley_value_sampling_cbk,
            "kernel_shap": kernel_shap_cbk,
            "gradient_shap": gradient_shap_cbk,
        }

    def _check_tensor(self, X: torch.Tensor) -> torch.Tensor:
        if isinstance(X, torch.Tensor):
            return X.to(DEVICE)
        else:
            return torch.from_numpy(np.asarray(X)).float().to(DEVICE)

    def explain(self, X: torch.Tensor) -> Dict:
        output = {}
        if self.baseline is None:
            self.baseline = torch.zeros(
                X.shape
            )  # Zero tensor as baseline if no baseline specified
        for name in self.explainer_list:
            X_test = self._check_tensor(X)
            self.baseline = self._check_tensor(self.baseline)
            X_test.requires_grad_()
            explainer = self.explainers[name]
            output[name] = explainer(X_test).detach().cpu().numpy()
        return output

    def plot(self, X: torch.Tensor) -> None:
        explanations = self.explain(X)

        fig, axs = plt.subplots(int((len(explanations) + 1) / 2), 2)

        idx = 0
        for name in explanations:
            x_pos = np.arange(len(self.feature_names))

            ax = axs[int(idx / 2), idx % 2]

            ax.bar(x_pos, np.mean(np.abs(explanations[name]), axis=0), align="center")
            ax.set_xlabel("Features")
            ax.set_title(f"{name}")

            idx += 1
        plt.tight_layout()