"""How a 3D surface is unrolled into the 2D plane a current map is binned on.

A `Projection` owns everything that depends on the surface's shape: which columns the
raster's axes are, how to derive them from the track positions, how the axes are
labelled, how default bins tile them, and how to tell whether the tracks actually lie
on the surface. `RSSAPlot` knows none of it.

The point of keeping this separate is that the projection is what makes a current map
*correct*: `particle_current()` divides the summed weight in each bin by `dx * dy`,
so a projection that doesn't preserve distances on the real surface (e.g. dropping a
tilted plane onto two global axes, which shrinks one of them by the cosine of the tilt)
produces plausible numbers that are wrong by a constant factor. Deriving the projection
from the header - `Projection.from_surface()` - is therefore the default; building one
by hand is the escape hatch, and it still asks for the full geometry so the area
normalization stays right.
"""

from __future__ import annotations

import math
from abc import ABC, abstractmethod
from dataclasses import dataclass

import numpy as np
import polars as pl

from .header import PERPENDICULAR_AXES, Axis, SurfaceParameters, SurfaceType

#: The direction-cosine column matching each position column.
DIRECTION_COLUMNS: dict[Axis, str] = {"x": "u", "y": "v", "z": "w"}


class Projection(ABC):
    """The 2D coordinate system a set of tracks is rastered on."""

    #: Overrides `default_description()`; set by `from_surface()` so plots derived from
    #: the header are titled with the surface they came from.
    description: str | None

    @property
    @abstractmethod
    def x_col(self) -> str:
        """Name of the column holding the raster's horizontal coordinate."""

    @property
    @abstractmethod
    def y_col(self) -> str:
        """Name of the column holding the raster's vertical coordinate."""

    def prepare(self, tracks: pl.LazyFrame) -> pl.LazyFrame:
        """Adds whatever derived columns `x_col`/`y_col` refer to. The default
        projection coordinates are plain track columns, so nothing is added."""
        return tracks

    @abstractmethod
    def labels(self) -> tuple[str, str]:
        """The x and y axis labels, units included."""

    @abstractmethod
    def default_description(self) -> str:
        """A one-line description of the projected geometry, used as the plot title."""

    def describe(self) -> str:
        return self.description or self.default_description()

    def residual(self) -> pl.Expr | None:
        """An expression giving each track's signed distance to the projected surface,
        in cm, or `None` when the geometry given isn't enough to place the surface in
        space (e.g. a plane built by hand with no `offset`). `None` means
        `RSSAPlot.verify()` has nothing to check, not that the tracks are fine."""
        return None

    def incidence_cosine(self) -> pl.Expr | None:
        """An expression giving `mu`, the cosine of the angle between each track's
        direction and the surface normal *where that track crosses it*, or `None` for a
        projection that doesn't know where its normal points.

        This is what separates the quantities that can be estimated from a set of
        crossings: weight alone gives the crossing rate, `weight / |mu|` gives the
        scalar flux, and `sign(mu)` gives the net current (see `RSSAPlot`). The
        normal is a function of position on a curved surface, so this is an expression
        over the track columns rather than a constant vector.

        The columns it reads are declared by `incidence_columns()`.
        """
        return None

    def incidence_columns(self) -> tuple[str, ...]:
        """The track columns `incidence_cosine()` reads.

        Declared rather than assumed so the aggregation can select exactly them - a
        curved surface needs position columns as well as direction cosines, and a
        plane may need only one of the three cosines, which keeps a flux computable on
        tracks read with a narrow `columns=` projection.
        """
        return ()

    def tangent_basis(self) -> tuple[pl.Expr, pl.Expr] | None:
        """Expressions giving each track's direction resolved onto the two directions
        *tangent* to the surface at its own position, or `None` for a projection that
        doesn't know how its surface sits in space.

        Together with `incidence_cosine()` these complete an orthonormal frame
        `(e1, e2, n)` in which a direction is fully described: `mu` says how head-on
        the track crosses, and these two say which way it leans while doing so.

        The invariant is alignment with the raster - `e1` points the way `x_col`
        increases, `e2` the way `y_col` does - so an angular distribution can be read
        against the map it was binned on: "up and to the right" in the angular plot
        means "up and to the right" on the current map. That is what makes the frame
        worth having on a curved surface, where it turns from bin to bin.

        Handedness is *not* the invariant, and follows from that alignment rather than
        being imposed on it: the planes come out right-handed, while `Cylinder` is
        left-handed because it deliberately rasters its surface as seen from inside.
        Anything that needs a consistent orientation should ask `incidence_cosine()`
        for the normal rather than reconstructing it from `e1 x e2`.

        The columns these read are declared by `tangent_columns()`.
        """
        return None

    def tangent_columns(self) -> tuple[str, ...]:
        """The track columns `tangent_basis()` reads, declared for the same reason as
        `incidence_columns()`."""
        return ()

    def tangent_labels(self) -> tuple[str, str]:
        """Short names for the two tangent directions, for labelling the azimuth of an
        angular plot: unambiguous enough to read, short enough to sit on a tick."""
        return ("+U", "+V")

    def scale(self) -> float:
        """A characteristic size of the surface, in cm, used to set the relative part
        of the verification tolerance."""
        return 1.0

    def bin_edges(
        self,
        bin_width: float,
        x_min: float,
        x_max: float,
        y_min: float,
        y_max: float,
    ) -> tuple[np.ndarray, np.ndarray]:
        """Bin edges of roughly `bin_width` cm covering the data's extent.

        Projections whose axes are periodic or bounded by the geometry itself (rather
        than by wherever tracks happen to land) override this.
        """
        return (
            _edges_covering(x_min, x_max, bin_width),
            _edges_covering(y_min, y_max, bin_width),
        )

    @staticmethod
    def from_surface(surface: SurfaceParameters) -> Projection:
        """The projection implied by a surface's own geometry, as recorded in the file
        header.

        - cylinders (`cx`, `cy`, `cz` and the off-axis `c/x`, `c/y`, `c/z`) give a
          `Cylinder`: the perimeter position around the cylinder's *own* axis and
          centre against the axial coordinate;
        - axis-aligned planes (`px`, `py`, `pz`) give an `AxisAlignedPlane`: the two
          coordinates spanning the plane;
        - general planes (`p`) give a `TiltedPlane`: an orthonormal basis of the plane
          itself.

        Any other surface type raises `NotImplementedError`: there is no meaningful 2D
        raster for e.g. a sphere or a torus without choosing a map projection - and
        binning one onto two global axes would not be a more general answer, just a
        wrong one, since the bin areas would no longer be areas on the surface.
        """
        description = surface.describe()
        if surface.is_cylinder:
            axis = surface.axis
            assert axis is not None  # every cylinder type has an axis
            return Cylinder(
                axis=axis,
                radius=surface.radius,
                center=surface.center,
                description=description,
            )
        if surface.surface_type in (SurfaceType.PX, SurfaceType.PY, SurfaceType.PZ):
            axis = surface.axis
            assert axis is not None
            return AxisAlignedPlane(
                axis=axis, offset=surface.offset, description=description
            )
        if surface.surface_type is SurfaceType.P:
            return TiltedPlane(
                normal=surface.normal, offset=surface.offset, description=description
            )
        raise NotImplementedError(
            f"{surface.describe()}: aci_rssa cannot derive a 2D projection for surfaces"
            f" of type {surface.mnemonic!r} yet. Pass a Projection explicitly if you"
            " know how this surface should be unrolled."
        )


