"""High-level plotting interface for Bayesian optimization results.

This module provides the `Plotter` class, which serves as the main plotting
interface for post-processing BOSS optimization results. It is typically used
by the `PPMain` object (defined in `boss/pp/pp_main.py`) to generate various
diagnostic and analysis plots.

The module builds on top of the lower-level graphics infrastructure:
- `boss.pp.graphics`: Provides fundamental plotting primitives (`Curve`, `Contour`,
  `Surface`) that handle mesh-based visualization of functions over input spaces.
- `boss.pp.elements`: Defines composable plot elements (`Acquisitions`, `Minimum`,
  `NextAcquisition`) that can be overlaid on graphics.
- `boss.pp.mesh`: Handles coordinate grid generation and function evaluation.

The `Plotter` class encapsulates common plotting workflows and provides convenient
methods for visualizing:
- Acquisition history and convergence tracking
- Gaussian process models and uncertainty
- Acquisition functions
- True objective functions (when available)
- Hyperparameter evolution

Typical Usage
-------------
The `Plotter` is usually created and used via `PPMain`:

>>> from boss.pp.pp_main import PPMain
>>> pp = PPMain(results, pp_models=True, pp_acq_funcs=True)
>>> pp.run()

Direct usage is also supported:

>>> from boss.pp.plotter import Plotter
>>> plotter = Plotter.from_grid_spec(results)
>>> plotter.plot_model(itr=-1, show=True)
>>> plotter.plot_acquisitions(save_as="acquisitions.png")
"""

from __future__ import annotations

import itertools
from pathlib import Path
from typing import TYPE_CHECKING, Literal

import matplotlib.pyplot as plt
import numpy as np
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.bo.results import BOResults
from boss.pp.elements import Acquisitions, Minimum, NextAcquisition
from boss.pp.graphics import Contour, Curve, MeshGraphic
from boss.pp.mesh import Mesh

Tableau20 = style.Tableau20


