from enum import Enum
from typing import Dict, Any, List, Optional
from pydantic import BaseModel, model_validator

class PlotLibrary(str, Enum):
    PLT = "PLT"
    SHAP = "SHAP"


class PlotType(str, Enum):
    BAR = "BAR"
    WATERFALL = "WATERFALL"
    PLOT = "PLOT"


class IndividualPlot(BaseModel):
    plot_type: PlotType
    kwargs: Optional[Dict] = {}
    args: Optional[List] = []
    xlabel: Optional[str] = None
    ylabel: Optional[str] = None
    title: Optional[str] = None

    @model_validator(mode="before")
    @classmethod
    def backwards_compatibility_plot_data(cls, data: Any) -> Any:
        if isinstance(data, dict):
            if "plot_data" in data:
                plot_data = data.pop("plot_data")
                data["kwargs"] = plot_data
        return data

    def model_dump(self, **kwargs: Any) -> Dict[str, Any]:
        # configure parent method
        kwargs.setdefault("exclude_none", True)
        kwargs.setdefault("exclude_defaults", True)
        return super().model_dump(**kwargs)

    def plot(self, plot_library: PlotLibrary):
        from matplotlib import pyplot as plt
        import shap
        import numpy as np

        shap_plot_function_mapping = {"WATERFALL": shap.plots.waterfall}
        plt_plot_function_mapping = {"PLOT": plt.plot, "BAR": plt.bar}

        plot_type = self.plot_type
        kwargs = self.kwargs
        args = self.args

        if plot_library == "SHAP":
            shap_values = shap._explanation.Explanation(
                values=np.array(kwargs["values"]),
                base_values=kwargs["base_value"],
                data=np.array(kwargs["data"]),
                feature_names=kwargs["feature_names"],
            )

            shap_plot_function_mapping[plot_type](shap_values, show=False)

        elif plot_library == "PLT":
            plt_plot_function_mapping[plot_type](*args, **kwargs)

            fig = plt.gcf()
            for axis in fig.axes:
                if self.title:
                    axis.set_title(self.title)
                if self.xlabel:
                    axis.set_xlabel(self.xlabel)
                if self.ylabel:
                    axis.set_ylabel(self.ylabel)

        else:
            raise ValueError(f"Unknown plot library {plot_library}")


class PlotData(BaseModel):
    plot_name: str
    plot_library: PlotLibrary
    plot_content: List[IndividualPlot]

    def create_plot(self) -> None:
        plot_library = self.plot_library
        plot_content = self.plot_content

        for individual_plot in plot_content:
            individual_plot.plot(plot_library=plot_library)