@dataclass(frozen=True)
class Cylinder(Projection):
    """A cylinder unrolled around its own axis: perimeter position against height.

    `radius` is the real radius in cm (not MCNP's stored square), and `center` is where
    the axis crosses the perpendicular plane, given as the two coordinates
    perpendicular to `axis` in cyclic order - `(x0, y0)` for `axis="z"`. Both matter:
    measuring the angle and radius about the coordinate origin instead smears an
    off-axis cylinder across the plot, and the radius is what makes the bins tile the
    circumference exactly.
    """

    axis: Axis
    radius: float
    center: tuple[float, float] = (0.0, 0.0)
    description: str | None = None

    def __post_init__(self) -> None:
        if self.axis not in PERPENDICULAR_AXES:
            raise ValueError(f"axis must be one of 'x', 'y', 'z', got {self.axis!r}")
        if not self.radius > 0:
            raise ValueError(f"a cylinder needs a positive radius, got {self.radius!r}")

    @property
    def x_col(self) -> str:
        return "perimeter_pos"

    @property
    def y_col(self) -> str:
        return self.axis

    def prepare(self, tracks: pl.LazyFrame) -> pl.LazyFrame:
        """Adds the `perimeter_pos` column: `theta * r`, where `theta` is each track's
        angle about the cylinder's axis and `r` its distance to it.

        `theta` comes from `arctan2`, measured from the *negative* first perpendicular
        axis, which flips the plot so it reads as a cylinder observed from the inside.
        """
        col_a, col_b = PERPENDICULAR_AXES[self.axis]
        offset_a = pl.col(col_a) - self.center[0]
        offset_b = pl.col(col_b) - self.center[1]
        radius = (offset_a.pow(2) + offset_b.pow(2)).sqrt()
        thetas = pl.arctan2(offset_b, -1 * offset_a)
        return tracks.with_columns((thetas * radius).alias("perimeter_pos"))

    def labels(self) -> tuple[str, str]:
        return (
            f"Perimeter position around the {self.axis} axis [cm]",
            f"{self.axis.upper()} [cm]",
        )

    def default_description(self) -> str:
        about = f"cylinder about the {self.axis} axis, radius {self.radius:.2f} cm"
        if self.center != (0.0, 0.0):
            about += f", centred at {self.center}"
        return about

    def residual(self) -> pl.Expr:
        col_a, col_b = PERPENDICULAR_AXES[self.axis]
        distance_to_axis = (
            (pl.col(col_a) - self.center[0]).pow(2)
            + (pl.col(col_b) - self.center[1]).pow(2)
        ).sqrt()
        return distance_to_axis - self.radius

    def incidence_cosine(self) -> pl.Expr:
        """The normal of a cylinder is radial, so it points differently at every track:
        `mu` is the track's direction projected onto the outward radial unit vector at
        its own position.

        The radius is taken from the track's own distance to the axis rather than from
        `self.radius`, so a track that sits fractionally off the surface still gets a
        unit normal (`RSSAPlot.verify()` is what checks it is only fractionally off).
        """
        col_a, col_b = PERPENDICULAR_AXES[self.axis]
        dir_a = DIRECTION_COLUMNS[col_a]
        dir_b = DIRECTION_COLUMNS[col_b]
        offset_a = pl.col(col_a) - self.center[0]
        offset_b = pl.col(col_b) - self.center[1]
        distance_to_axis = (offset_a.pow(2) + offset_b.pow(2)).sqrt()
        return (offset_a * pl.col(dir_a) + offset_b * pl.col(dir_b)) / distance_to_axis

    def incidence_columns(self) -> tuple[str, ...]:
        col_a, col_b = PERPENDICULAR_AXES[self.axis]
        return (col_a, col_b, DIRECTION_COLUMNS[col_a], DIRECTION_COLUMNS[col_b])

    def tangent_basis(self) -> tuple[pl.Expr, pl.Expr]:
        """The tangent frame of a cylinder turns with the track's own angle: `e1` runs
        along the perimeter, `e2` along the axis.

        `e1` is the derivative of the position with respect to the angle `prepare()`
        measures, so it points the way `perimeter_pos` grows. That angle is
        `theta = pi - alpha`, which *decreases* as the standard angle `alpha` grows -
        the flip that makes the plot read as a cylinder seen from the inside - so the
        tangent is `(offset_b, -offset_a) / r` and not its negative. Getting this
        backwards would mirror every angular plot east-west, which is precisely the
        kind of error such a plot is used to diagnose.

        A consequence of that same flip: `(e1, e2, n)` is left-handed with `n` outward,
        because the raster is drawn from behind the outward normal. See
        `Projection.tangent_basis()` on why alignment with the raster, not handedness,
        is the property being preserved.
        """
        col_a, col_b = PERPENDICULAR_AXES[self.axis]
        dir_a = DIRECTION_COLUMNS[col_a]
        dir_b = DIRECTION_COLUMNS[col_b]
        offset_a = pl.col(col_a) - self.center[0]
        offset_b = pl.col(col_b) - self.center[1]
        distance_to_axis = (offset_a.pow(2) + offset_b.pow(2)).sqrt()
        along_perimeter = (
            offset_b * pl.col(dir_a) - offset_a * pl.col(dir_b)
        ) / distance_to_axis
        return along_perimeter, pl.col(DIRECTION_COLUMNS[self.axis])

    def tangent_columns(self) -> tuple[str, ...]:
        col_a, col_b = PERPENDICULAR_AXES[self.axis]
        return (
            col_a,
            col_b,
            DIRECTION_COLUMNS[col_a],
            DIRECTION_COLUMNS[col_b],
            DIRECTION_COLUMNS[self.axis],
        )

    def tangent_labels(self) -> tuple[str, str]:
        return ("+perimeter", f"+{self.axis.upper()}")

    def scale(self) -> float:
        return max(self.radius, 1.0)

    def bin_edges(
        self,
        bin_width: float,
        x_min: float,
        x_max: float,
        y_min: float,
        y_max: float,
    ) -> tuple[np.ndarray, np.ndarray]:
        """The perimeter axis is binned over the cylinder's full circumference rather
        than over the extent of the data, so bins tile it exactly: the width is nudged
        to the nearest whole number of bins that fits `2*pi*r`, which avoids a ragged
        partial bin at the seam."""
        half_perimeter = math.pi * self.radius
        number_of_bins = max(1, round(2 * half_perimeter / bin_width))
        x_edges = np.linspace(-half_perimeter, half_perimeter, number_of_bins + 1)
        return x_edges, _edges_covering(y_min, y_max, bin_width)


