from __future__ import annotations

import abc
from typing import TYPE_CHECKING, Any

import numpy as np
from numpy.typing import ArrayLike, NDArray

if TYPE_CHECKING:
    from boss.pp.graphics import MeshGraphic

from boss.utils.sparselist import wrap_index


class Element(abc.ABC):
    """Abstract base class for all plot elements."""

    @property
    def graphic(self) -> MeshGraphic:
        """The graphic that the element is plotted on."""
        return self._graphic

    @graphic.setter
    def graphic(self, graphic: MeshGraphic) -> None:
        self._graphic = graphic

    @abc.abstractmethod
    def plot(self):
        """Plots the element on the graphic's axes."""
        pass


class Scatter(Element):
    """An element that creates a scatter plot."""

    def __init__(
        self,
        **scatter_kwargs: Any,
    ) -> None:
        """Initialize the Scatter element.

        Parameters
        ----------
        **scatter_kwargs : Any
            Keyword arguments passed to `matplotlib.axes.Axes.scatter`.
        """
        self.scatter_kwargs = scatter_kwargs

    @property
    @abc.abstractmethod
    def scatter_data(self) -> NDArray:
        """The data to be scattered."""
        pass

    def plot(self) -> None:
        """Creates the scatter plot on the graphic's axes."""
        _, ax = self.graphic.fig, self.graphic.ax
        kwargs = {'clip_on': False, 'zorder': 10}
        kwargs.update(self.scatter_kwargs)
        ax.scatter(*self.scatter_data.T, **kwargs)


class Acquisitions(Scatter):
    """An element for plotting acquired data points."""

    def __init__(
        self,
        itr: int | NDArray | None = None,
        **scatter_kwargs: Any,
    ) -> None:
        """Initialize the Acquisitions element.

        Parameters
        ----------
        itr : int | NDArray | None, optional
            The iteration(s) to plot. If None, all acquisitions up to the
            current iteration are plotted.
        **scatter_kwargs : Any
            Keyword arguments passed to `matplotlib.axes.Axes.scatter`.
        """
        self.itr = itr
        super().__init__(**scatter_kwargs)

    @property
    def scatter_data(self) -> NDArray:
        """The acquisition data to be plotted."""
        graphic = self.graphic
        # If an iteration was not specified we plot all acquisitions
        # from start to current iteration.
        if self.itr is None:
            itr_curr = wrap_index(graphic.itr, graphic.results.num_iters)
            self.itr = np.arange(itr_curr + 1) if itr_curr > 0 else 0

        results = graphic.results
        X = results.select("X", itr=self.itr)[:, graphic.active_input_dims]
        Y = results.select("Y", itr=self.itr)[:, graphic.active_target_dims]
        acq = np.hstack((X, Y))[:, : graphic.NUM_AXIS]
        return acq


class NextAcquisition(Scatter):
    """An element for plotting the next proposed acquisition point."""

    def __init__(
        self,
        **scatter_kwargs: Any,
    ) -> None:
        """Initialize the NextAcquisition element.

        Parameters
        ----------
        **scatter_kwargs : Any
            Keyword arguments passed to `matplotlib.axes.Axes.scatter`.
        """
        super().__init__(**scatter_kwargs)

    @property
    def scatter_data(self) -> NDArray:
        """The next acquisition point to be plotted."""
        graphic = self.graphic
        res = graphic.results
        X = res.get_next_acq(self.graphic.itr)[:, graphic.active_input_dims]
        return X


class Minimum(Scatter):
    """An element for plotting the inferred global minimum."""

    def __init__(
        self,
        itr: int | NDArray | None = None,
        **scatter_kwargs: Any,
    ) -> None:
        """Initialize the Minimum element.

        Parameters
        ----------
        itr : int | NDArray | None, optional
            The iteration to plot the minimum for. If None, the current
            iteration is used.
        **scatter_kwargs : Any
            Keyword arguments passed to `matplotlib.axes.Axes.scatter`.
        """
        self.itr = itr
        super().__init__(**scatter_kwargs)

    @property
    def scatter_data(self) -> NDArray:
        """The inferred global minimum to be plotted."""
        graphic = self.graphic
        results = graphic.results
        self.itr = self.itr or graphic.itr
        x_glmin = np.atleast_2d(results.select("x_glmin", itr=self.itr))
        mu_glmin = np.atleast_2d(results.select("mu_glmin", itr=self.itr)).T
        minimum = np.hstack((x_glmin[:, graphic.active_input_dims], mu_glmin))
        return minimum[:, : graphic.NUM_AXIS]


class BestAcquisition(Scatter):
    """An element for plotting the best acquisition so far."""

    def __init__(
        self,
        itr_max: int | None = None,
        **scatter_kwargs: Any,
    ) -> None:
        """Initialize the BestAcquisition element.

        Parameters
        ----------
        itr_max : int | None, optional
            The maximum iteration to consider. If None, the current
            iteration is used.
        **scatter_kwargs : Any
            Keyword arguments passed to `matplotlib.axes.Axes.scatter`.
        """
        self.itr_max = itr_max or self.graphic.itr
        super().__init__(**scatter_kwargs)

    @property
    def scatter_data(self) -> NDArray:
        """The best acquisition data to be plotted."""
        graphic = self.graphic
        res = graphic.results
        x_best, y_best = res.get_best_acq(self.itr_max)
        best = np.hstack((x_best, y_best))
        return best[:, : graphic.NUM_AXIS]
