"""The values the plotting classes compute, and how they render themselves.

Each of these is a finished result: holding one means the computation already ran, so
there is no "you forgot to call X first" state to guard against. They are immutable -
`with_parameters()` and `ratio_to()` return new objects - which is what makes it safe
to derive several renderings, or a ratio, from the same computed map.
"""

from __future__ import annotations

import dataclasses
from collections.abc import Sequence
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Literal, cast

import numpy as np
import polars as pl
from matplotlib import colors
from matplotlib import pyplot as plt
from matplotlib.axes import Axes
from matplotlib.figure import Figure
from matplotlib.projections.polar import PolarAxes
from typing_extensions import Self


@dataclass(frozen=True)
class PlotParameters:
    """How a `CurrentMap` is rendered."""

    title: str = ""
    xlabel: str = "X [cm]"
    ylabel: str = "Y [cm]"
    legend_label: str = ""
    legend_orientation: Literal["vertical", "horizontal"] = "horizontal"
    cmap: str = "plasma"
    number_of_colors: int = 10
    norm: Literal["linear", "log"] = "log"
    vmin: float | None = None
    vmax: float | None = None


@dataclass(frozen=True)
class VectorCurrentPlotParameters:
    """How a `VectorCurrent` is rendered."""

    title: str = "Dominant particle directions"
    #: Any Matplotlib colormap name.
    cmap: str = "plasma"
    colorbar_label: str = "Mean energy [MeV]"


@dataclass(frozen=True)
class AngularPlotParameters:
    """How an `AngularDistribution` is rendered."""

    title: str = "Angular distribution"
    #: Names of the two tangent directions, placed at 0 and 90 degrees of azimuth; the
    #: opposite two ticks are labelled by flipping the leading sign. Set from the
    #: projection, so the azimuth reads against the map the tracks were binned on.
    tangent_labels: tuple[str, str] = ("+U", "+V")
    cmap: str = "plasma"
    number_of_colors: int = 10
    norm: Literal["linear", "log"] = "log"
    vmin: float | None = None
    vmax: float | None = None
    legend_label: str = "Angular flux [/cm2/sr per unit source]"


@dataclass(frozen=True)
class SpectraPlotParameters:
    """How a `Spectrum` is rendered."""

    title: str = "Particle energy spectra"
    xlabel: str = "Energy [eV]"
    ylabel: str = "Normalized counts per unit of lethargy"
    label: str = ""
    min_energy: float | None = None
    max_energy: float | None = None


class _Renderable:
    """Shared `with_parameters()`/`show()`/`save()` behaviour of the result types.

    Subclasses are frozen dataclasses that hold their rendering options in a
    `parameters` field and implement `figure()`.
    """

    parameters: Any

    def figure(self) -> tuple[Figure, Axes]:  # pragma: no cover - overridden
        raise NotImplementedError

    def with_parameters(self, parameters: Any = None, /, **overrides: Any) -> Self:
        """Returns a copy rendered with different options.

        Pass a full parameters object to replace them wholesale, or keyword arguments
        to change individual ones - the latter keeps the labels and title derived from
        the surface's own geometry, which replacing them wholesale would discard:

        ```python
        current.with_parameters(vmin=1e6, vmax=1e12)
        ```
        """
        if parameters is not None and overrides:
            raise TypeError(
                "pass either a parameters object or keyword overrides, not both"
            )
        if parameters is None:
            parameters = dataclasses.replace(self.parameters, **overrides)
        return dataclasses.replace(cast("Any", self), parameters=parameters)

    def show(self) -> Self:
        """Renders the result and shows it in a window."""
        fig, _ax = self.figure()
        fig.show()
        return self

    def save(self, out_path: Path | str) -> Self:
        """Renders the result and writes it to `out_path` as a PNG."""
        fig, _ax = self.figure()
        out_path = Path(out_path)
        if out_path.suffix != ".png":
            out_path = out_path.with_suffix(".png")
        fig.savefig(out_path, dpi=300, bbox_inches="tight")
        return self


