from __future__ import annotations

import abc
from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, TypeVar

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.axes import Axes
from matplotlib.figure import Figure
from numpy.typing import ArrayLike, NDArray

import boss.pp.style as style
import boss.pp.tools as tools
from boss.pp.mesh import Mesh

style.load_style()

if TYPE_CHECKING:
    from pathlib import Path

    from boss.bo.results import BOResults
    from boss.pp.elements import Element

T = TypeVar("T", bound="MeshGraphic")

LegendLocationLiteral = Literal[
    "outside top",
    "outside right",
    "outside bottom",
    "best",
    "upper right",
    "upper left",
    "lower right",
    "lower left",
    "right",
    "center right",
    "center left",
    "lower center",
    "upper center",
    "center",
]


def _dispatch_target(
    results: BOResults,
    target: Literal["model_mean", "model_var", "model_std", "acqfn"],
    itr: int = -1,
) -> tuple[Callable, str]:
    """
    Dispatch target function and label based on target type.

    Parameters
    ----------
    results : BOResults
        Bayesian optimization results containing models and acquisition functions.
    target : {"model_mean", "model_var", "model_std", "acqfn"}
        Type of target to create function for.
    itr : int, default=-1
        Iteration number to use for model reconstruction.

    Returns
    -------
    target_func : Callable
        Function that evaluates the target on input arrays.
    label : str
        LaTeX label for the target function.
    """
    if target == "model_mean":
        model = results.reconstruct_model(itr)
        target_func = lambda X: model.predict(X)[0]
        label = r"$\mu$"
    elif target == "model_var":
        model = results.reconstruct_model(itr)
        target_func = lambda X: model.predict(X)[1]
        label = r"$\nu^2$"
    elif target == "model_std":
        model = results.reconstruct_model(itr)
        target_func = lambda X: np.sqrt(model.predict(X)[1])
        label = r"$\nu$"
    elif target == "acqfn":
        acqfn = results.reconstruct_acq_func(itr)
        target_func = acqfn.evaluate
        label = r"$\alpha$"
    return target_func, label


