"""Binning tracks into the quantities you plot.

The two classes here are *builders*: they hold a lazy set of tracks plus the bins to
aggregate it on, and every method returns a new object, so a plot can be reused as the
starting point for several derived ones. Computing something hands back a finished
result from `aci_rssa.results` - `CurrentMap`, `VectorCurrent`, `Spectrum` - which is
what knows how to render itself.

Row selection is *not* part of these builders: the tracks are a Polars `LazyFrame`, so
`filter()` takes ordinary Polars expressions, with `aci_rssa.expressions` covering the
two selections that need file-specific knowledge.
"""

from __future__ import annotations

import copy
import dataclasses
import warnings
from collections.abc import Sequence
from typing import Literal

import numpy as np
import polars as pl
from typing_extensions import Self

from .expressions import neutrons, on_surface, photons
from .header import FileParameters
from .projection import (
    AxisAlignedPlane,
    Cylinder,
    Projection,
    TiltedPlane,
    plane_basis,
    project_onto,
)
from .results import (
    VECTOR_CURRENT_COLUMNS,
    AngularDistribution,
    AngularPlotParameters,
    CurrentMap,
    PlotParameters,
    Spectrum,
    VectorCurrent,
    calculate_areas,
)

__all__ = [
    "MCNP_GRAZING_COSINE",
    "AxisAlignedPlane",
    "Cylinder",
    "Projection",
    "RSSAPlot",
    "RSSASpectraPlot",
    "TiltedPlane",
    "calculate_areas",
    "plane_basis",
    "project_onto",
]

#: How far a track may sit from the surface it was recorded on before
#: `RSSAPlot.verify()` considers the header's geometry to disagree with the data, as an
#: absolute distance in cm plus a term relative to the surface's own size (the
#: coordinates involved reach thousands of cm, so a purely absolute tolerance would be
#: either too tight or meaningless).
_SURFACE_ABSOLUTE_TOLERANCE = 1e-4
_SURFACE_RELATIVE_TOLERANCE = 1e-9

#: Number of tracks checked against the surface geometry. A mismatch is systematic
#: (wrong axis, wrong centre, radius vs squared radius), so it shows up immediately;
#: there is no point scanning a 100 GB file for it.
_SURFACE_VERIFICATION_SAMPLE = 10_000

#: Below this |mu|, a track counts as grazing and its flux contribution is capped -
#: `1/|mu|` would otherwise run to infinity as a track approaches tangency. This is
#: MCNP's own F2 threshold, and `particle_flux()` follows MCNP in scoring a capped
#: track at half of it (see `RSSAPlot._grazing_limited_cosine`).
MCNP_GRAZING_COSINE = 1e-3



class TrackPlot:
    """Common behaviour of the plot builders: a lazy set of tracks, the file's header,
    and the row selection applied to them."""

    tracks: pl.LazyFrame
    rssa_parameters: FileParameters

    def _replace(self, **changes: object) -> Self:
        """A copy of this builder with some attributes changed. Every public builder
        method goes through here: nothing mutates in place, so the same object can be
        reused as the base of several plots."""
        clone = copy.copy(self)
        for name, value in changes.items():
            setattr(clone, name, value)
        return clone

    def filter(self, *predicates: pl.Expr) -> Self:
        """Keeps only the tracks matching every predicate, which are ordinary Polars
        expressions over the track columns.

        `aci_rssa.expressions` covers the two selections that need knowledge of the
        file rather than of Polars - the packed neutron code, and surface ids:

        ```python
        from aci_rssa import neutrons, on_surface

        plot.filter(neutrons(), pl.col("energy") > 1.0)
        plot.filter(on_surface(rssa.parameters, 63))
        ```
        """
        return self._replace(tracks=self.tracks.filter(*predicates))

    # ------------------------------------------------------------------
    # Deprecated filter helpers, kept one release for the chained style.
    # ------------------------------------------------------------------

    def set_particle(self, particle_type: Literal["n", "p"]) -> Self:
        """Deprecated: use `filter(neutrons())` or `filter(photons())`."""
        _deprecated("set_particle()", "filter(neutrons()) / filter(photons())")
        if particle_type not in ("n", "p"):
            raise ValueError(f"particle_type must be 'n' or 'p', got {particle_type!r}")
        return self.filter(neutrons() if particle_type == "n" else photons())

    def set_surface_ids(self, surface_ids: list[int]) -> Self:
        """Deprecated: use `filter(on_surface(parameters, *ids))`."""
        _deprecated("set_surface_ids()", "filter(on_surface(parameters, *ids))")
        return self.filter(on_surface(self.rssa_parameters, *surface_ids))

    def set_axis_limits(
        self, column: Literal["x", "y", "z"], vmin: float, vmax: float
    ) -> Self:
        """Deprecated: use `filter(pl.col(column).is_between(vmin, vmax))`."""
        _deprecated("set_axis_limits()", "filter(pl.col(column).is_between(...))")
        return self.filter(pl.col(column).is_between(vmin, vmax, closed="both"))

    def set_z_limits(self, vmin: float, vmax: float) -> Self:
        """Deprecated: use `filter(pl.col("z").is_between(vmin, vmax))`."""
        _deprecated("set_z_limits()", 'filter(pl.col("z").is_between(...))')
        return self.filter(pl.col("z").is_between(vmin, vmax, closed="both"))

    def set_perimeter_limits(self, vmin: float, vmax: float) -> Self:
        """Deprecated: use `filter(pl.col("perimeter_pos").is_between(vmin, vmax))`."""
        _deprecated(
            "set_perimeter_limits()", 'filter(pl.col("perimeter_pos").is_between(...))'
        )
        if "perimeter_pos" not in self.tracks.collect_schema().names():
            raise ValueError(
                "there is no 'perimeter_pos' column to bound: it is added by a"
                " Cylinder projection, so this plot has to be built with one (e.g."
                " with RSSA.plot_surface() on a cylindrical surface)."
            )
        return self.filter(
            pl.col("perimeter_pos").is_between(vmin, vmax, closed="both")
        )