@dataclass(frozen=True)
class CurrentMap(_Renderable):
    """A 2D map of a per-bin quantity - a flux, a current, a relative error, or a ratio
    between two of them - together with the bins it was computed on.

    `values` is indexed `[y, x]`, matching the raster Matplotlib expects.

    `relative_errors` holds the statistical error of `values`, when it was asked for
    (`errors=True` on the method that produced this map). It is carried alongside the
    values rather than computed separately so it can only ever belong to the quantity
    it was estimated for - the error of a flux is not the error of a current, and the
    two are easy to mix up when they are computed by different calls.
    """

    values: np.ndarray
    x_bins: np.ndarray
    y_bins: np.ndarray
    parameters: PlotParameters = field(default_factory=PlotParameters)
    relative_errors: np.ndarray | None = None

    def error_map(self) -> CurrentMap:
        """The relative statistical error of this map, as a map in its own right.

        Rendered on a linear scale from 0: a relative error lives in a narrow range
        around 0.1 (MCNP's own convention treats values below that as reliable), not
        across decades. Bins no history reached are `NaN` - their error is undefined,
        not zero.
        """
        if self.relative_errors is None:
            raise ValueError(
                "This map carries no statistical errors. Pass errors=True to the"
                " method that produced it (e.g. plot.particle_flux(1e17,"
                " errors=True)) - the error has to be estimated from the tracks, at"
                " the same time and for the same quantity as the values."
            )
        return CurrentMap(
            values=self.relative_errors,
            x_bins=self.x_bins,
            y_bins=self.y_bins,
            parameters=dataclasses.replace(
                self.parameters,
                norm="linear",
                vmin=0.0,
                vmax=None,
                legend_label="Relative error",
                title=f"{self.parameters.title} (relative error)".strip(),
            ),
        )

    def ratio_to(self, other: CurrentMap) -> CurrentMap:
        """The relative difference to `other`, in %, computed as
        `(other - self) / self * 100`.

        Both maps must have been computed on the same bins, and should have undergone
        the same processing otherwise. The result can be negative, which a log colour
        scale can't render, so it defaults to a linear one.
        """
        if not (
            np.array_equal(self.x_bins, other.x_bins)
            and np.array_equal(self.y_bins, other.y_bins)
        ):
            raise ValueError(
                "the two maps were computed on different bins, so they cannot be"
                " compared bin by bin. Build both plots with the same with_bins() /"
                " with_bin_width() call."
            )
        with np.errstate(divide="ignore", invalid="ignore"):
            # Bins neither map reached divide 0 by 0, which is genuinely undefined:
            # NaN is the answer we want there, so the warning is noise.
            values = (other.values - self.values) / self.values * 100
        return CurrentMap(
            values=values,
            x_bins=self.x_bins,
            y_bins=self.y_bins,
            parameters=dataclasses.replace(
                self.parameters,
                norm="linear",
                legend_label="Relative difference [%]",
            ),
        )

    def figure(self) -> tuple[Figure, Axes]:
        """Returns the Matplotlib figure and axes for further manual customization."""
        if self.parameters.norm == "log":
            # Zero is deliberately not flagged here: it's the normal value for a
            # spatial bin with no crossings (e.g. outside a beam's footprint), and
            # LogNorm masking it as blank is the intended rendering for that case,
            # not a bug. A *negative* value is different - it can only come from
            # ratio_to(), and rendering it as the same blank as "no data" would hide
            # a meaningful decrease, so that's worth catching loudly instead.
            finite_values = self.values[np.isfinite(self.values)]
            if finite_values.size and finite_values.min() < 0:
                raise ValueError(
                    "plot parameters ask for a 'log' norm, but the data contains "
                    "negative values. A log-scale colormap can't represent those - it "
                    "would silently render them as blank/masked cells, "
                    "indistinguishable from bins with no data, instead of raising "
                    "here. Use with_parameters(norm='linear')."
                )
        norm = colors.LogNorm if self.parameters.norm == "log" else colors.Normalize

        fig, ax = plt.subplots()
        im = ax.pcolormesh(
            self.x_bins,
            self.y_bins,
            self.values,
            cmap=plt.get_cmap(self.parameters.cmap, self.parameters.number_of_colors),
            norm=norm(self.parameters.vmin, self.parameters.vmax),
        )
        fig.colorbar(
            im,
            ax=ax,
            label=self.parameters.legend_label,
            orientation=self.parameters.legend_orientation,
        )
        ax.set_aspect("equal")
        ax.set_title(self.parameters.title)
        ax.set_xlabel(self.parameters.xlabel)
        ax.set_ylabel(self.parameters.ylabel)
        return fig, ax