class MeshGraphic(abc.ABC, Generic[T]):
    """
    Abstract base class for creating mesh-based graphics from Bayesian optimization results.

    This class provides a framework for visualizing targets (model predictions,
    acquisition functions, etc.) over coordinate meshes. It handles mesh creation,
    target evaluation, and figure setup for various plotting backends.

    Parameters
    ----------
    results : BOResults
        Bayesian optimization results containing models and data.
    target : {"model_mean", "model_var", "model_std", "acqfn"} or Callable
        Target function to visualize. Can be a string identifier for built-in
        targets or a custom callable.
    active_input_dims : ArrayLike
        Dimensions of the input space that vary in the visualization.
    itr : int, default=-1
        Iteration number to use for model reconstruction.
    active_target_dims : ArrayLike, default=0
        Target output dimensions to visualize.
    fixed_inputs : {"min", "mid", "next_acq"} or dict or None, default="min"
        How to fix inactive input dimensions. String options use presets,
        dict maps dimension indices to fixed values.
    grid_pts : int or tuple, default=50
        Number of grid points per dimension.
    fig : Figure, optional
        Existing matplotlib figure to use.
    ax : Axes, optional
        Existing matplotlib axes to use.

    Attributes
    ----------
    NUM_AXIS : int
        Number of axes required for this graphic type (2D or 3D).
    """
    NUM_AXIS: int

    def __init__(
        self,
        results: BOResults,
        target: Literal["model_mean", "model_var", "model_std", "acqfn"] | Callable,
        active_input_dims: ArrayLike,
        itr: int = -1,
        active_target_dims: ArrayLike = 0,
        fixed_inputs: (
            Literal["min", "mid", "next_acq"] | dict[int, float] | None
        ) = "min",
        grid_pts: int | tuple[int, ...] = 50,
        fig: Figure | None = None,
        ax: Axes | None = None,
    ) -> None:
        mesh = Mesh(results.bounds, active_dims=active_input_dims, grid_pts=grid_pts)
        self._setup(
            results=results,
            target=target,
            mesh=mesh,
            itr=itr,
            active_target_dims=active_target_dims,
            fixed_inputs=fixed_inputs,
            fig=fig,
            ax=ax,
        )

    @classmethod
    def from_mesh(
        cls: type[T],
        results: BOResults,
        target: (
            Literal["model_mean", "model_var", "model_std", "acqfn"]
            | Callable
            | NDArray
        ),
        mesh: Mesh,
        itr: int = -1,
        active_target_dims: int = 0,
        fixed_inputs: (
            Literal["min", "mid", "next_acq"] | dict[int, float] | None
        ) = "min",
        fig: Figure | None = None,
        ax: Axes | None = None,
    ) -> T:
        """
        Create a MeshGraphic instance from an existing Mesh object.

        Parameters
        ----------
        results : BOResults
            Bayesian optimization results.
        target : {"model_mean", "model_var", "model_std", "acqfn"} or Callable or NDArray
            Target function or pre-computed target values.
        mesh : Mesh
            Pre-configured mesh object defining the coordinate grid.
        itr : int, default=-1
            Iteration number for model reconstruction.
        active_target_dims : int, default=0
            Target output dimension to visualize.
        fixed_inputs : {"min", "mid", "next_acq"} or dict or None, default="min"
            How to fix inactive input dimensions.
        fig : Figure, optional
            Existing matplotlib figure.
        ax : Axes, optional
            Existing matplotlib axes.

        Returns
        -------
        MeshGraphic
            Configured graphic instance.
        """
        self = cls.__new__(cls)
        self._setup(
            results=results,
            target=target,
            mesh=mesh,
            itr=itr,
            active_target_dims=active_target_dims,
            fixed_inputs=fixed_inputs,
            fig=fig,
            ax=ax,
        )
        return self

    def _setup(
        self,
        results: BOResults,
        target: (
            Literal["model_mean", "model_var", "model_std", "acqfn"]
            | Callable
            | NDArray
        ),
        mesh: Mesh,
        itr: int = -1,
        active_target_dims: ArrayLike = 0,
        fixed_inputs: (
            Literal["min", "mid", "next_acq"] | dict[int, float] | None
        ) = "min",
        fig: Figure | None = None,
        ax: Axes | None = None,
    ) -> None:
        self.results = results
        self.itr = itr
        self.mesh = mesh
        self.active_input_dims = self.mesh.active_dims
        self.active_target_dims = np.atleast_1d(active_target_dims)
        self.itr = itr
        self.fixed_inputs = fixed_inputs

        if isinstance(fixed_inputs, str):
            try:
                self.mesh.fix_dim_preset(results, fixed_inputs, itr=itr)
            except RuntimeError:
                self.fixed_inputs = "mid"
                self.mesh.fix_dim_preset(results, self.fixed_inputs, itr=itr)
        elif isinstance(fixed_inputs, dict):
            self.mesh.fix_dims(fixed_inputs)

        self.target_label = ""
        self.target_func = None
        if callable(target):
            self.target_func = target
        elif isinstance(target, np.ndarray):
            self.target_values = (target,)
        elif isinstance(target, str):
            self.target_func, self.target_label = _dispatch_target(results, target, itr)
        else:
            self.target_values = target

        if self.target_func:
            self.target_values = self.mesh.evaluate_func(self.target_func)

        subplot_kw = {"projection": "3d"} if self.NUM_AXIS == 3 else {}

        if not fig and not ax:
            self.fig, self.ax = plt.subplots(subplot_kw=subplot_kw)
        elif not fig and ax:
            self.fig = plt.gcf()
        elif fig and not ax:
            self.ax = fig.gca(**subplot_kw)
        elif fig and ax:
            self.fig = fig
            self.ax = ax

        self.elements = []

    def add(self, *elements: Element) -> None:
        """
        Add elements to the graphic for plotting.

        Parameters
        ----------
        *elements : Element
            Graphic elements to add (e.g., points, confidence intervals).
        """
        for elem in elements:
            elem.graphic = self
            self.elements.append(elem)

    def __iadd__(self: T, element: Element) -> T:
        self.add(element)
        return self

    def save(self, path: str | Path, **kwargs: Any) -> None:
        """
        Save the figure to file.

        Parameters
        ----------
        path : str or Path
            Output file path.
        **kwargs : Any
            Additional keyword arguments passed to plt.savefig().
        """
        plt.savefig(path, **kwargs)

    def show(self) -> None:
        """Display the figure."""
        plt.show()

    @abc.abstractmethod
    def plot(self) -> None:
        pass


