"""
Extensions and helpers for tutorials.
"""

from __future__ import annotations

import numpy as np
import xarray as xr
from matplotlib import pyplot as plt
from numpy.typing import ArrayLike

from .config import settings
from .frame import AzimuthConvention
from .scenes.atmosphere import AbstractHeterogeneousAtmosphere
from .spectral import SpectralIndex


def plot_polarfilm(
    da: xr.DataArray,
    levels: int | ArrayLike = 16,
    show_contour: bool = True,
    show_azimuth: bool = False,
    figsize: tuple = (3, 3),
    azimuth_convention: AzimuthConvention | str | None = None,
    vmin: float | None = None,
    vmax: float | None = None,
    theta_max: float = 90.0,
):
    """
    Plot film data generated by a hemispherical distant sensor as a polar
    contour plot.

    Parameters
    ----------
    da : xarray.DataArray
        Data array to be plotted.

    levels : int or array-like, default: 16
        Number of levels on the contour plot.

    show_contour : bool, default: True
        If ``True``, show contour outlines.

    show_azimuth : bool, default: False
        [Debug option] If ``True``, show azimuth values as points.

    figsize : tuple, default: (3, 3)
        Figure size.

    azimuth_convention : .AzimuthConvention or str, optional
        Azimuth convention of plotted data (will adjust the azimuth of the plot).

    vmin, vmax : float, optional
        Range of values to plot. If unset, automatically from the data.

    theta_max : float, default: 90.0
        Maximum value of the azimuth angle on plots.

    Returns
    -------
    matplotlib.Figure
    """

    import matplotlib.pyplot as plt
    import matplotlib.tri as tri

    # Filter grazing angles
    da = (
        da.where(da.vza <= theta_max)
        .dropna("x_index", how="all")
        .dropna("y_index", how="all")
    )
    theta_max = float(da.vza.max())
    rscale = 90.0 / theta_max

    # Prepare data for plotting
    azimuth_convention = (
        AzimuthConvention.convert(azimuth_convention)
        if azimuth_convention is not None
        else settings.AZIMUTH_CONVENTION
    )
    values = da.transpose("x_index", "y_index").values.ravel()
    zeniths = da["vza"].values.ravel()  # Degree
    azimuths = np.deg2rad(da["vaa"].values).ravel()  # Radian

    # Create triangulation
    x = zeniths * np.cos(azimuths) * rscale
    y = zeniths * np.sin(azimuths) * rscale
    triangles = tri.Triangulation(x, y)

    # Make plot
    fig = plt.figure(0, figsize=figsize)
    rect = [0, 0, 1, 1]

    # -- Main plot in Cartesian coordinates
    ax_cartesian = fig.add_axes(rect, aspect="equal")
    ax_cartesian.axis("off")  # Hide axis
    ctr = ax_cartesian.tricontourf(
        triangles,
        values,
        levels=levels,
        cmap="viridis",
        vmin=vmin,
        vmax=vmax,
    )

    if show_contour:
        ax_cartesian.tricontour(
            triangles,
            values,
            levels=levels,
            linewidths=0.5,
            colors="w",
            alpha=0.5,
            vmin=vmin,
            vmax=vmax,
        )

    if show_azimuth:
        ax_cartesian.scatter(x, y, c=azimuths, cmap="plasma", s=3)

    ax_cartesian.set_xlim([-90, 90])  ## Match limits with the full zenith range
    ax_cartesian.set_ylim([-90, 90])

    ## Polar axes
    ax_polar = fig.add_axes(rect, polar=True, facecolor="none")
    ax_polar.set_rlim([0, theta_max])  # Cover the full zenith value range
    ax_polar.grid(False, axis="x")
    ax_polar.grid(axis="y", alpha=0.25)
    ax_polar.tick_params(axis="y", colors="white")
    ax_polar.set_rlabel_position(120)

    ax_polar.set_theta_offset(azimuth_convention.value[0])
    ax_polar.set_theta_direction(azimuth_convention.value[1])

    # Add the color bar (important: both axes must be adjusted)
    fig.colorbar(ctr, ax=[ax_cartesian, ax_polar], shrink=0.75, pad=0.15)

    return fig


def plot_sigma_t(
    *atmospheres: AbstractHeterogeneousAtmosphere,
    labels: list[str] | None = None,
    altitude_extent: tuple[float, float] | None = None,
    si: SpectralIndex | None = None,
    show: bool = True,
) -> tuple[plt.Figure, plt.Axes]:
    """
    Display the extinction coefficient of one or several atmosphere objects for
    a single spectral context.

    Parameters
    ----------
    *atmospheres : .AbstractHeterogeneousAtmosphere
        One or several atmosphere objects for which to plot the extinction
        coefficient.

    labels : list of strings, optional
        Labels associated with the passed atmosphere objects. If unset, no
        legend will be added to the plot.

    altitude_extent : tuple of float
        A (min, max) altitude pair (in km) to which the plot is restricted.

    si : .SpectralIndex, optional
        The spectral index at which the extinction coefficient is evaluated.
        If unset, a default spectral index is created using
        :meth:`.SpectralIndex.new`.

    show : bool, optional
        If ``True``, return ``None`` and display the plot. Otherwise, return a
        (``Figure``, ``Axes``) pair.
    """

    from matplotlib.ticker import ScalarFormatter

    from eradiate.units import to_quantity

    if si is None:  # If none is found, fall back to the default (550 nm)
        si = SpectralIndex.new()
    else:  # Otherwise, apply conversion protocol
        si = SpectralIndex.convert(si)

    if labels is None:
        label_iter = iter([None for _ in atmospheres])
    else:
        label_iter = iter(labels)

    fig, ax = plt.subplots(1, 1)

    with plt.rc_context({"lines.linestyle": ":", "lines.marker": "."}):
        for atmosphere in atmospheres:
            altitude = to_quantity(atmosphere.eval_radprops(si=si).z_layer).m_as("km")
            sigma_t = to_quantity(atmosphere.eval_radprops(si=si).sigma_t).m_as("1/m")
            ax.plot(altitude, sigma_t, label=next(label_iter))

    formatter = ScalarFormatter(useMathText=True)
    formatter.set_powerlimits((-3, 2))
    ax.yaxis.set_major_formatter(formatter)
    ax.set_ylabel("Extinction coefficient [1/m]")

    if labels is not None:
        fig.legend(
            bbox_to_anchor=(1.0, 0.5),
            loc="center left",
            borderaxespad=0.0,
        )

    if altitude_extent is not None:
        ax.set_xlim(altitude_extent)
    ax.set_xlabel("Altitude [km]")

    plt.tight_layout()
    if show:
        plt.show()
    else:
        return fig, ax


def load_ipython_extension(ipython):
    """
    IPython extension for Eradiate tutorials. Should be loaded at the top of the
    tutorial notebook.

    It sets the Matplotlib style and prints the current date and Eradiate
    version as markdown.

    See Also
    --------
    :func:`eradiate.plot.set_style`

    Notes
    -----

    This extension should be loaded using the IPython magic:

    .. code::

       %load_ext eradiate.tutorials
    """

    import datetime

    from IPython.display import Markdown, display

    from . import __version__
    from .plot import set_style

    set_style()
    xr.set_options(display_expand_data=False)
    display(
        Markdown(
            f"*Last updated: {datetime.datetime.now():%Y-%m-%d %H:%M} "
            f"(eradiate v{__version__})*"
        )
    )