#: Columns of `VectorCurrent.table`, in order.
VECTOR_CURRENT_COLUMNS = (
    "pos_x",
    "pos_y",
    "pos_z",
    "dir_u",
    "dir_v",
    "dir_w",
    "magnitude",
    "net_current",
    "flux",
    "crossing_rate",
    "anisotropy",
    "mean_energy",
)


@dataclass(frozen=True)
class VectorCurrent(_Renderable):
    """The net vector current per spatial bin: see `RSSAPlot.vector_current()` for what
    each column of `table` means."""

    table: pl.DataFrame
    parameters: VectorCurrentPlotParameters = field(
        default_factory=VectorCurrentPlotParameters
    )

    def figure(self) -> tuple[Figure, Axes]:
        """Returns a 3D Matplotlib figure and axes: one arrow per spatial bin,
        positioned at the bin's mean location, oriented along its net direction, sized
        by current magnitude, and colored by mean energy."""
        fig = plt.figure()
        ax = fig.add_subplot(projection="3d")
        ax.set_xlabel("X [cm]")
        ax.set_ylabel("Y [cm]")
        ax.set_zlabel("Z [cm]")

        if self.table.is_empty():
            ax.set_title(f"{self.parameters.title} (no data)")
            return fig, ax

        positions = self.table.select(["pos_x", "pos_y", "pos_z"]).to_numpy()
        directions = self.table.select(["dir_u", "dir_v", "dir_w"]).to_numpy()
        energies = self.table["mean_energy"].to_numpy()
        magnitudes = self.table["magnitude"].to_numpy()

        e_min, e_max = float(energies.min()), float(energies.max())
        norm = colors.Normalize(
            vmin=e_min, vmax=e_max if e_max > e_min else e_min + 1.0
        )
        cmap = plt.get_cmap(self.parameters.cmap)
        arrow_colors = cmap(norm(energies))

        # Arrows are drawn in physical current units, which are usually tiny compared
        # to the geometry's spatial scale; rescale the longest arrow to a fixed
        # fraction of that scale so the plot is legible, rather than plotting raw
        # magnitudes.
        max_magnitude = float(magnitudes.max())
        spatial_extent = float(np.ptp(positions, axis=0).max())
        length_scale = (
            (0.2 * spatial_extent / max_magnitude) if max_magnitude > 0 else 1.0
        )

        ax.quiver(
            positions[:, 0],
            positions[:, 1],
            positions[:, 2],
            directions[:, 0],
            directions[:, 1],
            directions[:, 2],
            colors=arrow_colors,
            length=length_scale,
            normalize=False,
            arrow_length_ratio=0.3,
        )

        scalar_mappable = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
        scalar_mappable.set_array([])
        fig.colorbar(
            scalar_mappable,
            ax=ax,
            label=self.parameters.colorbar_label,
            shrink=0.6,
        )
        ax.set_title(self.parameters.title)
        return fig, ax