class Contour(MeshGraphic["Contour"]):
    """
    2D contour plot for visualizing targets over a 2D input space slice.

    Creates filled or line contour plots showing how a target function
    (model predictions, acquisition function, etc.) varies over two input
    dimensions while keeping other dimensions fixed.

    Parameters
    ----------
    results : BOResults
        Bayesian optimization results.
    target : {"model_mean", "model_var", "model_std", "acqfn"} or Callable
        Target function to visualize.
    itr : int, default=-1
        Iteration number for model reconstruction.
    active_input_dims : ArrayLike, default=(0, 1)
        Two input dimensions to vary in the contour plot.
    active_target_dim : int, default=0
        Target output dimension to visualize.
    fixed_inputs : {"min", "mid", "next_acq"} or dict or None, default="min"
        How to fix inactive input dimensions.
    grid_pts : tuple, default=(50, 50)
        Number of grid points in each dimension.
    fig : Figure, optional
        Existing matplotlib figure.
    ax : Axes, optional
        Existing matplotlib axes.
    """
    NUM_AXIS = 2

    def __init__(
        self,
        results: BOResults,
        target: Literal["model_mean", "model_var", "model_std", "acqfn"] | Callable,
        itr: int = -1,
        active_input_dims: ArrayLike = (0, 1),
        active_target_dim: int = 0,
        fixed_inputs: (
            Literal["min", "mid", "next_acq"] | dict[int, float] | None
        ) = "min",
        grid_pts: tuple = (50, 50),
        fig: Figure | None = None,
        ax: Axes | None = None,
    ) -> None:
        super().__init__(
            results,
            target,
            active_input_dims,
            itr=itr,
            active_target_dims=active_target_dim,
            fixed_inputs=fixed_inputs,
            grid_pts=grid_pts,
            fig=fig,
            ax=ax,
        )

    def plot(
        self,
        xlabel: str = "auto",
        ylabel: str = "auto",
        filled: bool = True,
        title: str = "",
        cbar: bool = True,
        cbar_kws: dict[str, Any] | None = None,
        legend: bool = True,
        legend_loc: LegendLocationLiteral = "outside top",
        **cont_kwargs: Any,
    ) -> None:
        """
        Create the contour plot.

        Parameters
        ----------
        xlabel : str, default="auto"
            X-axis label. If "auto", uses dimension index notation.
        ylabel : str, default="auto"
            Y-axis label. If "auto", uses dimension index notation.
        filled : bool, default=True
            Whether to create filled contours (True) or line contours (False).
        title : str, default=""
            Plot title.
        cbar : bool, default=True
            Whether to add a colorbar.
        cbar_kws : dict, optional
            Keyword arguments for colorbar customization.
        legend : bool, default=True
            Whether to display legend for added elements.
        legend_loc : str, default="outside top"
            Legend location.
        **cont_kwargs : Any
            Additional keyword arguments passed to contour/contourf.
        """
        cbar_kws = cbar_kws or {"label": self.target_label}
        cont_kwargs = cont_kwargs or {}
        ax = self.ax
        X1, X2 = self.mesh.grid

        plot_func = ax.contourf if filled else ax.contour
        target = self.target_values[self.active_target_dims[0]]
        conts = plot_func(X1, X2, target, **cont_kwargs)

        ax.set_xlim(X1.min(), X1.max())
        ax.set_ylim(X2.min(), X2.max())

        if xlabel == "auto":
            xlabel = "$x_%i$" % (self.active_input_dims[0] + 1)
        if ylabel == "auto":
            ylabel = "$x_%i$" % (self.active_input_dims[1] + 1)

        ax.set_xlabel(xlabel)
        ax.set_ylabel(ylabel)

        if title:
            ax.set_title(title)

        if cbar:
            self.colorbar = self.fig.colorbar(conts, **cbar_kws)

        for element in self.elements:
            element.plot()

        if legend:
            handles, labels = ax.get_legend_handles_labels()
            if len(labels) > 0:
                tools.set_legend(ax=ax, loc=legend_loc)