@dataclass(frozen=True)
class AxisAlignedPlane(Projection):
    """A `px`/`py`/`pz` plane, plotted on the two coordinate axes that span it.

    `offset` is where the plane sits along `axis`. It doesn't affect the projection -
    the raster is the same wherever the plane is - and is only used to check that the
    tracks really lie on it, so it may be left out; verification is then skipped.
    """

    axis: Axis
    offset: float | None = None
    description: str | None = None

    def __post_init__(self) -> None:
        if self.axis not in PERPENDICULAR_AXES:
            raise ValueError(f"axis must be one of 'x', 'y', 'z', got {self.axis!r}")

    @property
    def x_col(self) -> str:
        return PERPENDICULAR_AXES[self.axis][0]

    @property
    def y_col(self) -> str:
        return PERPENDICULAR_AXES[self.axis][1]

    def labels(self) -> tuple[str, str]:
        return f"{self.x_col.upper()} [cm]", f"{self.y_col.upper()} [cm]"

    def default_description(self) -> str:
        if self.offset is None:
            return f"plane perpendicular to the {self.axis} axis"
        return f"plane {self.axis} = {self.offset:.4g}"

    def residual(self) -> pl.Expr | None:
        if self.offset is None:
            return None
        return pl.col(self.axis) - self.offset

    def incidence_cosine(self) -> pl.Expr:
        """The normal is the axis the plane is perpendicular to, so `mu` is simply that
        axis' direction cosine. Unlike `residual()` this needs no `offset`: where the
        plane sits doesn't change which way it faces."""
        return pl.col(DIRECTION_COLUMNS[self.axis])

    def incidence_columns(self) -> tuple[str, ...]:
        return (DIRECTION_COLUMNS[self.axis],)

    def tangent_basis(self) -> tuple[pl.Expr, pl.Expr]:
        """The direction cosines of the two axes the plane is plotted on. They come
        from `PERPENDICULAR_AXES` in cyclic order, so `(e1, e2, n)` is right-handed."""
        col_a, col_b = PERPENDICULAR_AXES[self.axis]
        return pl.col(DIRECTION_COLUMNS[col_a]), pl.col(DIRECTION_COLUMNS[col_b])

    def tangent_columns(self) -> tuple[str, ...]:
        col_a, col_b = PERPENDICULAR_AXES[self.axis]
        return (DIRECTION_COLUMNS[col_a], DIRECTION_COLUMNS[col_b])

    def tangent_labels(self) -> tuple[str, str]:
        return (f"+{self.x_col.upper()}", f"+{self.y_col.upper()}")

    def scale(self) -> float:
        return max(abs(self.offset), 1.0) if self.offset is not None else 1.0