@dataclass(frozen=True)
class AngularDistribution(_Renderable):
    """The angular flux over the whole sphere of directions, in the surface's own
    frame: see `RSSAPlot.angular_distribution()` for how it is estimated.

    `values` is indexed `[mu, phi]` and holds the angular flux `psi(Omega)` in each
    cell, in units of `source_intensity` per cm2 per steradian. `mu_bins` runs from -1
    to 1 with an even number of bins (so `mu = 0`, the grazing direction, is always an
    edge and the two hemispheres are separable), and `phi_bins` from `-pi` to `pi`,
    measured from the first tangent direction towards the second.

    Both grids are uniform, which is deliberate: uniform in `mu` (not in the polar
    angle) means every cell subtends the same solid angle, so the values are directly
    comparable cell to cell with no per-cell correction, and the rendering can be
    equal-area too.
    """

    values: np.ndarray
    mu_bins: np.ndarray
    phi_bins: np.ndarray
    parameters: AngularPlotParameters = field(default_factory=AngularPlotParameters)

    @property
    def solid_angles(self) -> np.ndarray:
        """The solid angle of each cell, in steradians, shaped like `values`. Constant
        by construction; kept as an array so integrating over the sphere reads as the
        sum it is."""
        return np.outer(np.diff(self.mu_bins), np.diff(self.phi_bins))

    @property
    def scalar_flux(self) -> float:
        """`integral psi dOmega` over the sphere: the scalar flux, averaged over the
        area this distribution was spread across.

        Exact rather than approximate - the per-cell solid angle divided out to make
        `values` an angular flux is multiplied straight back in - so this equals the
        area-weighted mean of `RSSAPlot.particle_flux()` over the same tracks whenever
        the two are normalized by the same area. They are not, by default: a flux map
        divides each bin by its own area, padding included, while `angular_distribution`
        divides by the extent the tracks actually span (see it for why). Pass its
        `area` argument to line the two up exactly.
        """
        return float(np.nansum(self.values * self.solid_angles))

    def hemispheres(self) -> tuple[tuple[str, np.ndarray, np.ndarray], ...]:
        """The two halves of the sphere as `(label, radial edges, values)`, ready to
        `pcolormesh` on polar axes.

        The radius is the Lambert azimuthal equal-area one, `r = sqrt(2(1 - |mu|))`:
        head-on at the centre, grazing at the rim, and `r dr = d|mu|`, so a cell's area
        on the page is proportional to the solid angle it stands for. Plotting the
        polar angle itself as the radius instead would inflate the near-normal
        directions into a wide bullseye and squeeze the grazing ones against the rim,
        which is the same area-encoding bias that makes a bar rose overstate its
        longest petal.
        """
        half = (len(self.mu_bins) - 1) // 2
        # Outgoing: mu runs 0 -> 1 outwards in `values`, but the pole (mu = 1) is at
        # the *centre* of the disk, so both edges and rows are reversed.
        outgoing_radii = _equal_area_radii(self.mu_bins[half:])[::-1]
        incoming_radii = _equal_area_radii(self.mu_bins[: half + 1])
        return (
            ("Outgoing (mu > 0)", outgoing_radii, self.values[half:][::-1]),
            ("Incoming (mu < 0)", incoming_radii, self.values[:half]),
        )

    def figure(self) -> tuple[Figure, Axes]:
        """Returns a Matplotlib figure with one equal-area polar disk per hemisphere,
        and the outgoing axes. Both disks share one colour scale, so they can be
        compared to each other rather than only read individually; the other axes are
        reachable through `fig.axes`.
        """
        vmin, vmax = self._colour_limits()
        norm = colors.LogNorm if self.parameters.norm == "log" else colors.Normalize
        cmap = plt.get_cmap(self.parameters.cmap, self.parameters.number_of_colors)

        fig, axes = plt.subplots(
            1,
            2,
            figsize=(11, 5.0),
            subplot_kw={"projection": "polar"},
            # Room between the disks for the azimuth labels, which are words rather
            # than compass points and would otherwise run into each other.
            gridspec_kw={"wspace": 0.45},
        )
        mesh = None
        for ax, (label, radii, values) in zip(axes, self.hemispheres(), strict=True):
            mesh = ax.pcolormesh(
                self.phi_bins, radii, values, cmap=cmap, norm=norm(vmin, vmax)
            )
            # An empty hemisphere is a result, not a gap: on an `SSW` boundary every
            # track crosses the same way, and a blank disk that says so is how you see
            # at a glance that nothing is coming back.
            self._decorate(
                cast(PolarAxes, ax), label if values.any() else f"{label} - no tracks"
            )

        assert mesh is not None
        fig.colorbar(
            mesh,
            ax=axes,
            label=self.parameters.legend_label,
            fraction=0.045,
            pad=0.08,
        )
        fig.suptitle(self.parameters.title, y=0.99)
        fig.text(
            0.45,
            0.04,
            "rings: |mu| - centre head-on, rim grazing (equal-area)",
            ha="center",
            fontsize=8,
            color="0.35",
        )
        return fig, axes[0]

    def _colour_limits(self) -> tuple[float, float]:
        """One scale for both hemispheres. Resolved here rather than left to each
        `pcolormesh` to autoscale: two disks on independent scales look comparable and
        are not, which is the whole failure mode this plot exists to avoid."""
        logarithmic = self.parameters.norm == "log"
        vmin, vmax = self.parameters.vmin, self.parameters.vmax
        finite = self.values[np.isfinite(self.values)]
        if logarithmic:
            # Zero is the normal value for a direction nothing travels in, and
            # LogNorm renders it as blank, which is the right picture; it just can't
            # be an end of the scale.
            finite = finite[finite > 0]
        if vmax is None:
            vmax = float(finite.max()) if finite.size else 1.0
        if vmin is None:
            vmin = float(finite.min()) if finite.size else _below(vmax, logarithmic)
        if vmin >= vmax:
            # A distribution that is empty, or all one value; both norms reject a
            # degenerate range, so open it out rather than raise on a dull plot.
            vmin = _below(vmax, logarithmic)
        return vmin, vmax

    def _decorate(self, ax: PolarAxes, label: str) -> None:
        first, second = self.parameters.tangent_labels
        ax.set_xticks(np.deg2rad([0.0, 90.0, 180.0, 270.0]))
        ax.set_xticklabels(
            [first, second, _opposite(first), _opposite(second)], fontsize=9
        )
        # The rings are labelled with the |mu| they stand for rather than with the
        # radius, which is a projected quantity nobody reads directly.
        ring_cosines = np.array([0.75, 0.5, 0.25])
        ax.set_rgrids(
            _equal_area_radii(ring_cosines),
            labels=[f"{cosine:g}" for cosine in ring_cosines],
            angle=22.5,
            fontsize=7,
        )
        ax.set_rmax(float(np.sqrt(2.0)))
        ax.set_title(label, fontsize=10, pad=14)