class Surface(MeshGraphic["Surface"]):
    """
    3D surface plot for visualizing targets over a 2D input space slice.

    Creates 3D surface plots showing how a target function varies over two
    input dimensions, with the function value displayed as the third dimension.
    Other input dimensions are held fixed.

    Parameters
    ----------
    results : BOResults
        Bayesian optimization results.
    target : {"model_mean", "model_var", "model_std", "acqfn"} or Callable
        Target function to visualize.
    itr : int, default=-1
        Iteration number for model reconstruction.
    active_input_dims : ArrayLike, default=(0, 1)
        Two input dimensions to vary in the surface plot.
    active_target_dim : int, default=0
        Target output dimension to visualize.
    fixed_inputs : {"min", "mid", "next_acq"} or dict or None, default="min"
        How to fix inactive input dimensions.
    grid_pts : tuple, default=(50, 50)
        Number of grid points in each dimension.
    fig : Figure, optional
        Existing matplotlib figure.
    ax : Axes, optional
        Existing matplotlib 3D axes.
    """
    NUM_AXIS = 3

    def __init__(
        self,
        results: BOResults,
        target: Literal["model_mean", "model_var", "model_std", "acqfn"] | Callable,
        itr: int = -1,
        active_input_dims: ArrayLike = (0, 1),
        active_target_dim: int = 0,
        fixed_inputs: (
            Literal["min", "mid", "next_acq"] | dict[int, float] | None
        ) = "min",
        grid_pts: tuple = (50, 50),
        fig: Figure | None = None,
        ax: Axes | None = None,
    ) -> None:
        super().__init__(
            results,
            target,
            active_input_dims,
            itr=itr,
            active_target_dims=active_target_dim,
            fixed_inputs=fixed_inputs,
            grid_pts=grid_pts,
            fig=fig,
            ax=ax,
        )

    def plot(
        self,
        xlabel: str = "auto",
        ylabel: str = "auto",
        zlabel: str = "auto",
        title: str = "",
        cbar: bool = False,
        cbar_kws: dict[str, Any] | None = None,
        legend: bool = True,
        legend_loc: LegendLocationLiteral = "outside top",
        **surf_kwargs: Any,
    ) -> None:
        """
        Create the 3D surface plot.

        Parameters
        ----------
        xlabel : str, default="auto"
            X-axis label. If "auto", uses dimension index notation.
        ylabel : str, default="auto"
            Y-axis label. If "auto", uses dimension index notation.
        zlabel : str, default="auto"
            Z-axis label. If "auto", uses target label.
        title : str, default=""
            Plot title.
        cbar : bool, default=False
            Whether to add a colorbar.
        cbar_kws : dict, optional
            Keyword arguments for colorbar customization.
        legend : bool, default=True
            Whether to display legend for added elements.
        legend_loc : str, default="outside top"
            Legend location.
        **surf_kwargs : Any
            Additional keyword arguments passed to plot_surface.
        """
        cbar_kws = cbar_kws or {"label": self.target_label}
        surf_kwargs = surf_kwargs or {"cmap": "viridis"}
        ax = self.ax
        X1, X2 = self.mesh.grid

        target_vals = self.target_values[self.active_target_dims[0]]

        surf = ax.plot_surface(X1, X2, target_vals, **surf_kwargs)

        ax.set_xlim(X1.min(), X1.max())
        ax.set_ylim(X2.min(), X2.max())

        if xlabel == "auto":
            xlabel = "$x_%i$" % (self.active_input_dims[0] + 1)
        if ylabel == "auto":
            ylabel = "$x_%i$" % (self.active_input_dims[1] + 1)
        if zlabel == "auto":
            zlabel = self.target_label

        ax.set_xlabel(xlabel)
        ax.set_ylabel(ylabel)
        ax.set_zlabel(zlabel)

        if title:
            ax.set_title(title)

        if cbar:
            self.colorbar = self.fig.colorbar(surf, **cbar_kws)

        for element in self.elements:
            element.plot()

        if legend:
            handles, labels = ax.get_legend_handles_labels()
            if len(labels) > 0:
                tools.set_legend(ax=ax, loc=legend_loc)