@dataclass(frozen=True)
class TiltedPlane(Projection):
    """A general `p` plane, plotted on an orthonormal basis of the plane itself.

    Projecting a tilted plane onto two global axes instead would shrink one of them by
    the cosine of the tilt, and the current is per unit area - so the bin areas, and
    every number derived from them, would be wrong by that factor.

    The basis is added as the `plane_u`/`plane_v` columns and centred on the projection
    of the origin onto the plane. As for `AxisAlignedPlane`, `offset` is only needed to
    verify the tracks lie on the plane.
    """

    normal: tuple[float, float, float]
    offset: float | None = None
    description: str | None = None

    @property
    def x_col(self) -> str:
        return "plane_u"

    @property
    def y_col(self) -> str:
        return "plane_v"

    @property
    def basis(self) -> tuple[np.ndarray, np.ndarray]:
        return plane_basis(self.normal)

    def prepare(self, tracks: pl.LazyFrame) -> pl.LazyFrame:
        first, second = self.basis
        return tracks.with_columns(
            project_onto(first).alias("plane_u"),
            project_onto(second).alias("plane_v"),
        )

    def labels(self) -> tuple[str, str]:
        first, second = self.basis
        return (
            f"In-plane U {_format_direction(first)} [cm]",
            f"In-plane V {_format_direction(second)} [cm]",
        )

    def default_description(self) -> str:
        a, b, c = _unit(self.normal)
        equation = f"{a:.4g}x + {b:.4g}y + {c:.4g}z"
        if self.offset is None:
            return f"plane with normal {equation}"
        return f"plane {equation} = {self.offset:.4g}"

    def residual(self) -> pl.Expr | None:
        if self.offset is None:
            return None
        return project_onto(np.asarray(_unit(self.normal))) - self.offset

    def incidence_cosine(self) -> pl.Expr:
        """`mu` is the track's direction dotted with the plane's (unit) normal."""
        a, b, c = _unit(self.normal)
        return pl.col("u") * a + pl.col("v") * b + pl.col("w") * c

    def incidence_columns(self) -> tuple[str, ...]:
        return ("u", "v", "w")

    def tangent_basis(self) -> tuple[pl.Expr, pl.Expr]:
        """The plane's own basis, the same one `prepare()` rasters positions on.
        `plane_basis()` builds it as `first = reference x n`, `second = n x first`, so
        `first x second = n` and the frame is right-handed."""
        first, second = self.basis
        return direction_along(first), direction_along(second)

    def tangent_columns(self) -> tuple[str, ...]:
        return ("u", "v", "w")

    def scale(self) -> float:
        return max(abs(self.offset), 1.0) if self.offset is not None else 1.0