class RSSAPlot(TrackPlot):
    """A 2D map of the tracks crossing one surface.

    The `projection` decides which two coordinates the map is binned on and how they
    are derived from the track positions; it is normally taken from the file header by
    `for_surface()` (or `RSSA.plot_surface()`), which is also what makes the per-area
    normalization of the computed currents correct.
    """

    def __init__(
        self,
        tracks: pl.DataFrame | pl.LazyFrame,
        rssa_parameters: FileParameters,
        projection: Projection,
    ) -> None:
        self.rssa_parameters = rssa_parameters
        self.projection = projection
        self.tracks = projection.prepare(tracks.lazy())
        self._x_bins: np.ndarray | None = None
        self._y_bins: np.ndarray | None = None

    @classmethod
    def for_surface(
        cls,
        tracks: pl.DataFrame | pl.LazyFrame,
        rssa_parameters: FileParameters,
        surface_id: int | None = None,
        *,
        projection: Projection | None = None,
        verify: bool = True,
    ) -> RSSAPlot:
        """A plot of one recorded surface, projected according to its own geometry.

        This is what `RSSA.plot_surface()` calls; use it directly when plotting a
        `scan_tracks()` LazyFrame, which never goes through an `RSSA` object.

        Tracks are only filtered when the file recorded several surfaces - with a
        single surface there is nothing to separate, which keeps files that don't
        record a per-track surface id (see `FileParameters.records_surface_id`) usable.

        Parameters
        ----------
        surface_id: int | None
            Which recorded surface to plot, matched irrespective of sign. May be
            omitted only when the file recorded a single surface.
        projection: Projection | None
            Overrides the projection derived from the header, for a surface type
            `Projection.from_surface()` has no geometry for, or when the header's
            coefficients need reinterpreting. The tracks are still restricted to
            `surface_id`.
        verify: bool
            Check a sample of the tracks against the surface's equation, so a header
            whose parameters mean something other than what is assumed here fails
            loudly instead of producing a plausible, wrong plot.
        """
        surface = rssa_parameters.get_surface(surface_id)
        if projection is None:
            projection = Projection.from_surface(surface)

        plot = cls(tracks, rssa_parameters, projection)
        if len(rssa_parameters.surfaces) > 1:
            plot = plot.filter(on_surface(rssa_parameters, surface.id))
        if verify:
            plot = plot.verify()
        return plot

    @property
    def x_col(self) -> str:
        return self.projection.x_col

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

    @property
    def x_bins(self) -> np.ndarray:
        """The x bin edges."""
        if self._x_bins is None:
            raise ValueError(
                "Bins are not set. Call with_bins() or with_bin_width() first."
            )
        return self._x_bins

    @property
    def y_bins(self) -> np.ndarray:
        """The y bin edges."""
        if self._y_bins is None:
            raise ValueError(
                "Bins are not set. Call with_bins() or with_bin_width() first."
            )
        return self._y_bins

    def verify(self) -> RSSAPlot:
        """Checks a sample of the tracks against the projected surface's equation.

        Does nothing when the projection wasn't given enough geometry to place the
        surface in space (a hand-built plane with no `offset`, say); everything
        `Projection.from_surface()` builds can be checked.
        """
        residual = self.projection.residual()
        if residual is None:
            return self

        sample = (
            self.tracks.head(_SURFACE_VERIFICATION_SAMPLE)
            .select(residual.abs().max().alias("max_distance"))
            .collect(engine="streaming")
        )
        max_distance = sample.item(0, "max_distance")
        if max_distance is None:  # no tracks left on this surface, nothing to check
            return self
        tolerance = (
            _SURFACE_ABSOLUTE_TOLERANCE
            + _SURFACE_RELATIVE_TOLERANCE * self.projection.scale()
        )
        if max_distance > tolerance:
            raise ValueError(
                f"the tracks do not lie on {self.projection.describe()}: one sits"
                f" {max_distance:.4g} cm away (tolerance {tolerance:.4g} cm). Either"
                " the tracks were not filtered to this surface, or this package's"
                " interpretation of the header's surface parameters is wrong for this"
                " file. Pass verify=False to plot anyway."
            )
        return self

    def with_bins(
        self,
        x_edges: Sequence[float] | np.ndarray,
        y_edges: Sequence[float] | np.ndarray,
    ) -> RSSAPlot:
        """Bins the map on the given edges. See `with_bin_width()` to derive them from
        the data instead."""
        x_bins = np.sort(np.asarray(x_edges, dtype=np.float64))
        y_bins = np.sort(np.asarray(y_edges, dtype=np.float64))
        if x_bins.size < 2 or y_bins.size < 2:
            raise ValueError("at least two edges are needed on each axis")
        return self._replace(_x_bins=x_bins, _y_bins=y_bins)

    def with_bin_width(self, bin_width: float = 10.0) -> RSSAPlot:
        """Bins the map on square bins of roughly `bin_width` cm, covering the extent
        of the data - except where the projection knows better, e.g. a `Cylinder` bins
        its perimeter axis over the full circumference so the bins tile it exactly.
        """
        # Aggregate the extents directly instead of collecting every row: on a huge
        # (scan_tracks-backed) frame, pulling all rows just to call .min()/.max() in
        # Python would materialize the whole file for no reason.
        stats = self.tracks.select(
            pl.col(self.x_col).min().alias("x_min"),
            pl.col(self.x_col).max().alias("x_max"),
            pl.col(self.y_col).min().alias("y_min"),
            pl.col(self.y_col).max().alias("y_max"),
        ).collect(engine="streaming")
        x_min, x_max, y_min, y_max = stats.row(0)
        if x_min is None:
            raise ValueError("There are no tracks left to plot at this point.")

        x_edges, y_edges = self.projection.bin_edges(
            bin_width, x_min, x_max, y_min, y_max
        )
        return self.with_bins(x_edges, y_edges)

    def plot_parameters(self) -> PlotParameters:
        """The rendering options a result computed from this plot starts with: axis
        labels and title derived from the projection."""
        xlabel, ylabel = self.projection.labels()
        return PlotParameters(
            title=self.projection.describe(), xlabel=xlabel, ylabel=ylabel
        )

    def particle_flux(
        self,
        source_intensity: float,
        *,
        errors: bool = False,
        min_cosine: float = MCNP_GRAZING_COSINE,
    ) -> CurrentMap:
        """The scalar flux per bin, in units of `source_intensity` per cm2 (e.g.
        particles/cm2/s if `source_intensity` is in particles/s).

        This is MCNP's F2 (surface flux) quantity, and it is what response functions -
        flux-to-dose factors, KERMA, activation cross sections - are defined against.
        It is *not* what `particle_current()` gives: a track crossing at a shallow
        angle spends more path length in any thin volume straddling the surface than
        one crossing head-on, so each track is weighted by `1/|mu|`, where `mu` is the
        cosine of the angle between its direction and the surface normal there:

            flux = sum(w_i / |mu_i|) / (nps * A)

        The difference is not a detail. For an isotropic field the crossing directions
        are distributed as `mu dOmega`, and the mean of `1/|mu|` over that is exactly
        2 - so `particle_current()` understates the flux by a factor of two where the
        field is diffuse, while agreeing with it where the field is collimated. It
        therefore distorts the *shape* of a map, not just its scale.
        `VectorCurrent.anisotropy` bounds the discrepancy from below (see
        `vector_current()`).

        Grazing tracks (`|mu| -> 0`) get very large weight, which is real physics
        rather than a bug: it is the same effect behind MCNP's own grazing-angle
        variance problems on F2 tallies. Following MCNP, tracks below `min_cosine` are
        scored at `min_cosine / 2` instead of blowing up; `grazing_fraction()` reports
        how much of each bin's flux came from them. `errors=True` estimates how much a
        few such tracks are carrying.

        Requires the `u`, `v` and `w` columns.
        """
        weight = pl.col("weight") / self._grazing_limited_cosine(min_cosine)
        return self._current_map(
            weight,
            source_intensity,
            errors=errors,
            extra_columns=self.projection.incidence_columns(),
            legend_label="Flux [/cm2 per unit source]",
        )

    def particle_current(
        self, source_intensity: float, *, errors: bool = False
    ) -> CurrentMap:
        """The particle crossing rate per bin, in units of `source_intensity` per cm2
        (e.g. particles/cm2/s if `source_intensity` is in particles/s).

            current = sum(w_i) / (nps * A)

        This is MCNP's F1 quantity per unit area: the weight crossing each patch of
        surface, regardless of the angle it crosses at. Two things it is *not*:

        - it is not the flux. See `particle_flux()`, which applies the `1/|mu|`
          correction and can be twice this where the field is diffuse. If you are
          about to multiply a map by a dose or reaction-rate response function, you
          want that one.
        - it is not the *net* current: tracks crossing inward and outward both add to
          it, rather than cancelling. See `net_current()`. On a surface every track
          crosses the same way - the usual case for an `SSW` boundary - the two are
          equal, which makes comparing them a cheap check of that assumption.

        The per-track weights recorded in the RSSA file are relative to one source
        particle; this is turned into a physical rate by dividing the summed weight in
        each bin by the number of histories in the simulation (`abs(np1)`) and scaling
        by `source_intensity`.
        """
        return self._current_map(
            pl.col("weight"),
            source_intensity,
            errors=errors,
            legend_label="Crossing rate [/cm2 per unit source]",
        )

    def net_current(
        self, source_intensity: float, *, errors: bool = False
    ) -> CurrentMap:
        """The net current through the surface per bin, `J . n`, in units of
        `source_intensity` per cm2, signed:

            net current = sum(sign(mu_i) * w_i) / (nps * A)

        Tracks crossing one way count positive and the other way negative, so this is
        what actually leaks through each patch - as opposed to `particle_current()`,
        where both senses add up. A bin where as much comes back as goes out reads
        zero here and reads large there.

        Positive means "along the projection's normal": outward for a cylinder, along
        `axis` for an `AxisAlignedPlane`, along `normal` for a `TiltedPlane`.

        Unlike a flux, no `1/|mu|` weighting is involved - a crossing already *is* a
        cosine-weighted event - so this is immune to the grazing-angle blow-up.
        Rendered on a linear scale by default, since it can be negative.

        Requires the `u`, `v` and `w` columns.
        """
        signed = pl.when(self._incidence_cosine() >= 0).then(1.0).otherwise(
            -1.0
        ) * pl.col("weight")
        return self._current_map(
            signed,
            source_intensity,
            errors=errors,
            extra_columns=self.projection.incidence_columns(),
            norm="linear",
            legend_label="Net current [/cm2 per unit source]",
        )

    def grazing_fraction(self, min_cosine: float = MCNP_GRAZING_COSINE) -> CurrentMap:
        """The fraction of each bin's flux contributed by tracks crossing at or below
        `min_cosine`, i.e. by the ones `particle_flux()` had to cap.

        A diagnostic, not a physical quantity: values in `[0, 1]`, where anything much
        above 0 means that bin's flux rests on the capped approximation rather than on
        the estimator proper, and should be read together with the statistical error.
        Bins with no tracks are `NaN`.
        """
        cosine = self._grazing_limited_cosine(min_cosine)
        contribution = pl.col("weight") / cosine
        binned = self._binned(["weight", *self.projection.incidence_columns()])
        grid = (
            binned.lazy()
            .with_columns(contribution.alias("contribution"))
            .group_by(["bin_x", "bin_y"])
            .agg(
                pl.col("contribution").sum().alias("total"),
                pl.col("contribution")
                .filter(self._incidence_cosine().abs() < min_cosine)
                .sum()
                .alias("grazing"),
            )
            .with_columns((pl.col("grazing") / pl.col("total")).alias("fraction"))
            .collect()
        )
        return CurrentMap(
            self._raster(grid, "fraction", fill_value=np.nan),
            self.x_bins,
            self.y_bins,
            dataclasses.replace(
                self.plot_parameters(),
                norm="linear",
                vmin=0.0,
                vmax=1.0,
                legend_label=f"Flux fraction from tracks below |mu| = {min_cosine:g}",
            ),
        )

    def _incidence_cosine(self) -> pl.Expr:
        """`mu` for each track, or a clear error if the projection can't say."""
        cosine = self.projection.incidence_cosine()
        if cosine is None:
            raise ValueError(
                f"{type(self.projection).__name__} does not know where its surface"
                " normal points, so the angle each track crosses at is unknown and"
                " neither a flux nor a net current can be estimated. Only the"
                " direction-agnostic particle_current() is available."
            )
        self._require_columns(
            self.projection.incidence_columns(),
            "Weighting tracks by the angle they cross at",
        )
        return cosine

    def _require_columns(self, columns: Sequence[str], what: str) -> None:
        present = set(self.tracks.collect_schema().names())
        missing = [column for column in columns if column not in present]
        if missing:
            raise ValueError(
                f"{what} requires the {missing} column(s), which are not present in"
                " these tracks (they may have been dropped, or left out of"
                " read_from_file(columns=...))."
            )

    def _grazing_limited_cosine(self, min_cosine: float) -> pl.Expr:
        """`|mu|`, with grazing tracks pinned away from zero the way MCNP does it.

        MCNP scores a track whose `|mu|` falls below `1e-3` as if it were `5e-4` - the
        midpoint of the capped range, rather than its edge, which keeps the mean of a
        uniform spread of grazing directions right instead of systematically
        understating them. `particle_flux()` follows that, since the alternative is an
        infinity in the tally.
        """
        if not min_cosine > 0:
            raise ValueError(f"min_cosine must be positive, got {min_cosine!r}")
        cosine = self._incidence_cosine().abs()
        return pl.when(cosine < min_cosine).then(min_cosine / 2).otherwise(cosine)

    def _current_map(
        self,
        weight: pl.Expr,
        source_intensity: float,
        *,
        errors: bool,
        extra_columns: Sequence[str] = (),
        norm: Literal["linear", "log"] | None = None,
        legend_label: str = "",
    ) -> CurrentMap:
        """Sums `weight` per bin, normalizes it to a rate per unit area, and optionally
        estimates its statistical error in the same pass.

        Every per-bin quantity this class computes differs only in `weight` - raw
        track weight for a crossing rate, `w/|mu|` for a flux, `sign(mu)*w` for a net
        current - so they share the aggregation, and an error can only ever be
        estimated for the quantity it is returned with.
        """
        columns = ["weight", *extra_columns]
        if errors:
            columns.append("history")
            if "history" not in self.tracks.collect_schema().names():
                raise ValueError(
                    "Estimating the statistical error requires the 'history' column, "
                    "which is not present in these tracks (e.g. it may have been "
                    "dropped with `.drop(...)`)."
                )

        binned = self._binned(columns).lazy().with_columns(weight.alias("_w"))
        nps = abs(self.rssa_parameters.np1)
        normalization = (
            source_intensity / nps / calculate_areas(self.x_bins, self.y_bins)
        )

        if not errors:
            grid = binned.group_by(["bin_x", "bin_y"]).agg(pl.col("_w").sum()).collect()
            values = self._raster(grid, "_w") * normalization
            relative_errors = None
        else:
            grid = (
                binned
                # Sum contributions per history first: multiple tracks from the same
                # history landing in the same bin are correlated, not independent
                # samples, so they must be combined before estimating the variance.
                .group_by(["bin_x", "bin_y", "history"])
                .agg(pl.col("_w").sum().alias("history_weight"))
                .group_by(["bin_x", "bin_y"])
                .agg(
                    pl.col("history_weight").sum().alias("sum_w"),
                    (pl.col("history_weight") ** 2).sum().alias("sum_w2"),
                )
                .with_columns(
                    # Clip at 0: floating-point round-off can otherwise push a
                    # near-zero-variance bin's estimate very slightly negative.
                    (pl.col("sum_w2") / pl.col("sum_w") ** 2 - 1 / nps)
                    .clip(lower_bound=0.0)
                    .sqrt()
                    .alias("relative_error")
                )
                .collect()
            )
            values = self._raster(grid, "sum_w") * normalization
            # Bins with no contributing history are genuinely undefined, not zero
            # error; fill with NaN rather than silently claiming perfect statistics.
            relative_errors = self._raster(grid, "relative_error", fill_value=np.nan)

        parameters = self.plot_parameters()
        parameters = dataclasses.replace(parameters, legend_label=legend_label)
        if norm is not None:
            parameters = dataclasses.replace(parameters, norm=norm)
        return CurrentMap(values, self.x_bins, self.y_bins, parameters, relative_errors)

    def vector_current(
        self,
        source_intensity: float,
        *,
        min_cosine: float = MCNP_GRAZING_COSINE,
    ) -> VectorCurrent:
        """The net vector current `J` per spatial bin: the direction radiation actually
        flows through each patch of surface, and how directional that flow is.

            J = sum(w_i * Omega_i / |mu_i|) / (nps * A)

        The `1/|mu|` is the same correction `particle_flux()` applies, and it is what
        makes this the current density of transport theory rather than a descriptive
        average of crossing directions. A set of crossings samples the angular flux
        weighted by `|mu|` (head-on crossings are over-represented, because a patch of
        surface presents its full area to them), so summing raw `w * Omega` estimates
        the *second* angular moment; dividing that weighting back out gives the first,
        which is `J`. The two agree only at normal incidence.

        `VectorCurrent.table` has one row per non-empty spatial bin:

        - `pos_x`, `pos_y`, `pos_z`: mean physical position of the bin's tracks.
        - `dir_u`, `dir_v`, `dir_w`: the components of `J`, in the same units as
          `particle_flux()` (e.g. particles/cm2/s if `source_intensity` is in
          particles/s) - *not* normalized to a unit vector, so a bin whose particles
          point in inconsistent directions has a short vector here even if its raw
          particle count is high.
        - `magnitude`: `|J|`, i.e. `sqrt(dir_u^2 + dir_v^2 + dir_w^2)`.
        - `net_current`: `J . n`, the component through the surface, signed the same
          way as `net_current()`. Distinct from `magnitude`, which includes the flow
          *along* the surface.
        - `flux`: the scalar flux in the bin, the same quantity `particle_flux()`
          maps.
        - `crossing_rate`: the direction-agnostic quantity `particle_current()` maps,
          for reference.
        - `anisotropy`: `magnitude / flux`, in [0, 1] - the flux-to-current ratio. 1
          means every particle in the bin travels in (nearly) the same direction; 0.5
          is a half-isotropic (cosine) distribution; near 0 means directions cancel
          out, so this bin's direction shouldn't be read as meaningful.
        - `mean_energy`: the weighted mean energy in the bin, in MeV.

        `anisotropy` also bounds how much `particle_current()` understates the flux:
        since `1/|mu|` is convex, Jensen's inequality gives
        `flux / crossing_rate >= 1 / anisotropy`.

        Requires the `u`, `v` and `w` columns.
        """
        self._require_columns(("u", "v", "w"), "Summing track directions")
        cosine = self._grazing_limited_cosine(min_cosine)
        sign = pl.when(self._incidence_cosine() >= 0).then(1.0).otherwise(-1.0)

        binned = (
            self._binned(
                [
                    "x",
                    "y",
                    "z",
                    "u",
                    "v",
                    "w",
                    "weight",
                    "energy",
                    *self.projection.incidence_columns(),
                ]
            )
            .lazy()
            .with_columns(
                (pl.col("weight") / cosine).alias("_flux_weight"),
                sign.alias("_sign"),
            )
        )
        grid = (
            binned.group_by(["bin_x", "bin_y"])
            .agg(
                pl.col("x").mean().alias("pos_x"),
                pl.col("y").mean().alias("pos_y"),
                pl.col("z").mean().alias("pos_z"),
                (pl.col("_flux_weight") * pl.col("u")).sum().alias("sum_ju"),
                (pl.col("_flux_weight") * pl.col("v")).sum().alias("sum_jv"),
                (pl.col("_flux_weight") * pl.col("w")).sum().alias("sum_jw"),
                pl.col("_flux_weight").sum().alias("sum_flux"),
                (pl.col("_sign") * pl.col("weight")).sum().alias("sum_net"),
                pl.col("weight").sum().alias("sum_w"),
                (pl.col("weight") * pl.col("energy")).sum().alias("sum_w_energy"),
            )
            .collect()
        )

        if grid.is_empty():
            empty = pl.DataFrame(
                schema={name: pl.Float64 for name in VECTOR_CURRENT_COLUMNS}
            )
            return VectorCurrent(empty)

        nps = abs(self.rssa_parameters.np1)
        normalization = source_intensity / nps
        areas = calculate_areas(self.x_bins, self.y_bins)
        area_per_row = areas[grid["bin_y"].to_numpy(), grid["bin_x"].to_numpy()]
        per_area = normalization / area_per_row

        dir_u = grid["sum_ju"].to_numpy() * per_area
        dir_v = grid["sum_jv"].to_numpy() * per_area
        dir_w = grid["sum_jw"].to_numpy() * per_area
        flux = grid["sum_flux"].to_numpy() * per_area
        net_current = grid["sum_net"].to_numpy() * per_area
        crossing_rate = grid["sum_w"].to_numpy() * per_area
        magnitude = np.sqrt(dir_u**2 + dir_v**2 + dir_w**2)
        mean_energy = (grid["sum_w_energy"] / grid["sum_w"]).to_numpy()
        anisotropy = np.divide(
            magnitude, flux, out=np.zeros_like(magnitude), where=flux > 0
        )

        table = pl.DataFrame(
            {
                "pos_x": grid["pos_x"].to_numpy(),
                "pos_y": grid["pos_y"].to_numpy(),
                "pos_z": grid["pos_z"].to_numpy(),
                "dir_u": dir_u,
                "dir_v": dir_v,
                "dir_w": dir_w,
                "magnitude": magnitude,
                "net_current": net_current,
                "flux": flux,
                "crossing_rate": crossing_rate,
                "anisotropy": anisotropy,
                "mean_energy": mean_energy,
            }
        )
        return VectorCurrent(table)

    def angular_distribution(
        self,
        source_intensity: float,
        *,
        polar_bins: int = 8,
        azimuthal_bins: int = 16,
        min_cosine: float = MCNP_GRAZING_COSINE,
        area: float | None = None,
    ) -> AngularDistribution:
        """The angular flux `psi(Omega)` of these tracks, resolved in the surface's own
        frame and pooled over every spatial bin.

            psi(Omega) = sum(w_i / |mu_i|) / (nps * A * dOmega)

        This is the *distribution* of directions behind the single arrow per bin that
        `vector_current()` gives. The distinction matters because a mean vector cannot
        tell an isotropic field from two opposing beams - both average to `|J| = 0` -
        so where `VectorCurrent.anisotropy` is low, this is what says which of the two
        you have.

        The `1/|mu|` is not cosmetic. A set of crossings samples the angular flux
        weighted by `|mu|`, because a patch of surface presents its full area only to
        head-on tracks; summing raw weight per direction would therefore draw a dent
        at grazing angles that is an artefact of the sampling rather than a feature of
        the field - and grazing angles are exactly where an angular plot is read. It is
        the same correction, with the same `min_cosine` cap, that `particle_flux()`
        applies; `grazing_fraction()` is the diagnostic for how much rests on the cap.

        Directions are binned on `mu`, the cosine to the surface normal, and `phi`, the
        azimuth in the tangent plane measured from `Projection.tangent_basis()` - which
        is aligned with the raster, so `phi = 0` points the way this plot's x axis
        grows. `polar_bins` counts bins per hemisphere and `azimuthal_bins` the sectors
        around, centred on the tangent directions rather than divided by them; the `mu`
        grid is uniform (not the polar angle), so every cell subtends the same solid
        angle and the values need no per-cell correction to be compared.

        Only tracks inside the plot's bins contribute, so filtering or narrowing the
        bins zooms this in on a region.

        `A` is the area of surface those tracks are spread over, and it is taken from
        the extent they actually span rather than from the bin grid. That is what keeps
        the result independent of `bin_width`, which is a rendering choice and has no
        business changing a physical quantity: `with_bin_width()` pads its last bin out
        past the data, so on tracks spanning 60 cm a 100 cm-wide bin would claim a strip
        two thirds empty and dilute every value by 1.67. Only the level moves - the
        shape is a single histogram scaled by one number - but a level that drifts with
        a cosmetic parameter is worse than one that is merely arbitrary.

        Pass `area` to state the patch explicitly instead. That is the right thing to
        do when the region is a deliberate choice rather than wherever the tracks
        happen to land - characterizing a defined 200 x 200 window that the particles
        only half fill, say, where the empty half is part of the answer.

        Requires the `u`, `v` and `w` columns (a plane needs only some of them).
        """
        if polar_bins < 1 or azimuthal_bins < 1:
            raise ValueError(
                "polar_bins and azimuthal_bins must both be at least 1, got"
                f" {polar_bins!r} and {azimuthal_bins!r}"
            )
        tangent = self.projection.tangent_basis()
        if tangent is None:
            raise ValueError(
                f"{type(self.projection).__name__} does not know how its surface is"
                " oriented in space, so a track's direction cannot be resolved into an"
                " azimuth on it and no angular distribution can be estimated."
            )
        self._require_columns(
            self.projection.tangent_columns(),
            "Resolving track directions onto the surface's tangent plane",
        )
        along_x, along_y = tangent
        cosine = self._incidence_cosine()
        flux_weight = pl.col("weight") / self._grazing_limited_cosine(min_cosine)

        if area is not None and not area > 0:
            raise ValueError(f"area must be positive, got {area!r}")

        binned = self._binned(
            [
                "weight",
                *self.projection.incidence_columns(),
                *self.projection.tangent_columns(),
            ]
        )
        directions = (
            binned.lazy()
            .select(
                cosine.alias("_mu"),
                # A track crossing exactly normally has no tangential part and so no
                # defined azimuth; arctan2 puts it at phi = 0. That is the pole of the
                # plot, where every azimuth meets, so the choice is immaterial there.
                _wrapped_azimuth(pl.arctan2(along_y, along_x), azimuthal_bins).alias(
                    "_phi"
                ),
                flux_weight.alias("_weight"),
            )
            .collect(engine="streaming")
        )

        mu_bins = np.linspace(-1.0, 1.0, 2 * polar_bins + 1)
        # Offset by half a sector so the tangent directions are sector *centres*, not
        # edges: a beam aimed straight along the plot's x axis should read as one full
        # petal rather than as two half ones either side of a boundary.
        half_sector = np.pi / azimuthal_bins
        phi_bins = np.linspace(
            -np.pi - half_sector, np.pi - half_sector, azimuthal_bins + 1
        )
        summed, _, _ = np.histogram2d(
            # Clipped because a curved projection divides by the track's own distance
            # to the axis, which can hand back a cosine a few ulps outside [-1, 1] and
            # drop a legitimately normal track out of the histogram entirely.
            np.clip(directions["_mu"].to_numpy(), -1.0, 1.0),
            directions["_phi"].to_numpy(),
            bins=(mu_bins, phi_bins),
            weights=directions["_weight"].to_numpy(),
        )

        nps = abs(self.rssa_parameters.np1)
        covered = self._covered_area(binned) if area is None else area
        solid_angle = (2.0 / (2 * polar_bins)) * (2.0 * np.pi / azimuthal_bins)
        values = summed * (source_intensity / nps / covered / solid_angle)

        return AngularDistribution(
            values,
            mu_bins,
            phi_bins,
            AngularPlotParameters(
                title=f"Angular flux - {self.projection.describe()}",
                tangent_labels=self.projection.tangent_labels(),
            ),
        )

    def _covered_area(self, binned: pl.DataFrame) -> float:
        """The area of surface a pooled quantity is spread over, in cm2: the extent the
        binned tracks actually span on each raster axis.

        Taken from the tracks rather than from the bin grid because the grid is padded
        - `_edges_covering()` runs its last bin past the data so every bin is exactly
        `bin_width` wide - and a pooled quantity divided by the padded area drifts with
        `bin_width`. A per-bin map is not affected: there the padding lands in one edge
        bin, which really does cover surface no track reached, and dividing it by its
        own full area is right.

        An axis the tracks span no distance on - a single track, or several at one
        coordinate - carries no extent to measure, so the bins are used for that axis
        instead; they are the only statement of the region available.
        """
        spans = []
        for column, bins in ((self.x_col, self.x_bins), (self.y_col, self.y_bins)):
            values = binned.get_column(column).to_numpy()
            span = float(values.max() - values.min()) if values.size else 0.0
            spans.append(span or float(bins[-1]) - float(bins[0]))
        return spans[0] * spans[1]

    def _binned(self, columns: Sequence[str] = ()) -> pl.DataFrame:
        """Collects the tracks that fall inside the bins, with the `bin_x`/`bin_y`
        indices each one lands in.

        The requested columns are de-duplicated against the projection's own: a
        `Cylinder`'s `y_col` is a real track column (e.g. `"z"`) that may also be
        needed in its own right, and Polars rejects selecting the same name twice.
        """
        x_bins, y_bins = self.x_bins, self.y_bins
        wanted = list(dict.fromkeys([self.x_col, self.y_col, *columns]))
        frame = (
            self.tracks.select(wanted)
            .filter(
                pl.col(self.x_col).is_between(x_bins[0], x_bins[-1], closed="left"),
                pl.col(self.y_col).is_between(y_bins[0], y_bins[-1], closed="left"),
            )
            .collect(engine="streaming")
        )
        bin_x = np.searchsorted(x_bins, frame[self.x_col].to_numpy(), side="right") - 1
        bin_y = np.searchsorted(y_bins, frame[self.y_col].to_numpy(), side="right") - 1
        return frame.with_columns(pl.Series("bin_x", bin_x), pl.Series("bin_y", bin_y))

    def _raster(
        self, grid: pl.DataFrame, value_column: str, fill_value: float = 0.0
    ) -> np.ndarray:
        """Scatters a per-bin aggregate back into a dense `[y, x]` raster."""
        raster = np.full(
            (len(self.y_bins) - 1, len(self.x_bins) - 1), fill_value, dtype=np.float64
        )
        raster[
            grid.get_column("bin_y").to_numpy(), grid.get_column("bin_x").to_numpy()
        ] = grid.get_column(value_column).to_numpy()
        return raster