class Plotter:
    """High-level interface for creating plots from Bayesian optimization results.

    The `Plotter` class provides convenient methods for generating various diagnostic
    and analysis plots from a `BOResults` object. It uses `Mesh` objects to define
    coordinate grids and relies on the graphics primitives from `boss.pp.graphics`
    to create visualizations.

    Parameters
    ----------
    results : BOResults
        Bayesian optimization results object containing models, data, and settings.
    mesh : Mesh
        Mesh object defining the coordinate grid for plotting.
    style : str or Path or dict or None, optional
        Matplotlib style to use for plots. Can be a style name, path to style file,
        or dictionary of style parameters.

    Attributes
    ----------
    results : BOResults
        The BO results object.
    settings : Settings
        Settings object from the BO results.
    bounds : NDArray
        Optimization bounds from settings.
    colormap : str
        Default colormap for plots.
    mesh : Mesh
        The mesh object for coordinate grids.
    grid_spec : ArrayLike or None
        Grid specification from the mesh, if available.

    See Also
    --------
    boss.pp.pp_main.PPMain : Main post-processing interface
    boss.pp.graphics.MeshGraphic : Base class for mesh-based graphics
    boss.pp.mesh.Mesh : Coordinate grid generation
    """

    def __init__(
        self, results: BOResults, mesh: Mesh, style: str | Path | dict | None = None
    ) -> None:
        if style:
            plt.style.use(style)
        self.results = results
        self.settings = results.settings
        self.bounds = self.settings["bounds"]
        self.colormap = "viridis"
        self.mesh = mesh
        try:
            self.grid_spec = mesh.get_grid_spec()
        except RuntimeError:
            self.grid_spec = None

    @classmethod
    def from_grid_spec(
        cls,
        results: BOResults,
        grid_spec: ArrayLike | None = None,
        style: str | Path | dict | None = None,
    ):
        """Create a Plotter from a grid specification.

        This is a convenience constructor that creates a `Mesh` object from a grid
        specification before initializing the `Plotter`.

        Parameters
        ----------
        results : BOResults
            Bayesian optimization results object.
        grid_spec : ArrayLike or None, optional
            Grid specification for the mesh. If None, uses the grid spec from
            `results.settings["pp_model_slice"]`.
        style : str or Path or dict or None, optional
            Matplotlib style to use for plots.

        Returns
        -------
        Plotter
            Initialized plotter instance.
        """
        if grid_spec is None:
            grid_spec = results.settings["pp_model_slice"]
        mesh = Mesh.from_grid_spec(results.settings["bounds"], grid_spec)

        self = cls(results, mesh, style)
        return self

    def plot_acquisitions(
        self, save_as: str | Path | None = None, show: bool = True
    ) -> tuple[Figure, NDArray]:
        """Plot acquisition history over iterations.

        Creates a two-panel plot showing:
        - Top panel: Predicted minimum value with uncertainty bands and acquired y-values
        - Bottom panel: Acquisition locations and predicted minimum locations in input space

        Parameters
        ----------
        save_as : str or Path or None, optional
            Path to save the figure. If None, figure is not saved.
        show : bool, default=True
            Whether to display the figure.

        Returns
        -------
        fig : Figure
            The matplotlib figure object.
        ax_arr : NDArray
            Array of axes objects [ax1, ax2].
        """
        results = self.results
        iters_min = list(results["x_glmin"].keys())
        iters = results.batch_tracker.iteration_labels

        mu_glmin = results["mu_glmin"].to_array()
        nu_glmin = results["nu_glmin"].to_array()
        x_glmin = results["x_glmin"].to_array()

        colors = itertools.cycle(Tableau20)
        xticks = iters_min[:: max(1, round(len(iters_min) / 10))]
        fig, ax_arr = plt.subplots(
            2, 1, figsize=tools.scale_figsize(1.5, 1), sharex=True
        )
        ax1, ax2 = ax_arr

        ax1.fill_between(
            iters_min, mu_glmin + nu_glmin, mu_glmin, color=Tableau20.lightgrey
        )
        ax1.fill_between(
            iters_min, mu_glmin - nu_glmin, mu_glmin, color=Tableau20.lightgrey
        )
        ax1.plot(iters_min, mu_glmin + nu_glmin, Tableau20.grey)
        ax1.plot(iters_min, mu_glmin - nu_glmin, Tableau20.grey)
        ax1.plot(iters_min, mu_glmin, "black", label=r"$\mu(\hat{x})$")

        ax1.scatter(
            iters,
            results["Y"][:, 0],
            facecolors="none",
            edgecolors=next(colors),
            label=r"$acq_y$",
        )
        ax1.set_xlim(min(iters_min), max(iters_min))
        ax1.set_xticks(xticks)
        ax1.set_ylabel(r"$y$ with $\mu(\hat{x})$")

        for i in range(results.settings.dim):
            ax2.scatter(
                iters,
                results["X"][:, i],
                facecolors="none",
                edgecolors=next(colors),
                label=r"$acq_{x%i}$" % (i + 1),
            )

        for i in range(results.settings.dim):
            ax2.plot(
                iters_min,
                x_glmin[:, i],
                color=next(colors),
                label=r"$\hat{x}_%i$" % (i + 1),
            )

        ax2.set_xlim(iters_min[0], iters_min[-1])
        ax2.set_xticks(xticks)
        ax2.set_xlabel("Iteration")
        ax2.set_ylabel(r"$x_i$ with $\hat{x}_i$")
        fig.suptitle("Acquisitions")
        # --- unified save/show/close logic ---
        if save_as is not None:
            save_as = Path(save_as)
            fig.savefig(save_as, bbox_inches="tight")

        if show:
            fig.show()
        else:
            plt.close(fig)  # avoid auto-render in notebooks

        return fig, ax_arr

    def plot_convergence(
        self, save_as: str | Path | None = None, show: bool = True
    ) -> tuple[Figure, NDArray]:
        """Plot convergence diagnostics over iterations.

        Creates a two-panel plot showing convergence metrics:
        - Top panel: Change in predicted minimum location (Euclidean norm)
        - Bottom panel: Normalized change in predicted minimum value

        These metrics help assess whether the optimization has converged.

        Parameters
        ----------
        save_as : str or Path or None, optional
            Path to save the figure. If None, figure is not saved.
        show : bool, default=True
            Whether to display the figure.

        Returns
        -------
        fig : Figure
            The matplotlib figure object.
        ax_arr : NDArray
            Array of axes objects [ax1, ax2].
        """
        results = self.results
        iters_min = list(results["x_glmin"].keys())[1:]
        x_glmin = results["x_glmin"].to_array()
        mu_glmin = results["mu_glmin"].to_array()
        dmu_glmin = np.abs(np.diff(mu_glmin))
        dx_glmin = np.linalg.norm(np.diff(x_glmin, axis=0), axis=1)
        dy = np.zeros_like(dmu_glmin)
        for ind, itr in enumerate(iters_min):
            y_min, y_max = results.get_est_yrange(itr)
            dy[ind] = y_max - y_min

        fig, ax_arr = plt.subplots(
            2, 1, figsize=tools.scale_figsize(1.5, 1), sharex=True
        )
        ax1, ax2 = ax_arr
        xticks = iters_min[:: max(1, round(len(iters_min) / 10))]

        # dx_glmin
        ax1.plot(iters_min, dx_glmin, color=Tableau20.blue)
        ax1.set_ylabel(r"$\left|\Delta \hat{x}\right|$")
        ax1.set_xticks(xticks)
        tools.set_conditional_logscale(ax1, dx_glmin)

        # dmu_glmin
        ax2.plot(iters_min, dmu_glmin / dy, color=Tableau20.red)
        ax2.set_xlabel("Iteration")
        ax2.set_ylabel(r"$\left|\Delta \mu(\hat{x})\right|/\Delta y$")
        ax2.set_xticks(xticks)
        tools.set_conditional_logscale(ax2, dx_glmin)
        plt.suptitle("Convergence tracking")
        # ----- Unified save/show/close pattern -----
        if save_as is not None:
            save_as = Path(save_as)
            fig.savefig(save_as, bbox_inches="tight")

        if show:
            fig.show()
        else:
            plt.close(fig)  # prevents Jupyter inline auto-display

        return fig, ax_arr

    def plot_hyperparameters(
        self,
        save_as: str | Path | None = None,
        show: bool = True,
    ) -> tuple[Figure, NDArray]:
        """Plot Gaussian process hyperparameter evolution.

        Creates a two-panel plot showing how GP kernel hyperparameters change
        during optimization:
        - Top panel: Kernel variance over iterations
        - Bottom panel: Kernel lengthscales for each dimension over iterations

        Only plots hyperparameters that were not fixed during optimization.

        Parameters
        ----------
        save_as : str or Path or None, optional
            Path to save the figure. If None, figure is not saved.
        show : bool, default=True
            Whether to display the figure.

        Returns
        -------
        fig : Figure
            The matplotlib figure object.
        ax_arr : NDArray
            Array of axes objects [ax1, ax2].
        """
        results = self.results
        model_params = results["model_params"]
        iters = np.array(list(model_params.keys()))
        params = model_params.to_array(gaps=False)
        variances = params[:, 0]
        lengthscales = params[:, 1:]

        fig, ax_arr = plt.subplots(2, 1, figsize=tools.scale_figsize(2, 1), sharex=True)
        ax1, ax2 = ax_arr
        colors = itertools.cycle(Tableau20)
        markers = itertools.cycle(["s", "p", "*", "h", "^", "+", "x", "d", "v", "|"])

        # Variance
        ax1.plot(iters, variances, color=next(colors))
        ax1.set_ylabel("Variance")
        ax1.set_xticks(iters[:: max(1, round(len(iters) / 10))])
        tools.set_conditional_logscale(ax1, variances)

        # Lengthscales
        for i in range(results.settings.dim):
            ax2.plot(
                iters,
                lengthscales[:, i],
                color=next(colors),
                marker=next(markers),
                label=r"$\ell_ %i$" % (i + 1),
            )
        ax2.set_xlabel("Iteration")
        ax2.set_ylabel("Lengthscale")
        ax2.set_xticks(iters[:: max(1, round(len(iters) / 10))])
        tools.set_conditional_logscale(ax1, lengthscales.flatten())

        tools.set_legend(loc="best")
        plt.suptitle("Hyperparameter values")
        # --- NEW: consistent save/show/close logic ---
        if save_as is not None:
            save_as = Path(save_as)
            fig.savefig(save_as, bbox_inches="tight")

        if show:
            fig.show()
        else:
            plt.close(fig)   # prevents auto-display in notebooks

        return fig, ax_arr

    def plot_model(
        self,
        itr: int = -1,
        include_stdev: bool = True,
        save_as: str | Path | None = None,
        show: bool = True,
    ) -> list[MeshGraphic]:
        """Plot the Gaussian process model predictions.

        Creates visualizations of the GP model's mean predictions over the input
        space. For 1D problems, generates a curve plot; for higher dimensions,
        generates contour plots. Optionally includes uncertainty (standard deviation)
        visualizations.

        The plot includes overlays showing:
        - Acquisition points
        - Predicted global minimum
        - Next acquisition point (if available)

        Parameters
        ----------
        itr : int, default=-1
            Iteration number to plot. Negative indices count from the end.
        include_stdev : bool, default=True
            Whether to include uncertainty visualization. For 1D plots, shows
            uncertainty bands. For 2D plots, creates a separate contour plot.
        save_as : str or Path or None, optional
            Path to save the figure. If None, figure is not saved.
        show : bool, default=True
            Whether to display the figure.

        Returns
        -------
        list[MeshGraphic]
            List of graphic objects created (Curve or Contour instances).
            Contains 1-2 graphics depending on `include_stdev`.
        """

        model = self.results.reconstruct_model(itr=itr)
        mesh_dim = len(self.mesh.active_dims)
        bo_dim = len(self.mesh.bounds)

        # set defaults for fixed variables
        if var_defaults := self.settings["pp_var_defaults"]:
            dim_vals = {i: v for i, v in enumerate(var_defaults)}
            for i in self.mesh.active_dims:
                del dim_vals[i]
            self.mesh.fix_dims(dim_vals)
        else:
            self.mesh.fix_dim_preset(self.results, "min", itr=itr)

        Y, Y_var = self.mesh.evaluate_func(model.predict)
        Y_stdev = np.sqrt(Y_var)
        graphics = []

        if mesh_dim == 1:
            curve = Curve.from_mesh(
                results=self.results, target=Y, mesh=self.mesh, itr=itr
            )
            if bo_dim == 1:
                curve += Minimum(
                    s=30, marker="o", color=Tableau20.pink, label="Global min."
                )
                curve += Acquisitions(
                    s=30, marker="x", color=Tableau20.red, label="Acq."
                )

                X = self.results.get_next_acq(itr)
                curve.ax.axvline(x=X.item(), color=Tableau20.brown, label="Next acq.")
            else:
                # if we have fixed inputs to min the global min will
                # be on the 1d curve so we can plot it
                if curve.fixed_inputs == "min":
                    curve += Minimum(
                        s=30, marker="o", color=Tableau20.pink, label="Global min."
                    )

            if include_stdev:
                ax = curve.ax
                X = self.mesh.grid[0]
                ax.fill_between(X, Y + Y_stdev, Y, color=Tableau20.lightblue, alpha=0.7)
                ax.fill_between(X, Y - Y_stdev, Y, color=Tableau20.lightblue, alpha=0.7)

            curve.plot(
                grid="seaborn",
                ylabel=get_target_label("model_mean"),
                color=Tableau20.blue,
            )
            graphics.append(curve)

        else:
            cont = Contour.from_mesh(
                results=self.results,
                target=Y,
                mesh=self.mesh,
                itr=itr,
                fixed_inputs="min",
            )
            cont += Acquisitions(
                s=40, marker="x", color=Tableau20.red, label="Acquisitions"
            )
            cont += Minimum(s=40, marker="o", color=Tableau20.pink, label="Minimum")
            cont += NextAcquisition(
                s=40, marker="s", color=Tableau20.brown, label="Next Acquisition"
            )
            cont.plot()

            graphics.append(cont)
            if include_stdev:
                cont_stdev = Contour.from_mesh(
                    results=self.results,
                    target=Y_stdev,
                    mesh=self.mesh,
                    itr=itr,
                    fixed_inputs="min",
                )
                cont_stdev += Acquisitions(
                    s=40, marker="x", color=Tableau20.red, label="Acquisitions"
                )
                cont_stdev += Minimum(
                    s=40, marker="o", color=Tableau20.pink, label="Minimum"
                )
                cont_stdev += NextAcquisition(
                    s=40, marker="s", color=Tableau20.brown, label="Next Acquisition"
                )
                cont_stdev.plot()
                graphics.append(cont_stdev)
        # ---- handle all figures from all graphics ----
        figs: list[Figure] = []
        for g in graphics:
            if hasattr(g, "ax"):
                fig = g.ax.figure
                if fig not in figs:
                    figs.append(fig)

        # Saving
        if save_as is not None and figs:
            save_as = Path(save_as)
            # save first figure as requested
            figs[0].savefig(save_as, bbox_inches="tight")

            # if we have extra figures (e.g., stdev), save them with suffixes
            for idx, extra_fig in enumerate(figs[1:], start=1):
                # e.g. model_graph_0001_stdev.png
                extra_name = f"{save_as.stem}_stdev{idx}{save_as.suffix}"
                extra_path = save_as.with_name(extra_name)
                extra_fig.savefig(extra_path, bbox_inches="tight")

        # Showing / closing
        if show:
            for fig in figs:
                fig.show()
        else:
            for fig in figs:
                plt.close(fig)

        return graphics

    def plot_acq_func(
        self,
        itr=-1,
        save_as=None,
        show=True,
    ) -> MeshGraphic:
        """Plot the acquisition function.

        Visualizes the acquisition function over the input space at a specified
        iteration. For 1D problems, creates a curve plot; for higher dimensions,
        creates a contour plot. The plot includes overlays showing acquisition
        points, predicted minimum, and next acquisition point.

        Parameters
        ----------
        itr : int, default=-1
            Iteration number to plot. Negative indices count from the end.
        save_as : str or Path or None, optional
            Path to save the figure. If None, figure is not saved.
        show : bool, default=True
            Whether to display the figure.

        Returns
        -------
        MeshGraphic
            The graphic object created (Curve or Contour instance).
        """
        mesh_dim = len(self.mesh.active_dims)

        # set defaults for fixed variables
        if var_defaults := self.settings["pp_var_defaults"]:
            dim_vals = {i: v for i, v in enumerate(var_defaults)}
            for i in self.mesh.active_dims:
                del dim_vals[i]
            self.mesh.fix_dims(dim_vals)
        else:
            self.mesh.fix_dim_preset(self.results, "min", itr=itr)

        if mesh_dim == 1:
            graphic = Curve.from_mesh(
                results=self.results, target="acqfn", mesh=self.mesh, itr=itr
            )
            graphic.plot(
                grid="seaborn", ylabel=get_target_label("acqfn"), color=Tableau20.red
            )
        else:
            graphic = Contour.from_mesh(
                results=self.results,
                target="acqfn",
                mesh=self.mesh,
                itr=itr,
                fixed_inputs="min",
            )
            graphic += Acquisitions(
                s=40, marker="x", color=Tableau20.red, label="Acquisitions"
            )
            graphic += Minimum(s=40, marker="o", color=Tableau20.pink, label="Minimum")
            graphic += NextAcquisition(
                s=40, marker="s", color=Tableau20.brown, label="Next Acquisition"
            )
            graphic.plot()

        # get the figure from the graphic's axes
        fig = graphic.ax.figure if hasattr(graphic, "ax") else None

        if save_as is not None and fig is not None:
            save_as = Path(save_as)
            fig.savefig(save_as, bbox_inches="tight")

        if show and fig is not None:
            # fig.show() is a bit more explicit than global plt.show()
            fig.show()
        else:
            # close so Jupyter / inline backends don't auto-display
            if fig is not None:
                plt.close(fig)

        return graphic

    def plot_true_func(self, save_as: str | Path | None = None, show: bool = True):
        """Plot the true objective function.

        Visualizes the true objective function over the input space when it is
        available. For 1D problems, creates a curve plot; for higher dimensions,
        creates a contour plot. The plot includes overlays showing acquisition
        points and the predicted minimum.

        This is useful for comparing the GP model predictions against ground truth.

        Parameters
        ----------
        save_as : str or Path or None, optional
            Path to save the figure. If None, figure is not saved.
        show : bool, default=True
            Whether to display the figure.

        Returns
        -------
        MeshGraphic
            The graphic object created (Curve or Contour instance).
        """
        mesh_dim = len(self.mesh.active_dims)

        # set defaults for fixed variables
        if var_defaults := self.settings["pp_var_defaults"]:
            dim_vals = {i: v for i, v in enumerate(var_defaults)}
            for i in self.mesh.active_dims:
                del dim_vals[i]
            self.mesh.fix_dims(dim_vals)
        else:
            self.mesh.fix_dim_preset(self.results, "min", itr=-1)

        f_true = lambda X: self.settings.f.evaluate(X).Y
        Y = self.mesh.evaluate_func(f_true)

        if mesh_dim == 1:
            graphic = Curve.from_mesh(
                results=self.results,
                target=Y,
                mesh=self.mesh,
            )
            graphic.plot(
                grid="seaborn",
                ylabel="y",
                color=Tableau20.blue,
            )
        else:
            graphic = Contour.from_mesh(
                results=self.results,
                target=Y,
                mesh=self.mesh,
                fixed_inputs="min",
            )
            graphic += Acquisitions(
                s=40, marker="x", color=Tableau20.red, label="Acquisitions"
            )
            graphic += Minimum(
                s=40, marker="o", color=Tableau20.pink, label="Pred. Minimum"
            )
            graphic += NextAcquisition(
                s=40, marker="s", color=Tableau20.brown, label="Next Acquisition"
            )
            graphic.plot()

        # pick the figure from the first graphic (Curve or Contour)
        fig = None
        if graphic is not None and hasattr(graphic, "ax"):
            fig = graphic.ax.figure
        
        if save_as is not None and fig is not None:
            save_as = Path(save_as)
            fig.savefig(save_as, bbox_inches="tight")
        
        if show and fig is not None:
            fig.show()
        else:
            if fig is not None:
                plt.close(fig)
        
        return graphic


def get_target_label(
    target: Literal["model_mean", "model_var", "model_std", "acqfn"],
) -> str:
    """Get LaTeX label for a target function type.

    Parameters
    ----------
    target : {"model_mean", "model_var", "model_std", "acqfn"}
        Target function type.

    Returns
    -------
    str
        LaTeX-formatted label string for the target.
    """
    if target == "model_mean":
        label = r"$\mu$"
    elif target == "model_var":
        label = r"$\nu^2$"
    elif target == "model_std":
        label = r"$\nu$"
    elif target == "acqfn":
        label = r"$\alpha$"
    return label
