from __future__ import annotations

__all__ = ["plot_observables"]

from collections.abc import Mapping, Generator
import itertools

from matplotlib.axes import Axes
from matplotlib.lines import Line2D

from ..lattice import Lattice, Refpts, All
from ..latticetools import ElementObservable, ObservableList
from .synopt import plot_synopt

_SLICES = 400


class _RingSplitter:
    """Generate a ring split in small slices for plotting."""

    def __init__(self, ring: Lattice, s_range: tuple[float, float], slices: int):
        ring.s_range = s_range
        self._ring0 = ring
        self._ring1 = None
        self.s = None
        self.slices = slices

    @property
    def ring(self) -> Lattice:
        if self._ring1:
            return self._ring1
        else:
            rg = self._ring0.slice(slices=self.slices)
            self.s = rg.get_s_pos(range(len(rg) + 1))
            self._ring1 = rg
            return rg


def plot_observables(
    ring: Lattice,
    obsleft: ObservableList,
    *obsright: ObservableList,
    s_range: tuple[float, float] | None = None,
    axes: Axes | None = None,
    xlabel: str = "",
    ylabel: str = "",
    title: str = "",
    slices: int = _SLICES,
    dipole: dict | None = {},  # noqa: B006
    quadrupole: dict | None = {},  # noqa: B006
    sextupole: dict | None = {},  # noqa: B006
    multipole: dict | None = {},  # noqa: B006
    monitor: dict | None = {},  # noqa: B006
    labels: Refpts = None,
    **kwargs,
) -> tuple[Axes]:
    # noinspection PyUnresolvedReferences
    r"""Plot element observables along a lattice.

    Args:
        ring:       Lattice description
        obsleft:    List of :py:class:`.ElementObservable` plotted against the left
          axis. if refpts is :py:obj:`.All`, a line is drawn. Otherwise, markers are
          drawn. It is recommended to use Observables with scalar values. Otherwise,
          all the values are plotted but share the same line properties and legend,
        obsright:   Optional list of :py:class:`.ElementObservable` plotted against the
          right axis,
        axes:       :py:class:`~.matplotlib.axes.Axes` in which the observables are
          plotted. if :py:obj:`None`, a new figure is created,
        s_range:    Lattice range of interest, default: whole lattice,
        slices:     Number of lattice slices for getting smooth curves. Default: 400.
        xlabel:     x-axis label. Default: ``s [m]``.
        ylabel:     y-axis label. Default: :py:attr:`.ObservableList.axis_label`.
        title:      Plot title,

    The following keywords are transmitted to the :py:func:`.plot_synopt` function.They
    apply to the main (left) axis and are ignored when plotting in exising axes:

    Keyword Args:
        labels (Refpts):    Select elements for which the name is displayed.
          Default: :py:obj:`None`,
        dipole (dict):      Dictionary of properties overloading the default
          properties of the dipole representation.
          Example: :pycode:`{"facecolor": "xkcd:electric blue"}`. If :py:obj:`None`,
          dipoles are not shown.
        quadrupole (dict):  Same definition as for dipole,
        sextupole (dict):   Same definition as for dipole,
        multipole (dict):   Same definition as for dipole,
        monitor (dict):     Same definition as for dipole.

    Returns:
        axes: tuple of :py:class:`~.matplotlib.axes.Axes`. Contains 2 elements if there
          is a plot on the right y-axis, 1 element otherwise.

    .. note::

        The legend of the plot is controlled by the :py:attr:`.Observable.label`
        attributes. Default values are provided, but labels may explicitly set.

        Labels may contain LaTeX math code. Example: ``"$\beta_x$"`` will appear as
        ":math:`\beta_x`".

        Labels starting with an underscore do not appear in the legend.

    Examples:
        Minimal example using default values:

        >>> obsmu = at.ObservableList(
        ...     [
        ...         at.LocalOpticsObservable(at.All, "mu", plane=0),
        ...         at.LocalOpticsObservable(at.All, "mu", plane=1),
        ...     ]
        ... )
        >>>
        >>> (ax1,) = at.plot_observables(ring, obsmu)

        .. image:: /images/plot_observables/phase_obs.*
           :alt: phase advance plot

        This example demonstrates the use of the postfun post-processing attribute of
        observables to plot the beam envelopes for arbitrary emitttance values:

        >>> # Define the emittances
        >>> emit_x = 130.0e-12
        >>> emit_y = 10.0e-12
        >>>
        >>> # beam size computation
        >>> sigma_x = lambda x: 1.0e6 * np.sqrt(x * emit_x)  # result in um
        >>> sigma_y = lambda y: 1.0e6 * np.sqrt(y * emit_y)  # result in um
        >>>
        >>> # Observables
        >>> obsenv = at.ObservableList(
        ...     [
        ...         at.LocalOpticsObservable(
        ...             at.All, "beta", plane="x", postfun=sigma_x, label=r"$\sigma_x$"
        ...         ),
        ...         at.LocalOpticsObservable(
        ...             at.All, "beta", plane="y", postfun=sigma_y, label=r"$\sigma_y$"
        ...         ),
        ...     ]
        ... )
        >>> (ax2,) = at.plot_observables(
        ...     ring, obsenv, ylabel=r"Beam size [${\mu}m$]", title="Beam envelopes"
        ... )
        >>> )

        .. image:: /images/plot_observables/envelope_obs.*
           :alt: envelope plot
    """

    def get_format(obs: ElementObservable, default: str) -> tuple[tuple, Mapping]:
        fmt = getattr(obs, "plot_fmt", default)
        if isinstance(fmt, Mapping):
            return (), fmt
        else:
            return (fmt,), {}

    def axes1(axes: Axes, obs: ObservableList) -> Generator[Line2D, None, None]:
        """Plot all observables on a given axis."""

        def plot1(obs: ElementObservable):
            """Plot a single observable."""
            if obs.refpts is All:
                args, kwargs = get_format(obs, f"C{next(curve_counter)}")
                x = curve_s
            else:
                args, kwargs = get_format(obs, f"oC{next(dot_counter)}")
                x = dot_s[obs._boolrefs]

            return axes.plot(x, obs.value, *args, label=obs.label, **kwargs)

        axes.set_ylabel(obs.axis_label)
        for ob in obs:
            yield from plot1(ob)

    def evaluate(obs: ObservableList) -> None:
        """Evaluates one observable."""
        curves = ObservableList(**obs.eval_kw)
        dots = ObservableList(**obs.eval_kw)

        for ob in obs:
            if not isinstance(ob, ElementObservable):
                msg = f"{ob.name} is not an ElementObservable object."
                raise ValueError(msg)
            elif ob.refpts is All:
                curves.append(ob)
            else:
                dots.append(ob)

        if curves:
            # Evaluate curve data
            curves.evaluate(ring=splitter.ring, **kwargs)

        if dots:
            # Evaluate marker data
            dots.evaluate(ring=ring, **kwargs)

    splitter = _RingSplitter(ring, s_range, slices)

    if axes is None:
        _, axleft = plot_synopt(
            ring,
            title=title,
            dipole=dipole,
            quadrupole=quadrupole,
            sextupole=sextupole,
            multipole=multipole,
            monitor=monitor,
            labels=labels,
        )
    elif isinstance(axes, Axes):
        axleft = axes
    else:
        msg = "The 'axes' argument must be an Axes object."
        raise ValueError(msg)

    allaxes = (axleft, axleft.twinx()) if obsright else (axleft,)
    allobs = (obsleft, *obsright)

    for obs in allobs:
        evaluate(obs)

    curve_s = splitter.s
    dot_s = ring.get_s_pos(range(len(ring) + 1))
    curve_counter = itertools.count()
    dot_counter = itertools.count()

    lines = []
    for ax, obs in zip(allaxes, allobs, strict=True):
        lines.extend(axes1(ax, obs))

    axleft.set_xlabel(xlabel or "s [m]")
    if ylabel:
        axleft.set_ylabel(ylabel)
    axleft.legend(handles=[ln for ln in lines if not ln.get_label().startswith("_")])
    axleft.grid(True)
    return allaxes