#: A coarse default log-spaced energy grid (1e-5 eV to 20 MeV, 175 bins), used only
#: when no `energy_bins` are given to `RSSASpectraPlot`. Unlike the original F4Enix
#: implementation, this package does not bundle a specific standard group structure
#: (e.g. VITAMIN-J-175): pass `with_energy_bins()` explicitly with the group structure
#: your analysis requires.
DEFAULT_ENERGY_BINS = np.geomspace(1e-5, 2e7, 176)


class RSSASpectraPlot(TrackPlot):
    """The energy spectrum of a set of tracks.

    Unlike `RSSAPlot` this needs no geometry: an energy spectrum is well defined for
    any selection of tracks, including one spanning several surfaces.
    """

    def __init__(
        self,
        tracks: pl.DataFrame | pl.LazyFrame,
        rssa_parameters: FileParameters,
        *,
        energy_bins: Sequence[float] | np.ndarray = DEFAULT_ENERGY_BINS,
    ) -> None:
        self.tracks = tracks.lazy()
        self.rssa_parameters = rssa_parameters
        self.energy_bins = np.asarray(energy_bins, dtype=np.float64)

    def with_energy_bins(
        self, energy_bins: Sequence[float] | np.ndarray
    ) -> RSSASpectraPlot:
        """Bins the spectrum on the given energy edges, in eV."""
        return self._replace(energy_bins=np.asarray(energy_bins, dtype=np.float64))

    def spectrum(self) -> Spectrum:
        """Computes the spectrum: the summed track weight in each energy bin, divided
        by the bin's lethargy width and normalized so the total is 1."""
        data = self.tracks.select(
            (pl.col("energy") * 1e6).alias("energy"),  # MeV in the file, eV here
            pl.col("weight"),
        ).collect(engine="streaming")

        weighted_counts, _ = np.histogram(
            data["energy"],
            bins=self.energy_bins,
            weights=data["weight"],  # each contribution is weighted by its weight
        )

        lethargies = np.log(self.energy_bins[1:] / self.energy_bins[:-1])
        normalized_counts = (
            weighted_counts
            / lethargies  # The Y axis is in units of counts per lethargy
            / np.sum(weighted_counts)  # Normalize the total number of counts to 1
        )
        return Spectrum(
            normalized_counts=normalized_counts, energy_bins=self.energy_bins
        )


def _wrapped_azimuth(azimuth: pl.Expr, azimuthal_bins: int) -> pl.Expr:
    """`azimuth` folded into the half-sector-shifted range the bins cover.

    The sectors are centred on the tangent directions, so they run from
    `-pi - dphi/2` to `pi - dphi/2`; the sliver of `arctan2` output above that top edge
    belongs to the sector straddling the seam, and is brought back by a turn.
    """
    half_sector = np.pi / azimuthal_bins
    return (
        pl.when(azimuth >= np.pi - half_sector)
        .then(azimuth - 2.0 * np.pi)
        .otherwise(azimuth)
    )


def _deprecated(what: str, instead: str) -> None:
    warnings.warn(
        f"{what} is deprecated and will be removed in a future release; use {instead}"
        " instead. Track selection is now done with ordinary Polars expressions - see"
        " aci_rssa.expressions.",
        DeprecationWarning,
        stacklevel=3,
    )