def plane_basis(
    normal: tuple[float, float, float],
) -> tuple[np.ndarray, np.ndarray]:
    """Returns two orthonormal vectors spanning the plane with the given `normal`.

    The second one points as close to +z as the plane allows, so a wall-like surface
    is plotted upright; for a near-horizontal plane (where "up" is meaningless in
    plane) the basis falls back to being aligned with the y axis instead.
    """
    n = np.asarray(normal, dtype=np.float64)
    n = n / np.linalg.norm(n)
    reference = np.array([0.0, 0.0, 1.0])
    if abs(float(n @ reference)) > 0.9:
        reference = np.array([0.0, 1.0, 0.0])
    first = np.cross(reference, n)
    first /= np.linalg.norm(first)
    second = np.cross(n, first)
    return first, second


def project_onto(direction: np.ndarray) -> pl.Expr:
    """The (signed) length of each track's position vector along `direction`."""
    return (
        pl.col("x") * float(direction[0])
        + pl.col("y") * float(direction[1])
        + pl.col("z") * float(direction[2])
    )


def direction_along(direction: np.ndarray) -> pl.Expr:
    """The component of each track's direction cosine along `direction`.

    The counterpart of `project_onto()`, which does the same for positions: a track
    carries both, in the same frame but in different columns.
    """
    return (
        pl.col("u") * float(direction[0])
        + pl.col("v") * float(direction[1])
        + pl.col("w") * float(direction[2])
    )


def _unit(vector: tuple[float, float, float]) -> tuple[float, float, float]:
    length = math.sqrt(sum(component * component for component in vector))
    if length == 0.0:
        raise ValueError("a plane cannot have a degenerate (zero) normal")
    a, b, c = vector
    return a / length, b / length, c / length


def _format_direction(direction: np.ndarray) -> str:
    # Rounded before formatting so an axis that is only numerically-not-quite-zero
    # (MCNP writes e.g. -3.3e-11 for the z coefficient of a vertical plane) reads as 0
    # in the axis label rather than as noise. `+ 0.0` normalizes -0.0 away.
    rounded = np.round(direction, 3) + 0.0
    return "(" + ", ".join(f"{component:g}" for component in rounded) + ")"


def _edges_covering(minimum: float, maximum: float, bin_width: float) -> np.ndarray:
    """Bin edges of exactly `bin_width` starting at `minimum` and covering `maximum`.

    Built from a bin count rather than `np.arange(minimum, maximum + bin_width, ...)`,
    whose length depends on floating-point round-off at the last step and so could
    silently gain or lose a bin.
    """
    if bin_width <= 0:
        raise ValueError(f"bin_width must be positive, got {bin_width!r}")
    span = maximum - minimum
    number_of_bins = max(1, math.ceil(span / bin_width))
    return minimum + np.arange(number_of_bins + 1) * bin_width