def _equal_area_radii(cosines: np.ndarray) -> np.ndarray:
    """`sqrt(2(1 - |mu|))`, the Lambert equal-area radius of each cosine. Clipped
    because a track crossing exactly normally can come back from the arithmetic as
    `|mu|` a few ulps above 1, and a negative radicand there would be a `nan` ring."""
    return np.sqrt(np.clip(2.0 * (1.0 - np.abs(cosines)), 0.0, None))


def _below(vmax: float, logarithmic: bool) -> float:
    """A lower end for a colour scale whose data gave none: a decade under `vmax` on a
    log scale, where zero is not representable, and one unit under it on a linear one.
    """
    return vmax / 10.0 if logarithmic else vmax - 1.0


def _opposite(label: str) -> str:
    """The tick label for the direction opposite `label`."""
    return "-" + label[1:] if label.startswith("+") else f"-({label})"


@dataclass(frozen=True)
class Spectrum(_Renderable):
    """An energy spectrum: counts per unit lethargy, normalized so they sum to 1."""

    normalized_counts: np.ndarray
    energy_bins: np.ndarray
    parameters: SpectraPlotParameters = field(default_factory=SpectraPlotParameters)

    @property
    def label(self) -> str:
        """The legend label, from the plot parameters."""
        return self.parameters.label

    def scaled_by(self, factor: float) -> Spectrum:
        """The same spectrum with its counts multiplied by `factor`, e.g. to offset
        two otherwise identical curves in a comparison plot."""
        return dataclasses.replace(
            self, normalized_counts=self.normalized_counts * factor
        )

    def figure(self) -> tuple[Figure, Axes]:
        """Returns the Matplotlib figure and axes for further manual customization."""
        return self.overlay()

    def overlay(self, *others: Spectrum) -> tuple[Figure, Axes]:
        """Renders this spectrum together with `others` on shared axes, e.g. to
        compare filtering strategies or two different files."""
        fig, ax = plt.subplots()
        for spectrum in (self, *others):
            ax.step(
                spectrum.energy_bins[:-1],
                spectrum.normalized_counts,
                label=spectrum.label,
            )

        ax.set_xscale("log")
        ax.set_yscale("log")
        if any(spectrum.label for spectrum in (self, *others)):
            ax.legend()
        ax.set_xlabel(self.parameters.xlabel)
        ax.set_ylabel(self.parameters.ylabel)
        ax.set_title(self.parameters.title)
        ax.set_xlim(self.parameters.min_energy, self.parameters.max_energy)
        ax.grid(True)
        return fig, ax


def calculate_areas(
    x_bins: Sequence[float] | np.ndarray | pl.Series,
    y_bins: Sequence[float] | np.ndarray | pl.Series,
) -> np.ndarray:
    """Calculate the areas of the bins in a 2D histogram.

    Parameters
    ----------
    x_bins: Sequence[float] | np.ndarray | pl.Series
        The x-axis bin edges.
    y_bins: Sequence[float] | np.ndarray | pl.Series
        The y-axis bin edges.

    Returns
    -------
    np.ndarray
        A 2D array with the areas of each bin, indexed `[y, x]`.
    """
    dx = np.diff(np.asarray(x_bins))
    dy = np.diff(np.asarray(y_bins))
    return np.outer(dy, dx)
