"""Notebook diagnostics for a trained model's sample quality, fully streaming.

Every check here reads directly from a `giant predict --coord local` parquet
file (`pred_*`/`true_*` columns, denormalized but still local-frame/log-scaled)
and is computed with **lazy, streaming polars** so peak memory stays bounded no
matter how large the file is (production predict output runs to tens of GB).
There is no in-memory `SampleCollection`, no full-array materialization, and no
on-the-fly (checkpoint + live sampler) path: generate predictions once via the
CLI, then run every diagnostic below against that file.

Each public function takes `source: str | Path | pl.LazyFrame` — a path to the
predict parquet, or a pre-built `LazyFrame` with the same columns (for tests) —
and returns either a matplotlib figure or a small reduced table/dataclass.

Four tiers of checks::

    from giant.analysis import plot_kl_bars_pl, plot_marginals
    from giant.analysis import plot_correlation_matrices, plot_pairwise
    from giant.analysis import plot_direction_alignment, plot_constraint_violations

    FILE = "path/to/steps_predicted_local.parquet"

    # Tier 1 — stratified marginals (per-dimension real-vs-generated, sliced by
    #          pdg / material / energy so failures hidden by the aggregate show)
    for grouping in (None, "energy", "pdg", "material"):
        plot_kl_bars_pl(FILE, group_by=grouping)
    plot_marginals(FILE)
    plot_marginals(FILE, group_by="energy")

    # Tier 2 — joint structure (correlations + physically-coupled pairs + the
    #          post/travel direction alignment marginals can't see)
    plot_correlation_matrices(FILE)
    plot_pairwise(FILE)
    plot_direction_alignment(FILE)

    # Tier 3 — physical constraints (unit-norm directions, non-negative scalars;
    #          any violation is a pure generation artifact of the unconstrained MLP)
    plot_constraint_violations(FILE)

For event-level (shower) observables — total deposited energy, total length,
longitudinal/transverse profiles, shower-max depth — aggregated per `event_id`
in the world frame with physical units (mm, MeV)::

    from giant.analysis import compute_event_observables_pl
    from giant.analysis import plot_total_energy, plot_total_length
    from giant.analysis import plot_mean_energy_per_step, plot_mean_length_per_step
    from giant.analysis import plot_longitudinal_profile, plot_transverse_profile
    from giant.analysis import plot_shower_max_depth

    obs = compute_event_observables_pl(FILE)
    plot_total_energy(obs)
    plot_total_length(obs)
    plot_longitudinal_profile(obs)
    plot_transverse_profile(obs)
    plot_shower_max_depth(obs)

This re-aggregates one-step-ahead generations (each row generated conditioned on
the *real* preceding state) grouped by event — not a full autoregressive shower
rollout — so it won't surface covariate-shift failures that only appear under
true rollout, only how well one-step generation reconstructs aggregate shower
structure when fed real conditioning throughout.

For the dataset-wide breakdown of which particle species (pdg) contributed how
much of the total energy/length::

    from giant.analysis import pdg_contribution_table_pl
    from giant.analysis import plot_pdg_energy_share, plot_pdg_length_share

    table = pdg_contribution_table_pl(FILE)
    plot_pdg_energy_share(table)
    plot_pdg_length_share(table)
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import polars as pl

import pyarrow.parquet as pq

from giant.constants import (
    LOCAL_TARGET_NAMES,
    PREDICT_COORD_METADATA_KEY,
    PREDICT_SCHEMA_VERSION,
    PREDICT_SCHEMA_VERSION_KEY,
)

# The first 3 target dims are the scalar (non-direction) outputs. In raw/physical
# space they are step_length, delta_e, edep; in the model's native target space the
# first is log_step_length and the next two are the deposit/secondary ALR energy
# logits (see giant.constants.LOCAL_TARGET_NAMES and energy_simplex_encode).
_N_SCALAR_DIMS = 3
RAW_TARGET_NAMES = [
    "step_length",
    "delta_e",
    "edep",
    "post_dx",
    "post_dy",
    "post_dz",
    "travel_dx",
    "travel_dy",
    "travel_dz",
]
_LOG_EPS = (
    1e-8  # mirrors giant.data.transforms._EPS, duplicated for use in polars exprs
)


# ---------------------------------------------------------------------------
# Predict-parquet scanning, metadata verification, and physical-value exprs
# ---------------------------------------------------------------------------


def _check_predict_metadata(path: Path) -> None:
    """Verify a parquet file's giant-predict tag before trusting its column layout.

    Raises rather than warns: a wrong or missing tag means the column-layout
    assumptions below don't hold, so silently proceeding could mix up which
    columns are predictions vs. ground truth.
    """
    metadata = pq.read_schema(path).metadata or {}
    coord = metadata.get(PREDICT_COORD_METADATA_KEY.encode())
    if coord is None:
        raise ValueError(
            f"{path} has no '{PREDICT_COORD_METADATA_KEY}' parquet metadata — it wasn't "
            "written by `giant predict` (or predates schema tagging), so its column "
            "layout can't be verified"
        )
    if coord.decode() != "local":
        raise ValueError(
            f"{path} was written with --coord {coord.decode()!r}, not 'local' — "
            "giant.analysis only supports coord=local predict output, which is the "
            "only mode that also writes ground-truth columns"
        )
    version = metadata.get(PREDICT_SCHEMA_VERSION_KEY.encode())
    if version is not None and version.decode() != PREDICT_SCHEMA_VERSION:
        raise ValueError(
            f"{path} has predict schema version {version.decode()!r}, but "
            f"giant.analysis expects {PREDICT_SCHEMA_VERSION!r} — column layout may "
            "have changed; update giant.analysis to match"
        )


def _scan_predicted_local(source: str | Path | pl.LazyFrame) -> pl.LazyFrame:
    if isinstance(source, pl.LazyFrame):
        return source
    path = Path(source)
    _check_predict_metadata(path)
    return pl.scan_parquet(path)


def _edep_pl(prefix: str) -> pl.Expr:
    """Physical edep from a `{prefix}_edep_logit`/`{prefix}_sec_logit` pair + `pre_E`.

    Polars equivalent of `energy_simplex_decode(...)[0]` (the deposit component):
    a softmax over `[z_edep, z_sec, 0]` times pre_E.
    """
    z1, z2 = pl.col(f"{prefix}_edep_logit"), pl.col(f"{prefix}_sec_logit")
    m = pl.max_horizontal(z1, z2, pl.lit(0.0))
    e1, e2, e3 = (z1 - m).exp(), (z2 - m).exp(), (pl.lit(0.0) - m).exp()
    return (e1 / (e1 + e2 + e3)) * pl.col("pre_E")


def _delta_e_pl(prefix: str) -> pl.Expr:
    """Physical delta_e (= edep + e_sec = pre_E - post_E) from the ALR logits + pre_E."""
    z1, z2 = pl.col(f"{prefix}_edep_logit"), pl.col(f"{prefix}_sec_logit")
    m = pl.max_horizontal(z1, z2, pl.lit(0.0))
    e1, e2, e3 = (z1 - m).exp(), (z2 - m).exp(), (pl.lit(0.0) - m).exp()
    return ((e1 + e2) / (e1 + e2 + e3)) * pl.col("pre_E")


def _raw_dim_expr(prefix: str, j: int) -> pl.Expr:
    """Physical raw value of target dim j (`RAW_TARGET_NAMES[j]`) from a predict parquet.

    Dim 0 is de-logged step_length; dims 1–2 are the physical `delta_e`/`edep`
    decoded from the deposit/secondary ALR logits against `pre_E`; dims ≥3 are
    direction components, used as-is. `prefix` is "true" or "pred".
    """
    if j == 0:
        return pl.col(f"{prefix}_log_step_length").exp() - _LOG_EPS
    if j == 1:
        return _delta_e_pl(prefix)
    if j == 2:
        return _edep_pl(prefix)
    return pl.col(f"{prefix}_{LOCAL_TARGET_NAMES[j]}")


def _hist_edges(*arrays: np.ndarray, bins: int) -> np.ndarray:
    """Bin edges that don't blow up on near-constant data (e.g. a tight unit-norm cluster).

    Plain `np.linspace(lo, hi, bins+1)` raises when `hi - lo` is too small relative
    to float precision to support `bins` distinct edges, which is a real failure
    mode here (a well-trained model can push direction norms to within float32
    epsilon of 1.0), not just a test artifact. Used only on small per-event
    reduced arrays, never on raw file rows.
    """
    lo = min(a.min() for a in arrays)
    hi = max(a.max() for a in arrays)
    if not (hi - lo > 1e-6 * max(abs(hi), 1.0)):
        lo, hi = lo - 0.5, hi + 0.5
    return np.linspace(lo, hi, bins + 1)


# ---------------------------------------------------------------------------
# Streaming histogram / KL primitives
# ---------------------------------------------------------------------------


def _kl_from_counts(
    real_counts: np.ndarray, gen_counts: np.ndarray, eps: float = 1e-8
) -> float:
    """KL(real || gen) from two aligned histogram bin-count arrays.

    Add-eps smoothing then normalize both to distributions, as is standard for a
    histogram-estimated KL; the counts here come from a polars `group_by`
    aggregation rather than `np.histogram`.
    """
    p = real_counts.astype(np.float64) + eps
    q = gen_counts.astype(np.float64) + eps
    p /= p.sum()
    q /= q.sum()
    return float(np.sum(p * np.log(p / q)))


def _resolve_hist_range(
    lf: pl.LazyFrame, expr: pl.Expr, lo: float | None, hi: float | None
) -> tuple[float, float]:
    """Fill in any missing (lo, hi) from a streaming min/max pass over `expr`.

    Mirrors `_hist_edges`'s degenerate-range widening (±0.5 when the span is too
    tight to support distinct edges) so fixed-edge and data-ranged histograms
    agree on edge cases. Only runs the extra streaming pass when a bound is
    actually unknown.
    """
    if lo is not None and hi is not None:
        return lo, hi
    row = (
        lf.select(expr.min().alias("lo"), expr.max().alias("hi"))
        .collect(engine="streaming")
        .row(0, named=True)
    )
    data_lo, data_hi = row["lo"], row["hi"]
    if not (data_hi - data_lo > 1e-6 * max(abs(data_hi), 1.0)):
        data_lo, data_hi = data_lo - 0.5, data_hi + 0.5
    return (data_lo if lo is None else lo, data_hi if hi is None else hi)


def _streaming_hist1d(
    lf: pl.LazyFrame,
    expr: pl.Expr,
    bins: int,
    lo: float | None = None,
    hi: float | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    """Bounded-memory 1D histogram of `expr` over `lf`, returning (counts, edges).

    A single streaming `group_by(bin).len()` pass — a proper single hash-pass
    histogram — after (optionally) a streaming min/max pass to fix the range
    when `lo`/`hi` aren't supplied. Never materializes the underlying column.
    """
    lo, hi = _resolve_hist_range(lf, expr, lo, hi)
    width = hi - lo
    bin_expr = ((expr - lo) / width * bins).floor().cast(pl.Int64).clip(0, bins - 1)
    hist = (
        lf.select(bin_expr.alias("bin"))
        .group_by("bin")
        .agg(pl.len().alias("count"))
        .collect(engine="streaming")
    )
    counts = np.zeros(bins, dtype=np.int64)
    for bin_idx, count in hist.iter_rows():
        counts[bin_idx] = count
    return counts, np.linspace(lo, hi, bins + 1)


def _step_density(counts: np.ndarray, edges: np.ndarray) -> np.ndarray:
    """Density-normalized bin heights (integrate to 1), matching `hist(density=True)`."""
    total = counts.sum()
    if total == 0:
        return counts.astype(np.float64)
    width = edges[1] - edges[0]
    return counts.astype(np.float64) / (total * width)


# ---------------------------------------------------------------------------
# Tier 1: stratified marginals
# ---------------------------------------------------------------------------
#
# One native `group_by` per dim (covering every group's mean/std/n/histogram
# range at once) plus one pass for per-bin counts — runtime independent of group
# cardinality. Filtering and re-collecting once per (group, dim) pair instead
# would make runtime scale with the number of groups: on a 114M-row file with
# 138 distinct pdg codes that extrapolated to ~an hour vs ~1 minute here.


def _add_group_label(
    lf: pl.LazyFrame, group_by: str | None, n_energy_bins: int
) -> pl.LazyFrame:
    """Add a `_group` string column identifying each row's stratum.

    For "pdg"/"material" the label is a direct string expr over the existing
    column — `group_by("_group")` downstream discovers the distinct values
    itself. "energy" needs one streaming pass over `pre_E` to fix quantile bin
    edges before a label can be assigned per row.
    """
    if group_by is None:
        return lf.with_columns(pl.lit("all").alias("_group"))
    if group_by == "pdg":
        return lf.with_columns(
            (pl.lit("pdg=") + pl.col("pdg").cast(pl.Int64).cast(pl.Utf8)).alias(
                "_group"
            )
        )
    if group_by == "material":
        return lf.with_columns(
            (pl.lit("material=") + pl.col("material")).alias("_group")
        )
    if group_by == "energy":
        pre_E = lf.select("pre_E").collect(engine="streaming").to_series().to_numpy()
        edges = np.quantile(pre_E, np.linspace(0, 1, n_energy_bins + 1))
        edges[-1] += 1e-6
        labels = [
            f"E∈[{edges[i]:.3g},{edges[i + 1]:.3g})" for i in range(n_energy_bins)
        ]
        expr = pl.when(pl.col("pre_E") < edges[1]).then(pl.lit(labels[0]))
        for i in range(1, n_energy_bins - 1):
            expr = expr.when(pl.col("pre_E") < edges[i + 1]).then(pl.lit(labels[i]))
        expr = expr.otherwise(pl.lit(labels[-1]))
        return lf.with_columns(expr.alias("_group"))
    raise ValueError(f"unknown group_by={group_by!r}")


# On a file larger than RAM (never page-cached) the binding constraint is *peak
# memory*, not scan count: each streaming scan mmaps the columns it reads, so a
# scan that touches only ~3 columns keeps resident memory low, while one that
# reads all 9 dims' ~19 source columns at once mmaps almost the whole 32GB file
# and OOMs. So every scan here is kept *narrow* — stats and histogram are
# computed one dim at a time over just that dim's real/gen columns (2 scans per
# dim, real+gen histograms merged into one). More scans than a combined pass, but
# each is cheap in memory, which is what actually matters at this file size.


def _pad_range(lo: float, hi: float) -> tuple[float, float]:
    """Widen (lo, hi) by ±0.5 when too tight to support `bins` distinct edges.

    Scalar duplicate of `_hist_edges`'s degenerate-range handling, applied per
    (group, dim).
    """
    if not (hi - lo > 1e-6 * max(abs(hi), 1.0)):
        return lo - 0.5, hi + 0.5
    return lo, hi


def _dim_stats(lf: pl.LazyFrame, j: int) -> pl.DataFrame:
    """One narrow streaming pass: per-group n + mean/std/range for dim `j`.

    Reads only this dim's real/gen columns (see the note above on why scans stay
    narrow). `n` (= `pl.len()`) is the group size; ddof=0 matches numpy's
    population std.
    """
    r, g = _raw_dim_expr("true", j), _raw_dim_expr("pred", j)
    return (
        lf.select("_group", r.alias("real"), g.alias("gen"))
        .group_by("_group")
        .agg(
            pl.len().alias("n"),
            pl.col("real").mean().alias("real_mean"),
            pl.col("gen").mean().alias("gen_mean"),
            pl.col("real").std(ddof=0).alias("real_std"),
            pl.col("gen").std(ddof=0).alias("gen_std"),
            pl.min_horizontal(pl.col("real").min(), pl.col("gen").min()).alias("lo"),
            pl.max_horizontal(pl.col("real").max(), pl.col("gen").max()).alias("hi"),
        )
        .collect(engine="streaming")
    )


def _dim_hist_counts(
    lf: pl.LazyFrame, j: int, lo_hi: pl.DataFrame, bins: int
) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
    """One streaming pass: real & gen histogram bin counts for dim `j`, by group.

    Projects to just `(_group, real, gen)` for this dim (an explicit up-front
    `.select` — projection pushdown doesn't reliably prune across the join in
    this polars version), joins each row to its group's padded `(lo, hi)` range,
    bins with the same floor/clip formula the original per-dim path used (so
    counts are bit-identical), and counts real and gen together by unpivoting only
    the two bin-index columns — a bounded 2× explosion, unlike unpivoting all 18
    raw values at once (which OOMs on a >RAM file). The inner join also drops rows
    whose group had n < 2.
    """
    narrow = lf.select(
        "_group",
        _raw_dim_expr("true", j).alias("real"),
        _raw_dim_expr("pred", j).alias("gen"),
    )
    width = pl.col("hi") - pl.col("lo")

    def _bin(col: str) -> pl.Expr:
        return (
            ((pl.col(col) - pl.col("lo")) / width * bins)
            .floor()
            .cast(pl.Int64)
            .clip(0, bins - 1)
        )

    hist = (
        narrow.join(lo_hi.lazy(), on="_group")
        .select("_group", _bin("real").alias("real"), _bin("gen").alias("gen"))
        .unpivot(index="_group", variable_name="kind", value_name="bin")
        .group_by(["_group", "kind", "bin"])
        .agg(pl.len().alias("count"))
        .collect(engine="streaming")
    )

    real_by: dict[str, np.ndarray] = {}
    gen_by: dict[str, np.ndarray] = {}
    for group, kind, bin_idx, count in hist.iter_rows():
        target = real_by if kind == "real" else gen_by
        target.setdefault(group, np.zeros(bins, dtype=np.int64))[bin_idx] = count
    return real_by, gen_by


# Per (group, dim) histogram payload: (real_counts, gen_counts, lo, hi).
_DimHists = dict[str, dict[str, tuple[np.ndarray, np.ndarray, float, float]]]


def _marginal_histograms(
    source: str | Path | pl.LazyFrame,
    group_by: str | None,
    n_energy_bins: int,
    bins: int,
) -> tuple[pl.DataFrame, _DimHists]:
    """Stratified per-dim marginal stats table + the per-(group, dim) histograms.

    The single source of truth for both `marginal_table_pl` (which wants only the
    table) and `plot_marginals` (which also wants the bin counts). Two narrow
    scans per dim (stats, then merged real+gen histogram) — see the note above
    these helpers on why scans stay narrow rather than combined.
    """
    lf = _add_group_label(_scan_predicted_local(source), group_by, n_energy_bins)

    empty = np.zeros(bins, dtype=np.int64)
    rows = []
    hists: _DimHists = {name: {} for name in RAW_TARGET_NAMES}
    for j, name in enumerate(RAW_TARGET_NAMES):
        stats = _dim_stats(lf, j).filter(pl.col("n") >= 2)
        range_by = {
            row["_group"]: _pad_range(row["lo"], row["hi"])
            for row in stats.iter_rows(named=True)
        }
        lo_hi = pl.DataFrame(
            [{"_group": g, "lo": lo, "hi": hi} for g, (lo, hi) in range_by.items()],
            schema={"_group": pl.Utf8, "lo": pl.Float64, "hi": pl.Float64},
        )
        real_by, gen_by = _dim_hist_counts(lf, j, lo_hi, bins)
        for row in stats.iter_rows(named=True):
            group = row["_group"]
            real_counts = real_by.get(group, empty)
            gen_counts = gen_by.get(group, empty)
            lo, hi = range_by[group]
            hists[name][group] = (real_counts, gen_counts, lo, hi)
            rows.append(
                {
                    "group": group,
                    "dim": name,
                    "n": row["n"],
                    "real_mean": row["real_mean"],
                    "gen_mean": row["gen_mean"],
                    "real_std": row["real_std"],
                    "gen_std": row["gen_std"],
                    "kl_real_gen": _kl_from_counts(real_counts, gen_counts),
                }
            )

    table = pl.DataFrame(rows).sort("kl_real_gen", descending=True)
    return table, hists


def marginal_table_pl(
    source: str | Path | pl.LazyFrame,
    group_by: str | None = None,
    n_energy_bins: int = 4,
    bins: int = 50,
) -> pl.DataFrame:
    """Per-dimension real-vs-generated summary stats + KL(real||gen), in raw units.

    `source` is a path to a `giant predict --coord local` parquet file, or an
    already-built LazyFrame with the same `pred_*`/`true_*`/`pdg`/`material`/
    `pre_E` columns (e.g. for testing). `group_by`: None for an aggregate table,
    or one of "pdg", "material", "energy" to stratify. Sorted worst-KL first, so
    failure modes hidden by the aggregate surface at the top.
    """
    table, _ = _marginal_histograms(source, group_by, n_energy_bins, bins)
    return table


def _limit_groups_by_kl_n(table: pd.DataFrame, max_groups: int) -> pd.DataFrame:
    """Keep only the `max_groups` groups with the largest max(kl) * n.

    Ranks by how badly wrong *and* how common a stratum is, rather than by KL
    alone, so a rare pdg/material with a noisy, high-variance KL estimate from a
    handful of samples doesn't crowd out groups that actually matter.
    """
    by_group = table.groupby("group").agg(
        kl_max=("kl_real_gen", "max"), n=("n", "first")
    )
    worst_first = (by_group["kl_max"] * by_group["n"]).sort_values(ascending=False)
    keep = set(worst_first.index[:max_groups])
    return table[table["group"].isin(keep)]


def _worst_first_groups(table: pd.DataFrame, max_groups: int) -> list[str]:
    """Group labels ranked by max(kl) * n (worst first), capped at `max_groups`."""
    by_group = table.groupby("group").agg(
        kl_max=("kl_real_gen", "max"), n=("n", "first")
    )
    worst_first = (by_group["kl_max"] * by_group["n"]).sort_values(ascending=False)
    return list(worst_first.index[:max_groups])


def plot_marginals(
    source: str | Path | pl.LazyFrame,
    dims: list[str] | None = None,
    group_by: str | None = None,
    n_energy_bins: int = 4,
    bins: int = 50,
    max_groups: int = 6,
    figsize_per_axis: tuple[float, float] = (3.5, 2.8),
):
    """Overlaid real-vs-generated histograms: one row per group, one column per dim.

    Without `group_by`, a single row over the whole file. With "pdg", "material",
    or "energy", one row per stratum, ranked by max(kl) * n (capped at
    `max_groups`) so groups that are both badly wrong and common surface first,
    rather than rare groups with a noisy, high-variance KL estimate. Every
    histogram is computed with streaming bin counts (`_marginal_histograms`) — no
    row is ever materialized.
    """
    dims = dims or RAW_TARGET_NAMES
    table, hists = _marginal_histograms(source, group_by, n_energy_bins, bins)

    if group_by is None:
        groups = ["all"]
    else:
        groups = _worst_first_groups(table.to_pandas(), max_groups)

    n_rows, n_cols = len(groups), len(dims)
    fig, axes = plt.subplots(
        n_rows,
        n_cols,
        squeeze=False,
        figsize=(figsize_per_axis[0] * n_cols, figsize_per_axis[1] * n_rows),
    )
    for row, label in enumerate(groups):
        for col, name in enumerate(dims):
            ax = axes[row][col]
            entry = hists[name].get(label)
            if entry is not None:
                real_counts, gen_counts, lo, hi = entry
                edges = np.linspace(lo, hi, bins + 1)
                ax.stairs(_step_density(real_counts, edges), edges, label="real")
                ax.stairs(_step_density(gen_counts, edges), edges, label="generated")
                ax.set_yscale("log")
            if row == 0:
                ax.set_title(name, fontsize=9)
            if col == 0:
                ax.set_ylabel(label, fontsize=8)
            if row == 0 and col == n_cols - 1:
                ax.legend(fontsize=7)
    fig.tight_layout()
    return fig


def _plot_kl_bars(table: pd.DataFrame, figsize: tuple[float, float]):
    """Shared bar-plot body for `plot_kl_bars_pl`.

    `table` is a `marginal_table_pl` result converted to pandas with
    `group`/`dim`/`kl_real_gen` columns.
    """
    pivot = table.pivot(index="dim", columns="group", values="kl_real_gen")
    pivot = pivot.reindex(RAW_TARGET_NAMES)
    groups = sorted(pivot.columns)
    pivot = pivot[groups]

    n_dims, n_groups = len(RAW_TARGET_NAMES), len(groups)
    x = np.arange(n_dims)
    width = 0.8 / n_groups

    fig, ax = plt.subplots(figsize=figsize)
    for i, group in enumerate(groups):
        offset = (i - (n_groups - 1) / 2) * width
        ax.bar(x + offset, pivot[group].to_numpy(), width=width, label=group)
    ax.set_xticks(x)
    ax.set_xticklabels(RAW_TARGET_NAMES, rotation=45, ha="right")
    ax.set_ylabel("KL(real || gen)")
    if n_groups > 1:
        ax.legend(fontsize=7)
    fig.tight_layout()
    return fig


def plot_kl_bars_pl(
    source: str | Path | pl.LazyFrame,
    group_by: str | None = None,
    n_energy_bins: int = 4,
    bins: int = 50,
    max_groups: int = 6,
    figsize: tuple[float, float] = (8, 4),
):
    """Bar plot of KL(real||gen) per target dimension, optionally stratified.

    One bar cluster per dimension; with `group_by` set, one bar per stratum
    within each cluster, so which dimension/stratum combination drives a KL
    regression is visible at a glance. With "pdg" or "material" — open-ended
    vocabularies — groups are capped at `max_groups`, ranked by max(kl) * n;
    "energy" is already bounded by `n_energy_bins` and isn't capped.
    """
    table = marginal_table_pl(
        source, group_by=group_by, n_energy_bins=n_energy_bins, bins=bins
    ).to_pandas()
    if group_by in ("pdg", "material"):
        table = _limit_groups_by_kl_n(table, max_groups)
    return _plot_kl_bars(table, figsize=figsize)


# ---------------------------------------------------------------------------
# Tier 2: joint structure
# ---------------------------------------------------------------------------


def correlation_matrices_pl(
    source: str | Path | pl.LazyFrame,
) -> tuple[np.ndarray, np.ndarray]:
    """9×9 Pearson correlation matrices of the raw targets, (real, generated).

    Computed from streaming sufficient statistics — one pass accumulating
    per-dim sums and all i≤j cross-products (in float64) — rather than loading
    the target arrays and calling `np.corrcoef`. `cov_ij = E[x_i x_j] -
    E[x_i]E[x_j]` (population/N normalization, matching `np.corrcoef`), then
    `corr = cov / sqrt(cov_ii cov_jj)`.
    """
    lf = _scan_predicted_local(source)
    n_dims = len(RAW_TARGET_NAMES)
    narrow = lf.select(
        *[
            _raw_dim_expr("true", j).cast(pl.Float64).alias(f"r{j}")
            for j in range(n_dims)
        ],
        *[
            _raw_dim_expr("pred", j).cast(pl.Float64).alias(f"g{j}")
            for j in range(n_dims)
        ],
    )

    aggs: list[pl.Expr] = [pl.len().alias("n")]
    for j in range(n_dims):
        aggs.append(pl.col(f"r{j}").sum().alias(f"rs{j}"))
        aggs.append(pl.col(f"g{j}").sum().alias(f"gs{j}"))
    for i in range(n_dims):
        for j in range(i, n_dims):
            aggs.append((pl.col(f"r{i}") * pl.col(f"r{j}")).sum().alias(f"rp{i}_{j}"))
            aggs.append((pl.col(f"g{i}") * pl.col(f"g{j}")).sum().alias(f"gp{i}_{j}"))
    row = narrow.select(aggs).collect(engine="streaming").row(0, named=True)
    n = row["n"]

    def _build(sum_prefix: str, prod_prefix: str) -> np.ndarray:
        mean = np.array([row[f"{sum_prefix}{j}"] for j in range(n_dims)]) / n
        cov = np.empty((n_dims, n_dims), dtype=np.float64)
        for i in range(n_dims):
            for j in range(i, n_dims):
                c = row[f"{prod_prefix}{i}_{j}"] / n - mean[i] * mean[j]
                cov[i, j] = cov[j, i] = c
        std = np.sqrt(np.diag(cov))
        return cov / np.outer(std, std)

    return _build("rs", "rp"), _build("gs", "gp")


def plot_correlation_matrices(source: str | Path | pl.LazyFrame):
    """Side-by-side real/generated correlation heatmaps, plus their difference."""
    real_corr, gen_corr = correlation_matrices_pl(source)
    diff = gen_corr - real_corr

    fig, axes = plt.subplots(1, 3, figsize=(13, 4))
    for ax, mat, title, cmap, vlim in [
        (axes[0], real_corr, "real", "coolwarm", (-1, 1)),
        (axes[1], gen_corr, "generated", "coolwarm", (-1, 1)),
        (axes[2], diff, "generated − real", "PuOr", (-0.5, 0.5)),
    ]:
        im = ax.imshow(mat, vmin=vlim[0], vmax=vlim[1], cmap=cmap)
        ax.set_xticks(range(len(RAW_TARGET_NAMES)))
        ax.set_xticklabels(RAW_TARGET_NAMES, rotation=90, fontsize=7)
        ax.set_yticks(range(len(RAW_TARGET_NAMES)))
        ax.set_yticklabels(RAW_TARGET_NAMES, fontsize=7)
        ax.set_title(title)
        fig.colorbar(im, ax=ax, fraction=0.046)
    fig.tight_layout()
    return fig


_DEFAULT_PAIRS = [
    ("step_length", "delta_e"),
    ("delta_e", "edep"),
    ("step_length", "edep"),
]


def plot_pairwise(
    source: str | Path | pl.LazyFrame,
    pairs: list[tuple[str, str]] | None = None,
    n_sample: int = 3000,
    seed: int = 0,
):
    """Real-vs-generated scatter for physically coupled target pairs.

    Marginals matching doesn't imply the joint does — these pairs are coupled by
    the underlying physics (energy loss tracks distance, edep is part of
    delta_e), so a model that decorrelates them shows up here even with clean
    per-dimension marginals. Only ~`n_sample` rows are ever materialized: a cheap
    streaming count fixes a hash-subsample fraction, and only the dims referenced
    by `pairs` are read for the kept rows.
    """
    pairs = pairs or _DEFAULT_PAIRS
    names = sorted({name for pair in pairs for name in pair})
    lf = _scan_predicted_local(source)

    n_total = lf.select(pl.len()).collect(engine="streaming").item()
    if n_total == 0:
        raise ValueError("no rows in predict source")
    frac = min(1.0, n_sample / n_total)
    threshold = int(frac * 2**32)

    select_exprs = []
    for name in names:
        j = RAW_TARGET_NAMES.index(name)
        select_exprs.append(_raw_dim_expr("true", j).alias(f"real_{name}"))
        select_exprs.append(_raw_dim_expr("pred", j).alias(f"gen_{name}"))

    sampled = (
        lf.select(*select_exprs)
        .with_row_index("_ri")
        .filter((pl.col("_ri").hash(seed=seed) % 2**32) < threshold)
        .drop("_ri")
        .collect(engine="streaming")
    )
    if sampled.height > n_sample:
        idx = np.random.default_rng(seed).choice(
            sampled.height, size=n_sample, replace=False
        )
        sampled = sampled[idx]

    fig, axes = plt.subplots(2, len(pairs), squeeze=False, figsize=(4 * len(pairs), 7))
    for col, (a, b) in enumerate(pairs):
        for row, (prefix, title) in enumerate([("real", "real"), ("gen", "generated")]):
            ax = axes[row][col]
            ax.scatter(
                sampled[f"{prefix}_{a}"].to_numpy(),
                sampled[f"{prefix}_{b}"].to_numpy(),
                s=3,
                alpha=0.3,
            )
            ax.set_xlabel(a)
            ax.set_ylabel(b)
            if col == 0:
                ax.set_title(title, loc="left", fontsize=9)
    fig.tight_layout()
    return fig


def _cos_alignment_expr(prefix: str) -> pl.Expr:
    """cos(angle) between `{prefix}` post_dir and travel_dir, per row."""
    post = [pl.col(f"{prefix}_{n}") for n in ("post_dx", "post_dy", "post_dz")]
    travel = [pl.col(f"{prefix}_{n}") for n in ("travel_dx", "travel_dy", "travel_dz")]
    dot = sum(p * t for p, t in zip(post, travel))
    post_norm = pl.sum_horizontal([p**2 for p in post]).sqrt()
    travel_norm = pl.sum_horizontal([t**2 for t in travel]).sqrt()
    return dot / (post_norm * travel_norm + 1e-8)


def plot_direction_alignment(source: str | Path | pl.LazyFrame, bins: int = 50):
    """cos(angle) between post_dir and travel_dir, real vs generated.

    These two unit vectors are coupled through the scattering physics, so their
    joint alignment is a check the per-dimension marginals can't see. Fixed
    `[-1, 1]` edges, one streaming bin-count pass per side — no materialization.
    """
    lf = _scan_predicted_local(source)
    real_counts, edges = _streaming_hist1d(
        lf, _cos_alignment_expr("true"), bins, lo=-1.0, hi=1.0
    )
    gen_counts, _ = _streaming_hist1d(
        lf, _cos_alignment_expr("pred"), bins, lo=-1.0, hi=1.0
    )

    fig, ax = plt.subplots(figsize=(5, 4))
    ax.stairs(_step_density(real_counts, edges), edges, label="real")
    ax.stairs(_step_density(gen_counts, edges), edges, label="generated")
    ax.set_yscale("log")
    ax.set_xlabel("cos(angle) between post_dir and travel_dir")
    ax.legend()
    fig.tight_layout()
    return fig


# ---------------------------------------------------------------------------
# Tier 3: physical constraints
# ---------------------------------------------------------------------------


def _pred_norm_expr(start: int) -> pl.Expr:
    """||pred direction|| for the 3 components starting at target dim `start`."""
    cols = [pl.col(f"pred_{LOCAL_TARGET_NAMES[k]}") for k in range(start, start + 3)]
    return pl.sum_horizontal([c**2 for c in cols]).sqrt()


def constraint_report_pl(
    source: str | Path | pl.LazyFrame, norm_tol: float = 0.05
) -> pl.DataFrame:
    """Rate of physical-constraint violations in the generated raw-space samples.

    The model is an unconstrained MLP, so nothing forces post_dir / travel_dir to
    stay unit-norm or step_length/delta_e/edep to stay non-negative — all hold by
    construction in the real data, so any violation rate here is purely a
    generation artifact. A single streaming aggregation over the `pred_*` columns;
    never materializes the target arrays.
    """
    lf = _scan_predicted_local(source)
    post_norm = _pred_norm_expr(3)
    travel_norm = _pred_norm_expr(6)
    raw_scalar_dims = [_raw_dim_expr("pred", j) for j in range(_N_SCALAR_DIMS)]

    agg = (
        lf.select(
            [
                ((post_norm - 1).abs() > norm_tol)
                .mean()
                .alias("post_dir_violation_rate"),
                (post_norm - 1).abs().mean().alias("post_dir_mean_abs_error"),
                ((travel_norm - 1).abs() > norm_tol)
                .mean()
                .alias("travel_dir_violation_rate"),
                (travel_norm - 1).abs().mean().alias("travel_dir_mean_abs_error"),
                *[
                    (raw < 0).mean().alias(f"{name}_violation_rate")
                    for raw, name in zip(
                        raw_scalar_dims, RAW_TARGET_NAMES[:_N_SCALAR_DIMS]
                    )
                ],
                *[
                    raw.clip(upper_bound=0).abs().mean().alias(f"{name}_mean_abs_error")
                    for raw, name in zip(
                        raw_scalar_dims, RAW_TARGET_NAMES[:_N_SCALAR_DIMS]
                    )
                ],
            ]
        )
        .collect(engine="streaming")
        .row(0, named=True)
    )

    rows = [
        {
            "check": "post_dir unit norm",
            "violation_rate": agg["post_dir_violation_rate"],
            "mean_abs_error": agg["post_dir_mean_abs_error"],
        },
        {
            "check": "travel_dir unit norm",
            "violation_rate": agg["travel_dir_violation_rate"],
            "mean_abs_error": agg["travel_dir_mean_abs_error"],
        },
    ]
    for name in RAW_TARGET_NAMES[:_N_SCALAR_DIMS]:
        rows.append(
            {
                "check": f"{name} >= 0",
                "violation_rate": agg[f"{name}_violation_rate"],
                "mean_abs_error": agg[f"{name}_mean_abs_error"],
            }
        )
    return pl.DataFrame(rows)


def plot_constraint_violations(source: str | Path | pl.LazyFrame, bins: int = 50):
    """Histograms backing `constraint_report_pl`: direction norms and sign of the scalars.

    Each panel is a streaming histogram (`_streaming_hist1d`) over the generated
    columns — the direction norms and the three physical scalar dims — so nothing
    is materialized. The dashed line marks the constraint boundary (1.0 for norms,
    0.0 for the non-negative scalars).
    """
    lf = _scan_predicted_local(source)

    panels: list[tuple[pl.Expr, str, float]] = [
        (_pred_norm_expr(3), "||post_dir||", 1.0),
        (_pred_norm_expr(6), "||travel_dir||", 1.0),
    ]
    for k, name in enumerate(RAW_TARGET_NAMES[:_N_SCALAR_DIMS]):
        panels.append((_raw_dim_expr("pred", k), f"generated {name}", 0.0))

    fig, axes = plt.subplots(1, len(panels), figsize=(4 * len(panels), 3.5))
    for ax, (expr, title, boundary) in zip(axes, panels):
        counts, edges = _streaming_hist1d(lf, expr, bins)
        ax.stairs(counts, edges)
        ax.set_yscale("log")
        ax.axvline(boundary, color="k", linestyle="--", linewidth=1)
        ax.set_title(title)
    fig.tight_layout()
    return fig


# ---------------------------------------------------------------------------
# Tier 4: event-level (shower) observables
#
# Built directly on `giant predict --coord local` parquet output, streamed in
# two passes rather than materialized: per-event sums (total deposited energy,
# etc.) run into the tens of millions of rows.
# ---------------------------------------------------------------------------

_PRE_COLS = ["pre_x", "pre_y", "pre_z", "pre_dx", "pre_dy", "pre_dz", "pre_E"]


@dataclass
class EventObservables:
    event_table: pl.DataFrame  # one row per event_id; real_*/gen_* columns, mm/MeV
    depth_edges: np.ndarray  # (depth_bins+1,) mm, along shower axis
    transverse_edges: np.ndarray  # (transverse_bins+1,) mm, perpendicular to axis
    real_depth_profile: np.ndarray  # (depth_bins,) mean edep/event/bin, MeV
    gen_depth_profile: np.ndarray
    real_depth_profile_std: np.ndarray  # event-to-event RMS per bin
    gen_depth_profile_std: np.ndarray
    real_transverse_profile: np.ndarray
    gen_transverse_profile: np.ndarray
    real_transverse_profile_std: np.ndarray
    gen_transverse_profile_std: np.ndarray


# Every per-row use of the per-event entry point / shower axis attaches those
# six values with `replace_strict` (a 10⁴-entry hash map applied as an
# expression) rather than a `.join`: polars 1.4's streaming join buffers the
# whole 323M-row left side and OOMs on a file larger than RAM, whereas
# `replace_strict` streams in bounded memory.


def _entry_axis_exprs(entry_df: pl.DataFrame) -> list[pl.Expr]:
    """`replace_strict` expressions mapping event_id → each entry_*/axis_* value."""
    event_ids = entry_df["event_id"].to_numpy()
    return [
        pl.col("event_id")
        .replace_strict(event_ids, entry_df[col].to_numpy(), return_dtype=pl.Float64)
        .alias(col)
        for col in ("entry_x", "entry_y", "entry_z", "axis_x", "axis_y", "axis_z")
    ]


def _entry_axis_and_bin_edges(
    lf: pl.LazyFrame,
    depth_bins: int,
    transverse_bins: int,
    sample_rows: int = 1_000_000,
    seed: int = 0,
) -> tuple[pl.DataFrame, np.ndarray, np.ndarray]:
    """Per-event shower axis (highest-pre_E row) plus depth/transverse bin edges.

    The per-event entry point / axis is a `group_by(event_id)` (a bounded
    streaming aggregation, no join). Bin edges are sized from `pre_pos` alone (no
    post_pos reconstruction needed) via robust quantiles of the depth/transverse
    proxies — estimated on a hash-subsample of ~`sample_rows` rows rather than the
    whole file, because an exact `quantile` is holistic (materializes every row
    and OOMs on a >RAM file) and outlier-clipping bin bounds don't need more than
    a sample.
    """
    entry_df = (
        lf.select(["event_id", *_PRE_COLS])
        .group_by("event_id")
        .agg(
            pl.col("pre_x").get(pl.col("pre_E").arg_max()).alias("entry_x"),
            pl.col("pre_y").get(pl.col("pre_E").arg_max()).alias("entry_y"),
            pl.col("pre_z").get(pl.col("pre_E").arg_max()).alias("entry_z"),
            pl.col("pre_dx").get(pl.col("pre_E").arg_max()).alias("axis_x"),
            pl.col("pre_dy").get(pl.col("pre_E").arg_max()).alias("axis_y"),
            pl.col("pre_dz").get(pl.col("pre_E").arg_max()).alias("axis_z"),
        )
        .collect(engine="streaming")
        .sort("event_id")
    )

    n_total = lf.select(pl.len()).collect(engine="streaming").item()
    frac = min(1.0, sample_rows / max(n_total, 1))
    threshold = int(frac * 2**32)

    dx = pl.col("pre_x") - pl.col("entry_x")
    dy = pl.col("pre_y") - pl.col("entry_y")
    dz = pl.col("pre_z") - pl.col("entry_z")
    depth = dx * pl.col("axis_x") + dy * pl.col("axis_y") + dz * pl.col("axis_z")
    tx = dx - depth * pl.col("axis_x")
    ty = dy - depth * pl.col("axis_y")
    tz = dz - depth * pl.col("axis_z")
    transverse = (tx**2 + ty**2 + tz**2).sqrt()

    sample = (
        lf.select(["event_id", "pre_x", "pre_y", "pre_z"])
        .filter((pl.col("pre_x").hash(seed=seed) % 2**32) < threshold)
        .with_columns(*_entry_axis_exprs(entry_df))
        .select(depth.alias("depth_proxy"), transverse.alias("transverse_proxy"))
        .collect(engine="streaming")
    )
    depth_proxy = sample["depth_proxy"].to_numpy()
    transverse_proxy = sample["transverse_proxy"].to_numpy()

    depth_lo = float(np.quantile(depth_proxy, 0.001))
    depth_hi = float(np.quantile(depth_proxy, 0.999))
    if not (depth_hi - depth_lo > 1e-6 * max(abs(depth_hi), 1.0)):
        depth_lo, depth_hi = depth_lo - 0.5, depth_hi + 0.5
    depth_edges = np.linspace(depth_lo, depth_hi, depth_bins + 1)
    transverse_hi = max(float(np.quantile(transverse_proxy, 0.999)), 1e-6)
    transverse_edges = np.linspace(0.0, transverse_hi, transverse_bins + 1)
    return entry_df, depth_edges, transverse_edges


# Approximate per-event medians are estimated from fixed log10-spaced value bins
# (median = the bin the running count crosses half at). Exact medians need a
# holistic per-group aggregation that can't stream and materializes every row
# (~18 GB / OOM on a >RAM file), whereas this histogram is a bounded streaming
# group_by. The range is generous enough to cover physical edep/step values.
_MED_LOG10_LO = -8.0
_MED_LOG10_HI = 6.0
_MED_BINS = 280
_MED_QUANTITIES = (
    "real_median_edep",
    "gen_median_edep",
    "real_median_length",
    "gen_median_length",
)


def _uniform_bin(value: pl.Expr, lo: float, hi: float, nbins: int) -> pl.Expr:
    """Bin `value` into `[0, nbins)` over uniform edges [lo, hi].

    Matches `np.digitize(value, linspace(lo, hi, nbins+1)[1:-1])` for uniform
    edges — floor of the scaled offset, clipped to the valid range.
    """
    return ((value - lo) / (hi - lo) * nbins).floor().cast(pl.Int64).clip(0, nbins - 1)


def _geometry_exprs(src: str, out: str) -> list[pl.Expr]:
    """Per-row edep / step_length / depth / transverse for one side (`true`/`pred`).

    Reconstructs world-frame post_pos as polars expressions — the inverse
    local-frame (Rodrigues) rotation of `giant.data.transforms`, closed-form for
    the constant ẑ axis — then projects `post_pos - entry` onto the per-event
    shower axis (`depth`) and its perpendicular (`transverse`). Requires the
    per-event `entry_*`/`axis_*` columns to be joined on already. Doing this in
    expressions (rather than a numpy loop over pyarrow batches) lets the whole
    per-row pass parallelize across the streaming engine's cores.
    """
    norm = (
        pl.col("pre_dx") ** 2 + pl.col("pre_dy") ** 2 + pl.col("pre_dz") ** 2
    ).sqrt()
    ux, uy, uz = (
        pl.col("pre_dx") / norm,
        pl.col("pre_dy") / norm,
        pl.col("pre_dz") / norm,
    )
    cos_t = uz.clip(-1.0, 1.0)
    sin_t = (1.0 - cos_t**2).clip(lower_bound=0.0).sqrt()

    # Rodrigues axis pre_dir × ẑ = [uy, -ux, 0], normalized; x̂ when pre_dir ∥ ẑ.
    axis_norm = (uy**2 + ux**2).sqrt()
    degenerate = axis_norm < 1e-7
    ax = pl.when(degenerate).then(pl.lit(1.0)).otherwise(uy / axis_norm)
    ay = pl.when(degenerate).then(pl.lit(0.0)).otherwise(-ux / axis_norm)

    vx, vy, vz = (
        pl.col(f"{src}_travel_dx"),
        pl.col(f"{src}_travel_dy"),
        pl.col(f"{src}_travel_dz"),
    )
    # axis × v (axis_z = 0); axis · v; inverse rotation applies R^T (−angle).
    kxv_x, kxv_y, kxv_z = ay * vz, -(ax * vz), ax * vy - ay * vx
    kdv = ax * vx + ay * vy
    omc = 1.0 - cos_t
    wx = vx * cos_t - kxv_x * sin_t + ax * kdv * omc
    wy = vy * cos_t - kxv_y * sin_t + ay * kdv * omc
    wz = vz * cos_t - kxv_z * sin_t  # axis_z * kdv * omc == 0

    step = pl.col(f"{src}_log_step_length").exp() - _LOG_EPS
    post_x = pl.col("pre_x") + step * wx
    post_y = pl.col("pre_y") + step * wy
    post_z = pl.col("pre_z") + step * wz

    dx = post_x - pl.col("entry_x")
    dy = post_y - pl.col("entry_y")
    dz = post_z - pl.col("entry_z")
    depth = dx * pl.col("axis_x") + dy * pl.col("axis_y") + dz * pl.col("axis_z")
    tx = dx - depth * pl.col("axis_x")
    ty = dy - depth * pl.col("axis_y")
    tz = dz - depth * pl.col("axis_z")
    transverse = (tx**2 + ty**2 + tz**2).sqrt()

    return [
        _edep_pl(src).alias(f"{out}_edep"),
        step.alias(f"{out}_step"),
        depth.alias(f"{out}_depth"),
        transverse.alias(f"{out}_transverse"),
    ]


def _spatial_grid(
    geo: pl.LazyFrame,
    out: str,
    depth_edges: np.ndarray,
    transverse_edges: np.ndarray,
    depth_bins: int,
    transverse_bins: int,
) -> pl.DataFrame:
    """One streaming pass: per-(event, depth-bin, transverse-bin) edep/length sums.

    The reduced grid (at most n_events × depth_bins × transverse_bins rows) carries
    everything the event-level observables need: summing its cells recovers exact
    per-event totals and the edep-weighted centroid/RMS numerators (since a sum of
    per-cell sums is the full per-event sum), and marginalizing one axis gives each
    longitudinal/transverse profile — so the whole 323M-row reduction happens in a
    parallel `group_by`, not a serial python loop.
    """
    depth_bin = _uniform_bin(
        pl.col(f"{out}_depth"), depth_edges[0], depth_edges[-1], depth_bins
    )
    transverse_bin = _uniform_bin(
        pl.col(f"{out}_transverse"),
        transverse_edges[0],
        transverse_edges[-1],
        transverse_bins,
    )
    edep = pl.col(f"{out}_edep")
    return (
        geo.with_columns(depth_bin.alias("db"), transverse_bin.alias("tb"))
        .group_by(["event_id", "db", "tb"])
        .agg(
            pl.len().alias("cnt"),
            edep.sum().alias("s_edep"),
            (edep * pl.col(f"{out}_depth")).sum().alias("s_edep_depth"),
            (edep * pl.col(f"{out}_transverse") ** 2).sum().alias("s_edep_t2"),
            pl.col(f"{out}_step").sum().alias("s_step"),
        )
        .collect(engine="streaming")
    )


def _reduce_spatial_grid(
    grid: pl.DataFrame,
    event_ids: np.ndarray,
    depth_bins: int,
    transverse_bins: int,
    depth_edges: np.ndarray,
) -> dict[str, np.ndarray]:
    """Collapse a `_spatial_grid` table into per-event arrays (aligned to event_ids).

    Pure numpy over the small reduced grid: `np.bincount` scatter-adds recover the
    per-event totals, centroid/RMS, dense (event, bin) profile matrices, and the
    shower-max depth — identical to the previous full-file numpy accumulation, but
    fed pre-summed grid cells instead of every row.
    """
    n_events = len(event_ids)
    idx = np.searchsorted(event_ids, grid["event_id"].to_numpy())
    db = grid["db"].to_numpy()
    tb = grid["tb"].to_numpy()
    cnt = grid["cnt"].to_numpy().astype(np.float64)
    s_edep = grid["s_edep"].to_numpy().astype(np.float64)
    s_edep_depth = grid["s_edep_depth"].to_numpy().astype(np.float64)
    s_edep_t2 = grid["s_edep_t2"].to_numpy().astype(np.float64)
    s_step = grid["s_step"].to_numpy().astype(np.float64)

    total_edep = np.bincount(idx, weights=s_edep, minlength=n_events)
    total_length = np.bincount(idx, weights=s_step, minlength=n_events)
    n_steps = np.bincount(idx, weights=cnt, minlength=n_events)
    sum_edep_depth = np.bincount(idx, weights=s_edep_depth, minlength=n_events)
    sum_edep_t2 = np.bincount(idx, weights=s_edep_t2, minlength=n_events)
    depth_mat = np.bincount(
        idx * depth_bins + db, weights=s_edep, minlength=n_events * depth_bins
    ).reshape(n_events, depth_bins)
    transverse_mat = np.bincount(
        idx * transverse_bins + tb, weights=s_edep, minlength=n_events * transverse_bins
    ).reshape(n_events, transverse_bins)

    safe_total = np.where(total_edep > 0, total_edep, 1.0)
    safe_n = np.where(n_steps > 0, n_steps, 1)
    depth_centers = 0.5 * (depth_edges[:-1] + depth_edges[1:])
    return {
        "total_edep": total_edep,
        "total_length": total_length,
        "n_steps": n_steps.astype(np.int64),
        "mean_edep": total_edep / safe_n,
        "mean_length": total_length / safe_n,
        "centroid_depth": sum_edep_depth / safe_total,
        "transverse_rms": np.sqrt(sum_edep_t2 / safe_total),
        "max_depth": depth_centers[np.argmax(depth_mat, axis=1)],
        "depth_mat": depth_mat,
        "transverse_mat": transverse_mat,
    }


def _approx_medians(geo: pl.LazyFrame, event_ids: np.ndarray) -> dict[str, np.ndarray]:
    """Per-event approximate median edep/step for real & gen, via a streaming pass.

    Histograms each quantity into fixed log10 bins per event (one streaming
    `group_by` after unpivoting only the four bin-index columns — a bounded 4×
    expansion), then reads off the bin its running count crosses half at. See the
    `_MED_*` note for why exact medians are avoided.
    """

    def _log_bin(col: str) -> pl.Expr:
        log10 = pl.col(col).clip(lower_bound=1e-30).log10()
        return (
            ((log10 - _MED_LOG10_LO) / (_MED_LOG10_HI - _MED_LOG10_LO) * _MED_BINS)
            .floor()
            .cast(pl.Int64)
            .clip(0, _MED_BINS - 1)
        )

    hist = (
        geo.select(
            "event_id",
            _log_bin("real_edep").alias("real_median_edep"),
            _log_bin("gen_edep").alias("gen_median_edep"),
            _log_bin("real_step").alias("real_median_length"),
            _log_bin("gen_step").alias("gen_median_length"),
        )
        .unpivot(index="event_id", variable_name="q", value_name="bin")
        .group_by(["q", "event_id", "bin"])
        .agg(pl.len().alias("c"))
        .collect(engine="streaming")
    )

    n_events = len(event_ids)
    log_width = (_MED_LOG10_HI - _MED_LOG10_LO) / _MED_BINS
    rows = np.arange(n_events)
    out: dict[str, np.ndarray] = {}
    for name in _MED_QUANTITIES:
        sub = hist.filter(pl.col("q") == name)
        idx = np.searchsorted(event_ids, sub["event_id"].to_numpy())
        counts = np.bincount(
            idx * _MED_BINS + sub["bin"].to_numpy(),
            weights=sub["c"].to_numpy().astype(np.float64),
            minlength=n_events * _MED_BINS,
        ).reshape(n_events, _MED_BINS)
        cum = np.cumsum(counts, axis=1)
        total = cum[:, -1]
        half = total / 2.0
        median_bin = (cum >= half[:, None]).argmax(axis=1)
        # Linear interpolation of the log-CDF within the median bin, for sub-bin
        # resolution (a bare bin center is only ~12% granular).
        cum_before = np.where(median_bin > 0, cum[rows, median_bin - 1], 0.0)
        count_in = counts[rows, median_bin]
        frac = np.where(count_in > 0, (half - cum_before) / count_in, 0.5)
        log_val = _MED_LOG10_LO + (median_bin + frac) * log_width
        median = 10.0**log_val
        median[total == 0] = 0.0
        out[name] = median
    return out


def compute_event_observables_pl(
    source: str | Path | pl.LazyFrame,
    depth_bins: int = 20,
    transverse_bins: int = 20,
) -> EventObservables:
    """Stream a `giant predict --coord local` parquet file into event-level observables.

    For each event, the highest-`pre_E` row is the primary's entry step
    (secondaries always carry less energy than their parent), fixing a shower
    axis/entry point shared by real and generated rows. Every row's `post_pos`/
    `edep` is reconstructed into the world frame in physical units (mm, MeV) and
    projected onto depth-along-axis / transverse-distance-from-axis.

    Everything is computed with parallel streaming polars aggregations — no serial
    python-over-pyarrow-batch loop, and no holistic per-event median (medians are
    approximated from a streaming log-bin histogram) — so the 323M-row reduction
    both uses all cores and stays within a bounded memory budget on a file larger
    than RAM. `event_table` carries `real_total_length`/`gen_total_length`
    (`sum(step_length)` per event) alongside the deposited-energy observables.
    """
    lf = _scan_predicted_local(source)
    entry_df, depth_edges, transverse_edges = _entry_axis_and_bin_edges(
        lf, depth_bins, transverse_bins
    )
    event_ids = entry_df["event_id"].to_numpy()

    # Project to just the geometry inputs *before* the join, and to just the
    # geometry outputs right after: projection pushdown doesn't reliably prune
    # across a join in this polars version, so without the explicit narrowing the
    # join buffers all ~30 columns × 323M rows and OOMs (the same trap the
    # marginal path documents). Narrowed, only ~9 columns cross the join.
    geo_inputs = [
        "event_id",
        "pre_x",
        "pre_y",
        "pre_z",
        "pre_dx",
        "pre_dy",
        "pre_dz",
        "pre_E",
        *[
            f"{prefix}_{name}"
            for prefix in ("true", "pred")
            for name in (
                "log_step_length",
                "edep_logit",
                "sec_logit",
                "travel_dx",
                "travel_dy",
                "travel_dz",
            )
        ],
    ]
    geo = (
        lf.select(geo_inputs)
        .join(entry_df.lazy(), on="event_id")
        .select(
            "event_id",
            *_geometry_exprs("true", "real"),
            *_geometry_exprs("pred", "gen"),
        )
    )

    real = _reduce_spatial_grid(
        _spatial_grid(
            geo, "real", depth_edges, transverse_edges, depth_bins, transverse_bins
        ),
        event_ids,
        depth_bins,
        transverse_bins,
        depth_edges,
    )
    gen = _reduce_spatial_grid(
        _spatial_grid(
            geo, "gen", depth_edges, transverse_edges, depth_bins, transverse_bins
        ),
        event_ids,
        depth_bins,
        transverse_bins,
        depth_edges,
    )
    medians = _approx_medians(geo, event_ids)

    event_table = pl.DataFrame(
        {
            "event_id": event_ids,
            "n_steps": real["n_steps"],
            "real_total_edep": real["total_edep"],
            "gen_total_edep": gen["total_edep"],
            "real_total_length": real["total_length"],
            "gen_total_length": gen["total_length"],
            "real_mean_edep": real["mean_edep"],
            "gen_mean_edep": gen["mean_edep"],
            "real_mean_length": real["mean_length"],
            "gen_mean_length": gen["mean_length"],
            "real_median_edep": medians["real_median_edep"],
            "gen_median_edep": medians["gen_median_edep"],
            "real_median_length": medians["real_median_length"],
            "gen_median_length": medians["gen_median_length"],
            "real_centroid_depth": real["centroid_depth"],
            "gen_centroid_depth": gen["centroid_depth"],
            "real_transverse_rms": real["transverse_rms"],
            "gen_transverse_rms": gen["transverse_rms"],
            "real_max_depth": real["max_depth"],
            "gen_max_depth": gen["max_depth"],
        }
    )

    return EventObservables(
        event_table=event_table,
        depth_edges=depth_edges,
        transverse_edges=transverse_edges,
        real_depth_profile=real["depth_mat"].mean(axis=0),
        gen_depth_profile=gen["depth_mat"].mean(axis=0),
        real_depth_profile_std=real["depth_mat"].std(axis=0),
        gen_depth_profile_std=gen["depth_mat"].std(axis=0),
        real_transverse_profile=real["transverse_mat"].mean(axis=0),
        gen_transverse_profile=gen["transverse_mat"].mean(axis=0),
        real_transverse_profile_std=real["transverse_mat"].std(axis=0),
        gen_transverse_profile_std=gen["transverse_mat"].std(axis=0),
    )


def plot_total_energy(observables: EventObservables, bins: int = 50):
    """Real-vs-generated histogram of total deposited energy per event, with resolution."""
    table = observables.event_table
    real = table["real_total_edep"].to_numpy()
    gen = table["gen_total_edep"].to_numpy()

    fig, ax = plt.subplots(figsize=(6, 4))
    edges = _hist_edges(real, gen, bins=bins).tolist()
    ax.hist(
        real,
        bins=edges,
        density=True,
        histtype="step",
        label=f"real (σ/μ={real.std() / real.mean():.3f})",
    )
    ax.hist(
        gen,
        bins=edges,
        density=True,
        histtype="step",
        label=f"generated (σ/μ={gen.std() / gen.mean():.3f})",
    )
    ax.set_yscale("log")
    ax.set_xlabel("total deposited energy per event [MeV]")
    ax.legend(fontsize=8)
    fig.tight_layout()
    return fig


def plot_total_length(observables: EventObservables, bins: int = 50):
    """Real-vs-generated histogram of total length traveled per event (sum of step_length)."""
    table = observables.event_table
    real = table["real_total_length"].to_numpy()
    gen = table["gen_total_length"].to_numpy()

    fig, ax = plt.subplots(figsize=(6, 4))
    edges = _hist_edges(real, gen, bins=bins).tolist()
    ax.hist(
        real,
        bins=edges,
        density=True,
        histtype="step",
        label=f"real (σ/μ={real.std() / real.mean():.3f})",
    )
    ax.hist(
        gen,
        bins=edges,
        density=True,
        histtype="step",
        label=f"generated (σ/μ={gen.std() / gen.mean():.3f})",
    )
    ax.set_yscale("log")
    ax.set_xlabel("total length traveled per event [mm]")
    ax.legend(fontsize=8)
    fig.tight_layout()
    return fig


def _plot_mean_median_per_step(
    real_mean: np.ndarray,
    gen_mean: np.ndarray,
    real_median: np.ndarray,
    gen_median: np.ndarray,
    mean_xlabel: str,
    median_xlabel: str,
    bins: int,
    median_bins: int,
):
    """Side-by-side real-vs-generated histograms: per-event mean (left), median (right).

    The mean is pulled down by a compressed/under-sampled right tail (rare large
    values), while the median is robust to that tail. Splitting them into separate
    panels keeps each comparison legible; if a tail-compression bias explains a
    low generated mean, the median panel should overlap far more closely. Each
    panel gets its own bin edges (`median_bins` larger, since the median is far
    less spread than the mean).
    """
    mean_edges = _hist_edges(real_mean, gen_mean, bins=bins)
    median_edges = _hist_edges(real_median, gen_median, bins=median_bins)
    prop_colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
    real_color, gen_color = prop_colors[0], prop_colors[1]

    fig, (ax_mean, ax_median) = plt.subplots(1, 2, figsize=(12, 4))
    for ax, real, gen, edges, title, xlabel in [
        (ax_mean, real_mean, gen_mean, mean_edges, "mean", mean_xlabel),
        (ax_median, real_median, gen_median, median_edges, "median", median_xlabel),
    ]:
        ax.hist(
            real,
            bins=edges,
            density=True,
            histtype="step",
            color=real_color,
            label=f"real (σ/μ={real.std() / real.mean():.3f})",
        )
        ax.hist(
            gen,
            bins=edges,
            density=True,
            histtype="step",
            color=gen_color,
            label=f"generated (σ/μ={gen.std() / gen.mean():.3f})",
        )
        ax.set_yscale("log")
        ax.set_title(title)
        ax.set_xlabel(xlabel)
        ax.legend(fontsize=8)
    fig.tight_layout()
    return fig


def plot_mean_energy_per_step(
    observables: EventObservables, bins: int = 50, median_bins: int = 150
):
    """Real-vs-generated histograms of mean/median deposited energy per step, per event.

    Per event: `total_edep / n_steps` (mean) and the per-step median edep —
    distinct from `plot_total_energy`, which histograms the per-event *total*;
    these instead ask whether the typical step's energy deposit is right,
    independent of how many steps the event happened to have.
    """
    table = observables.event_table
    return _plot_mean_median_per_step(
        table["real_mean_edep"].to_numpy(),
        table["gen_mean_edep"].to_numpy(),
        table["real_median_edep"].to_numpy(),
        table["gen_median_edep"].to_numpy(),
        mean_xlabel="mean deposited energy per step, per event [MeV]",
        median_xlabel="median deposited energy per step, per event [MeV]",
        bins=bins,
        median_bins=median_bins,
    )


def plot_mean_length_per_step(
    observables: EventObservables, bins: int = 50, median_bins: int = 150
):
    """Real-vs-generated histograms of mean/median step length per step, per event.

    Per event: `total_length / n_steps` (mean) and the per-step median step length
    — distinct from `plot_total_length`, which histograms the per-event *total*;
    these instead ask whether the typical step length is right, independent of how
    many steps the event happened to have.
    """
    table = observables.event_table
    return _plot_mean_median_per_step(
        table["real_mean_length"].to_numpy(),
        table["gen_mean_length"].to_numpy(),
        table["real_median_length"].to_numpy(),
        table["gen_median_length"].to_numpy(),
        mean_xlabel="mean step length per step, per event [mm]",
        median_xlabel="median step length per step, per event [mm]",
        bins=bins,
        median_bins=median_bins,
    )


def _plot_profile(
    centers: np.ndarray,
    real_mean: np.ndarray,
    gen_mean: np.ndarray,
    real_std: np.ndarray,
    gen_std: np.ndarray,
    xlabel: str,
):
    fig, ax = plt.subplots(figsize=(6, 4))
    ax.errorbar(centers, real_mean, yerr=real_std, fmt="o-", label="real", capsize=2)
    ax.errorbar(centers, gen_mean, yerr=gen_std, fmt="s-", label="generated", capsize=2)
    ax.set_xlabel(xlabel)
    ax.set_ylabel("mean edep per event per bin [MeV]")
    ax.legend(fontsize=8)
    fig.tight_layout()
    return fig


def plot_longitudinal_profile(observables: EventObservables):
    """E_dep(depth) mean ± event-to-event RMS, real vs generated."""
    centers = 0.5 * (observables.depth_edges[:-1] + observables.depth_edges[1:])
    return _plot_profile(
        centers,
        observables.real_depth_profile,
        observables.gen_depth_profile,
        observables.real_depth_profile_std,
        observables.gen_depth_profile_std,
        "depth along shower axis [mm]",
    )


def plot_transverse_profile(observables: EventObservables):
    """E_dep(transverse distance) mean ± event-to-event RMS, real vs generated (Molière-style)."""
    centers = 0.5 * (
        observables.transverse_edges[:-1] + observables.transverse_edges[1:]
    )
    return _plot_profile(
        centers,
        observables.real_transverse_profile,
        observables.gen_transverse_profile,
        observables.real_transverse_profile_std,
        observables.gen_transverse_profile_std,
        "transverse distance from shower axis [mm]",
    )


def plot_shower_max_depth(observables: EventObservables, bins: int = 30):
    """Real-vs-generated histogram of per-event shower-maximum depth."""
    table = observables.event_table
    real = table["real_max_depth"].to_numpy()
    gen = table["gen_max_depth"].to_numpy()

    fig, ax = plt.subplots(figsize=(6, 4))
    edges = _hist_edges(real, gen, bins=bins).tolist()
    ax.hist(real, bins=edges, density=True, histtype="step", label="real")
    ax.hist(gen, bins=edges, density=True, histtype="step", label="generated")
    ax.set_xlabel("depth of shower maximum [mm]")
    ax.legend(fontsize=8)
    fig.tight_layout()
    return fig


# ---------------------------------------------------------------------------
# Particle-species (pdg) contribution shares
#
# Dataset-wide (not per-event) breakdown of which pdg species contributed how
# much of the total deposited energy / total length traveled. Only needs scalar
# sums — no post_pos reconstruction, no shower axis — so it's a single lazy
# polars group_by.
# ---------------------------------------------------------------------------

_PDG_NAMES = {11: "e-", -11: "e+", 22: "gamma", 2112: "n", 2212: "p"}


def _pdg_label(pdg: int) -> str:
    if pdg in _PDG_NAMES:
        return _PDG_NAMES[pdg]
    if abs(pdg) > 1_000_000_000:
        return f"ion{pdg}"
    return str(pdg)


def pdg_contribution_table_pl(source: str | Path | pl.LazyFrame) -> pl.DataFrame:
    """Total edep / step_length contributed by each pdg species, real vs generated.

    One row per pdg code, sorted by pdg. Pure lazy polars `group_by` over the
    whole file — `edep`/`step_length` are scalars unaffected by the local-frame
    rotation. `step_length` is a simple `exp(...) - eps` de-log; `edep` is decoded
    from the deposit/secondary energy logits against pre_E (`_edep_pl`).
    """
    lf = _scan_predicted_local(source)

    def _delog(col: str) -> pl.Expr:
        return pl.col(col).exp() - _LOG_EPS

    return (
        lf.group_by("pdg")
        .agg(
            _edep_pl("true").sum().alias("real_total_edep"),
            _edep_pl("pred").sum().alias("gen_total_edep"),
            _delog("true_log_step_length").sum().alias("real_total_length"),
            _delog("pred_log_step_length").sum().alias("gen_total_length"),
        )
        .collect(engine="streaming")
        .sort("pdg")
    )


def _pdg_pie_shares(
    table: pl.DataFrame, real_col: str, gen_col: str, max_slices: int
) -> tuple[list[str], np.ndarray, np.ndarray]:
    """Pie-ready (labels, real_values, gen_values), lumping small contributors into 'other'.

    Ranked by combined real+gen contribution so the same species end up in the
    same slice position in both pies, making them easier to compare.
    """
    pdg = table["pdg"].to_numpy()
    real = table[real_col].to_numpy()
    gen = table[gen_col].to_numpy()

    order = np.argsort(-(real + gen))
    pdg, real, gen = pdg[order], real[order], gen[order]

    if len(pdg) > max_slices:
        keep = max_slices - 1
        labels = [_pdg_label(int(p)) for p in pdg[:keep]] + ["other"]
        real = np.append(real[:keep], real[keep:].sum())
        gen = np.append(gen[:keep], gen[keep:].sum())
    else:
        labels = [_pdg_label(int(p)) for p in pdg]

    return labels, real, gen


def _plot_pdg_pie(
    table: pl.DataFrame, real_col: str, gen_col: str, suptitle: str, max_slices: int
):
    labels, real_vals, gen_vals = _pdg_pie_shares(table, real_col, gen_col, max_slices)

    fig, axes = plt.subplots(1, 2, figsize=(9, 4.5))
    for ax, vals, title in [
        (axes[0], real_vals, "real"),
        (axes[1], gen_vals, "generated"),
    ]:
        ax.pie(vals, labels=labels, autopct="%1.1f%%", startangle=90)
        ax.set_title(title)
    fig.suptitle(suptitle)
    fig.tight_layout()
    return fig


def plot_pdg_energy_share(table: pl.DataFrame, max_slices: int = 6):
    """Real-vs-generated pies of total deposited energy share by pdg species."""
    return _plot_pdg_pie(
        table,
        "real_total_edep",
        "gen_total_edep",
        "deposited energy share by particle type",
        max_slices,
    )


def plot_pdg_length_share(table: pl.DataFrame, max_slices: int = 6):
    """Real-vs-generated pies of total length-traveled share by pdg species."""
    return _plot_pdg_pie(
        table,
        "real_total_length",
        "gen_total_length",
        "length traveled share by particle type",
        max_slices,
    )