class Curve(MeshGraphic["Curve"]):
    """
    1D curve plot for visualizing targets along a single input dimension.

    Creates line plots showing how a target function varies along one input
    dimension while keeping all other dimensions fixed. Useful for examining
    slices through high-dimensional functions.

    Parameters
    ----------
    results : BOResults
        Bayesian optimization results.
    target : {"model_mean", "model_var", "model_std", "acqfn"} or Callable
        Target function to visualize.
    itr : int, default=-1
        Iteration number for model reconstruction.
    active_input_dims : ArrayLike, default=0
        Input dimension to vary in the curve plot.
    active_target_dim : int, default=0
        Target output dimension to visualize.
    fixed_inputs : {"min", "mid", "next_acq"} or dict or None, default="min"
        How to fix inactive input dimensions.
    grid_pts : int, default=100
        Number of points along the curve.
    fig : Figure, optional
        Existing matplotlib figure.
    ax : Axes, optional
        Existing matplotlib axes.
    """
    NUM_AXIS = 2

    def __init__(
        self,
        results: BOResults,
        target: Literal["model_mean", "model_var", "model_std", "acqfn"] | Callable,
        itr: int = -1,
        active_input_dims: ArrayLike = 0,
        active_target_dim: int = 0,
        fixed_inputs: (
            Literal["min", "mid", "next_acq"] | dict[int, float] | None
        ) = "min",
        grid_pts: int = 100,
        fig: Figure | None = None,
        ax: Axes | None = None,
    ) -> None:
        super().__init__(
            results,
            target,
            active_input_dims,
            itr=itr,
            active_target_dims=active_target_dim,
            fixed_inputs=fixed_inputs,
            grid_pts=grid_pts,
            fig=fig,
            ax=ax,
        )

    def plot(
        self,
        xlabel: str = "auto",
        ylabel: str = "auto",
        title: str = "",
        legend: bool = True,
        legend_loc: LegendLocationLiteral = "outside top",
        grid: Literal["seaborn", "default"] | None = None,
        **kwargs: Any,
    ) -> None:
        """
        Create the 1D curve plot.

        Parameters
        ----------
        xlabel : str, default="auto"
            X-axis label. If "auto", uses dimension index notation.
        ylabel : str, default="auto"
            Y-axis label. If "auto", uses target label.
        title : str, default=""
            Plot title.
        legend : bool, default=True
            Whether to display legend for added elements.
        legend_loc : str, default="outside top"
            Legend location.
        grid : {"seaborn", "default"} or None, optional
            Grid style to apply.
        **kwargs : Any
            Additional keyword arguments passed to plot.
        """
        kwargs = kwargs or {}
        ax = self.ax
        (x,) = self.mesh.grid
        y = self.target_values[self.active_target_dims[0]]
        ax.plot(x, y, **kwargs)
        if xlabel == "auto":
            xlabel = "$x_%i$" % (self.active_input_dims[0] + 1)
        if ylabel == "auto":
            ylabel = self.target_label
        ax.set_xlabel(xlabel)
        ax.set_ylabel(ylabel)
        if title:
            ax.set_title(title)

        for element in self.elements:
            element.plot()

        if grid:
            tools.set_grid(ax, seaborn_style=(grid == "seaborn"))

        if legend:
            handles, labels = ax.get_legend_handles_labels()
            if len(labels) > 0:
                tools.set_legend(ax=ax, loc=legend_loc)
