"""Notebook-friendly diagnostics for a trained model's sample quality.

All checks here read from `giant predict --coord local` parquet output
(`pred_*`/`true_*` columns, denormalized but still local-frame/log-scaled) —
there is no on-the-fly (checkpoint + live sampler) path; generate predictions
once via the CLI, then run every diagnostic below against that file::

    from giant.analysis import load_predicted_local
    from giant.analysis import plot_marginals, plot_kl_bars, plot_correlation_matrices
    from giant.analysis import plot_pairwise, plot_direction_alignment
    from giant.analysis import plot_constraint_violations

    samples = load_predicted_local("path/to/steps_predicted_local.parquet")

    plot_marginals(samples, group_by="energy")
    plot_kl_bars(samples, group_by="energy")
    plot_correlation_matrices(samples)
    plot_pairwise(samples)
    plot_direction_alignment(samples)
    plot_constraint_violations(samples)

(`--coord global` output isn't supported here — it has no ground-truth columns
to compare against.)

For event-level (shower) observables on the same `--coord local` predict
output, see `compute_event_observables_pl` — it streams the file directly
(no row subsampling, no `SampleCollection`), since per-event sums would be
corrupted by partial events::

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

    obs = compute_event_observables_pl("path/to/steps_predicted_local.parquet")
    plot_total_energy(obs)
    plot_total_length(obs)
    plot_mean_energy_per_step(obs)
    plot_mean_length_per_step(obs)
    plot_longitudinal_profile(obs)
    plot_transverse_profile(obs)
    plot_shower_max_depth(obs)

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

    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("path/to/steps_predicted_local.parquet")
    plot_pdg_energy_share(table)
    plot_pdg_length_share(table)

Four tiers of checks, building on the aggregate marginal/KL check in
`giant.validate.validate_marginals`:

1. stratified marginals  — per-dimension real-vs-generated comparison, sliced by
   pdg / material / energy so failures hidden by the aggregate don't go unnoticed.
2. joint structure        — correlation matrices, physically-coupled pairwise
   plots, and post/travel direction alignment, since marginals can match while
   the model decorrelates targets that are coupled by the underlying physics.
3. physical constraints   — unit-norm direction vectors and non-negative raw
   step_length/delta_e/edep, checked in denormalized physical units; nothing in
   the unconstrained MLP output enforces these, so violations are a pure
   generation artifact.
4. event-level observables — total deposited energy, total length traveled,
   longitudinal/transverse shower profiles, and shower-max depth, aggregated
   per `event_id` in the world frame with physical units (mm, MeV). 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.
"""

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,
)
from giant.data.transforms import (
    energy_simplex_decode,
    inv_log_transform,
    reconstruct_post_pos,
)
from giant.validate import _histogram_kl

# 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
)


def _decode_raw_targets(target_local: np.ndarray, pre_E: np.ndarray) -> np.ndarray:
    """Map a model-native target array (N, 9) to physical raw units.

    Column 0 (log_step_length) is de-logged; columns 1–2 (the deposit/secondary
    ALR energy logits) are decoded against `pre_E` into physical `delta_e` and
    `edep` (energy_simplex_decode); the six direction components pass through
    unchanged. Output columns therefore line up with `RAW_TARGET_NAMES`.
    """
    raw = target_local.astype(np.float32).copy()
    raw[:, 0] = inv_log_transform(target_local[:, 0])
    edep, _e_sec, _post_E, delta_e = energy_simplex_decode(target_local[:, 1:3], pre_E)
    raw[:, 1] = delta_e
    raw[:, 2] = edep
    return raw


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 _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.
    """
    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)


@dataclass
class SampleCollection:
    cond_cont_raw: np.ndarray  # (N, 9) denormalized conditioning (pre_E delogged)
    pdg: np.ndarray  # (N,) raw PDG codes
    material: np.ndarray  # (N,) raw material names
    real_raw: np.ndarray  # (N, 9) denormalized + delogged real targets
    gen_raw: np.ndarray  # (N, 9) denormalized + delogged generated targets


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' — "
            "load_predicted_local 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 load_predicted_local to match"
        )


_COND_CONT_COLS = [
    "pre_x",
    "pre_y",
    "pre_z",
    "pre_E",
    "pre_dx",
    "pre_dy",
    "pre_dz",
    "layer_id",
    "n_sec",
]


def load_predicted_local(
    path: str | Path,
    sample_frac: float = 1.0,
    seed: int = 0,
    batch_size: int = 1_000_000,
) -> SampleCollection:
    """Build a SampleCollection from a `giant predict --coord local` parquet file.

    Reads the `pred_*`/`true_*` columns directly — no checkpoint or model needed,
    since the predict CLI already denormalized them into this log-scaled,
    local-frame space. Requires the file to carry the `giant predict` metadata
    tag (see `_check_predict_metadata`); raises if it's missing or from
    --coord global, rather than guessing from column names.

    Reads via `_iter_predicted_local_batches` (pyarrow's row-batch reader)
    with the needed columns selected, so at most `batch_size` rows are ever
    materialized at once — a plain lazy-polars `.collect()` with the sampling
    filter applied afterward looks lazy but doesn't push the row reduction
    into the scan (see its `.explain()`), so it still peaks at the full
    file's memory footprint even for a small `sample_frac`; streaming keeps
    peak memory to one batch regardless of file size or `sample_frac`.

    `sample_frac` (0 < sample_frac <= 1) randomly keeps only that fraction of
    rows after the column projection — useful for files too large to
    comfortably hold as numpy arrays in `real_raw`/`gen_raw`. The kept/dropped
    decision is a hash of each row's position in the file (`seed`-dependent),
    computed against a running offset across batches so it's equivalent to
    hashing a single row index over the whole file rather than restarting at
    each batch boundary.
    """
    if not (0 < sample_frac <= 1):
        raise ValueError(f"sample_frac must be in (0, 1], got {sample_frac}")
    _check_predict_metadata(Path(path))
    pred_cols = [f"pred_{name}" for name in LOCAL_TARGET_NAMES]
    true_cols = [f"true_{name}" for name in LOCAL_TARGET_NAMES]
    columns = pred_cols + true_cols + _COND_CONT_COLS + ["pdg", "material"]

    threshold = int(sample_frac * 2**32) if sample_frac < 1.0 else None
    frames = []
    offset = 0
    for batch_df in _iter_predicted_local_batches(path, columns, batch_size):
        n = batch_df.height
        if threshold is not None:
            row_idx = pl.arange(offset, offset + n, eager=True).cast(pl.UInt32)
            batch_df = batch_df.filter((row_idx.hash(seed=seed) % 2**32) < threshold)
        offset += n
        frames.append(batch_df)
    df = pl.concat(frames) if len(frames) != 1 else frames[0]

    gen_log_local = df.select(pred_cols).to_numpy().astype(np.float32)
    real_log_local = df.select(true_cols).to_numpy().astype(np.float32)

    cond_cont_raw = df.select(_COND_CONT_COLS).to_numpy().astype(np.float32)
    pre_E = cond_cont_raw[:, 3]  # pre_E is already physical in predict output

    return SampleCollection(
        cond_cont_raw=cond_cont_raw,
        pdg=df["pdg"].to_numpy(),
        material=df["material"].to_numpy(),
        real_raw=_decode_raw_targets(real_log_local, pre_E),
        gen_raw=_decode_raw_targets(gen_log_local, pre_E),
    )


# ---------------------------------------------------------------------------
# Tier 1: stratified marginals
# ---------------------------------------------------------------------------


def _group_labels(
    collection: SampleCollection,
    group_by: str | None,
    n_energy_bins: int,
) -> list[tuple[str, np.ndarray]]:
    n = len(collection.pdg)
    if group_by is None:
        return [("all", np.ones(n, dtype=bool))]
    if group_by == "pdg":
        return [(f"pdg={v}", collection.pdg == v) for v in np.unique(collection.pdg)]
    if group_by == "material":
        return [
            (f"material={v}", collection.material == v)
            for v in np.unique(collection.material)
        ]
    if group_by == "energy":
        pre_E = collection.cond_cont_raw[:, 3]
        edges = np.quantile(pre_E, np.linspace(0, 1, n_energy_bins + 1))
        edges[-1] += 1e-6
        bin_idx = np.digitize(pre_E, edges[1:-1])
        return [
            (f"E∈[{edges[i]:.3g},{edges[i + 1]:.3g})", bin_idx == i)
            for i in range(n_energy_bins)
        ]
    raise ValueError(f"unknown group_by={group_by!r}")


def marginal_table(
    collection: SampleCollection,
    group_by: str | None = None,
    n_energy_bins: int = 4,
    bins: int = 50,
) -> pd.DataFrame:
    """Per-dimension real-vs-generated summary stats + KL(real||gen), in raw units.

    `group_by`: None for an aggregate table, or one of "pdg", "material", "energy"
    to stratify rows by that conditioning variable. Sorted worst-KL first, so
    failure modes hidden by the aggregate surface at the top.
    """
    rows = []
    for label, mask in _group_labels(collection, group_by, n_energy_bins):
        if mask.sum() < 2:
            continue
        real, gen = collection.real_raw[mask], collection.gen_raw[mask]
        for j, name in enumerate(RAW_TARGET_NAMES):
            rows.append(
                {
                    "group": label,
                    "dim": name,
                    "n": int(mask.sum()),
                    "real_mean": real[:, j].mean(),
                    "gen_mean": gen[:, j].mean(),
                    "real_std": real[:, j].std(),
                    "gen_std": gen[:, j].std(),
                    "kl_real_gen": _histogram_kl(real[:, j], gen[:, j], bins=bins),
                }
            )
    return (
        pd.DataFrame(rows)
        .sort_values("kl_real_gen", ascending=False)
        .reset_index(drop=True)
    )


# ---------------------------------------------------------------------------
# Polars-native duplicates
#
# `marginal_table`/`constraint_report` above require a `SampleCollection`
# with the full real/gen arrays already materialized in numpy. The functions
# below instead take a parquet path (or LazyFrame) directly and stay lazy,
# reading the file a small constant number of times — one native `group_by`
# (covering every group and every dim's mean/std/n/histogram-range at once)
# plus one more pass per dim for histogram bin counts — rather than filtering
# and re-collecting once per (group, dim) pair. The latter (the previous
# implementation) meant runtime scaled with the number of *groups*: on a
# 114M-row file with 138 distinct pdg codes, `group_by="pdg"` extrapolated to
# roughly an hour, against ~1 minute for `group_by=None`. This version's
# runtime is independent of group cardinality.
# ---------------------------------------------------------------------------


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.

    Same smoothing/normalization as `giant.validate._histogram_kl`, just
    taking counts directly instead of raw samples (the counts here come from
    a polars `group_by` aggregation, not `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 _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 _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 _add_group_label(
    lf: pl.LazyFrame, group_by: str | None, n_energy_bins: int
) -> pl.LazyFrame:
    """Add a `_group` string column matching `_group_labels`'s label format.

    For "pdg"/"material" the label is a direct string expr over the existing
    column — no upfront pass to enumerate distinct values needed, since
    `group_by("_group")` downstream discovers them itself. "energy" still
    needs one 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}")


def _dim_narrow_lf(lf: pl.LazyFrame, j: int) -> pl.LazyFrame:
    """Project down to just `(_group, real, gen)` for one dim, before any join/group_by.

    Polars' projection pushdown doesn't reliably prune columns across a
    `.join()` in this version — without this explicit `.select()` up front,
    `_dim_hist_counts`'s join ends up materializing every column of the file
    (~30 columns × 114M rows ≈ 13GB) instead of just the ~2-3 this dim needs,
    even though nothing downstream references the others. Selecting first
    guarantees the narrow projection regardless of what the optimizer would
    otherwise infer.
    """
    return lf.select(
        "_group",
        _raw_dim_expr("true", j).alias("real"),
        _raw_dim_expr("pred", j).alias("gen"),
    )


def _dim_stats(narrow: pl.LazyFrame) -> pl.DataFrame:
    """Per-group mean/std/n plus histogram-range (min/max of real+gen) for one dim.

    A single `group_by("_group")` pass over `narrow` (see `_dim_narrow_lf`),
    mirroring the proven-cheap shape of `pdg_contribution_table_pl`, rather
    than folding all 9 dims' source columns into one query — the latter reads
    ~19 columns and builds every dim's intermediate (softmax, log, etc.)
    arrays for the whole file at once, which is memory-heavy enough to OOM on
    a 114M-row file even though it's still just one pass.
    """
    return (
        narrow.group_by("_group")
        .agg(
            pl.len().alias("n"),
            pl.col("real").mean().alias("real_mean"),
            pl.col("gen").mean().alias("gen_mean"),
            # ddof=0 to match numpy's (population-std) default used by marginal_table
            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 _pad_degenerate_range(lo_hi: pl.DataFrame) -> pl.DataFrame:
    """Widen a (lo, hi) pair by ±0.5 when too tight to support `bins` distinct edges.

    Polars duplicate of `_hist_edges`'s degenerate-range handling, applied
    per group row instead of per call.
    """
    ok = (pl.col("hi") - pl.col("lo")) > 1e-6 * pl.max_horizontal(
        pl.col("hi").abs(), pl.lit(1.0)
    )
    return lo_hi.with_columns(
        pl.when(ok).then(pl.col("lo")).otherwise(pl.col("lo") - 0.5).alias("lo"),
        pl.when(ok).then(pl.col("hi")).otherwise(pl.col("hi") + 0.5).alias("hi"),
    )


def _dim_hist_counts(
    narrow: pl.LazyFrame, lo_hi: pl.DataFrame, bins: int
) -> tuple[pl.DataFrame, pl.DataFrame]:
    """Per-group real/gen histogram bin counts for one dim, as two long tables.

    Joins each row (from `narrow`, see `_dim_narrow_lf`) to its group's
    (lo, hi) range (from `lo_hi`, already collected and tiny — one row per
    group), bins real/gen into `[0, bins)`, and counts via a two-key
    `group_by(["_group", "bin"])` — a proper single hash-pass histogram.
    (An earlier version aggregated with one `(bin == k).sum()` expression per
    bin, i.e. `bins` separate boolean-compare-and-reduce passes over every
    row; that's O(N × bins) work — 50 bins meant ~50x more comparisons than
    necessary — and was the dominant cost, not the join.)
    """
    width = pl.col("hi") - pl.col("lo")
    real_bin = (
        ((pl.col("real") - pl.col("lo")) / width * bins).floor().cast(pl.Int64).clip(0, bins - 1)
    )
    gen_bin = (
        ((pl.col("gen") - pl.col("lo")) / width * bins).floor().cast(pl.Int64).clip(0, bins - 1)
    )

    joined = narrow.join(lo_hi.lazy(), on="_group")
    real_hist = (
        joined.select("_group", real_bin.alias("bin"))
        .group_by(["_group", "bin"])
        .agg(pl.len().alias("count"))
        .collect(engine="streaming")
    )
    gen_hist = (
        joined.select("_group", gen_bin.alias("bin"))
        .group_by(["_group", "bin"])
        .agg(pl.len().alias("count"))
        .collect(engine="streaming")
    )
    return real_hist, gen_hist


def _hist_counts_by_group(hist: pl.DataFrame, bins: int) -> dict[str, np.ndarray]:
    """Long `(_group, bin, count)` table -> `{group: dense (bins,) count array}`."""
    out: dict[str, np.ndarray] = {}
    for group, bin_idx, count in hist.iter_rows():
        out.setdefault(group, np.zeros(bins, dtype=np.int64))[bin_idx] = count
    return out


def marginal_table_pl(
    source: str | Path | pl.LazyFrame,
    group_by: str | None = None,
    n_energy_bins: int = 4,
    bins: int = 50,
) -> pl.DataFrame:
    """Polars duplicate of `marginal_table`, reading straight from a predict parquet.

    `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). Never builds a `SampleCollection` —
    see the module-level note above on why this stays lazy, and on why this
    reads the file three passes per dim (stats, real histogram, gen
    histogram) rather than once per (group, dim) pair.
    """
    lf = _add_group_label(_scan_predicted_local(source), group_by, n_energy_bins)

    rows = []
    for j, name in enumerate(RAW_TARGET_NAMES):
        narrow = _dim_narrow_lf(lf, j)
        stats = _dim_stats(narrow)
        stats = stats.filter(pl.col("n") >= 2)
        lo_hi = _pad_degenerate_range(stats.select("_group", "lo", "hi"))
        real_hist, gen_hist = _dim_hist_counts(narrow, lo_hi, bins)
        real_by_group = _hist_counts_by_group(real_hist, bins)
        gen_by_group = _hist_counts_by_group(gen_hist, bins)
        for row in stats.iter_rows(named=True):
            group = row["_group"]
            real_counts = real_by_group.get(group, np.zeros(bins, dtype=np.int64))
            gen_counts = gen_by_group.get(group, np.zeros(bins, dtype=np.int64))
            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),
                }
            )
    return pl.DataFrame(rows).sort("kl_real_gen", descending=True)


def plot_marginals(
    collection: SampleCollection,
    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 val set. 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
    in the dataset surface first, rather than rare groups with a noisy,
    high-variance KL estimate from just one or two samples.
    """
    dims = dims or RAW_TARGET_NAMES
    dim_idx = [RAW_TARGET_NAMES.index(d) for d in dims]
    groups = _group_labels(collection, group_by, n_energy_bins)

    if group_by is not None:
        table = marginal_table(
            collection, group_by=group_by, n_energy_bins=n_energy_bins, bins=bins
        )
        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)
        order = {label: rank for rank, label in enumerate(worst_first.index)}
        groups = [g for g in groups if g[0] in order]
        groups = sorted(groups, key=lambda g: order[g[0]])[: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, mask) in enumerate(groups):
        real, gen = collection.real_raw[mask], collection.gen_raw[mask]
        for col, j in enumerate(dim_idx):
            ax = axes[row][col]
            edges = _hist_edges(real[:, j], gen[:, j], bins=bins)
            ax.hist(real[:, j], bins=edges, density=True, histtype="step", label="real")
            ax.hist(
                gen[:, j], bins=edges, density=True, histtype="step", label="generated"
            )
            ax.set_yscale("log")
            if row == 0:
                ax.set_title(dims[col], 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 _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, like `plot_marginals`.

    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 _plot_kl_bars(table: pd.DataFrame, figsize: tuple[float, float]):
    """Shared bar-plot body for `plot_kl_bars`/`plot_kl_bars_pl`.

    `table` is a `marginal_table`/`marginal_table_pl` result (already converted
    to pandas in the polars case) 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(
    collection: SampleCollection,
    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 rather than buried in `marginal_table`'s
    sorted rows. With "pdg" or "material" — open-ended vocabularies that can
    run to many distinct values — groups are capped at `max_groups`, ranked by
    max(kl) * n as in `plot_marginals`; "energy" is already bounded by
    `n_energy_bins` and isn't capped. Built on top of `marginal_table`; see
    `plot_kl_bars_pl` for the polars-native, parquet-direct equivalent.
    """
    table = marginal_table(
        collection, group_by=group_by, n_energy_bins=n_energy_bins, bins=bins
    )
    if group_by in ("pdg", "material"):
        table = _limit_groups_by_kl_n(table, max_groups)
    return _plot_kl_bars(table, figsize=figsize)


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),
):
    """Polars duplicate of `plot_kl_bars`, reading straight from a predict parquet.

    Built on top of `marginal_table_pl`; see that function for the `source`
    argument and why it stays lazy until the final per-(group, dim) collect.
    See `plot_kl_bars` for the `max_groups` capping behavior.
    """
    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(collection: SampleCollection) -> tuple[np.ndarray, np.ndarray]:
    """Pearson correlation matrices of the raw targets, real vs generated."""
    return (
        np.corrcoef(collection.real_raw, rowvar=False),
        np.corrcoef(collection.gen_raw, rowvar=False),
    )


def plot_correlation_matrices(collection: SampleCollection):
    """Side-by-side real/generated correlation heatmaps, plus their difference."""
    real_corr, gen_corr = correlation_matrices(collection)
    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(
    collection: SampleCollection,
    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.
    """
    pairs = pairs or _DEFAULT_PAIRS
    rng = np.random.default_rng(seed)
    n = len(collection.pdg)
    idx = rng.choice(n, size=min(n_sample, n), replace=False)

    fig, axes = plt.subplots(2, len(pairs), squeeze=False, figsize=(4 * len(pairs), 7))
    for col, (a, b) in enumerate(pairs):
        ia, ib = RAW_TARGET_NAMES.index(a), RAW_TARGET_NAMES.index(b)
        for row, (data, title) in enumerate(
            [(collection.real_raw, "real"), (collection.gen_raw, "generated")]
        ):
            ax = axes[row][col]
            ax.scatter(data[idx, ia], data[idx, ib], 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 direction_alignment(collection: SampleCollection) -> tuple[np.ndarray, np.ndarray]:
    """cos(angle) between post_dir_local and travel_dir_local, 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.
    """

    def cos_angle(raw: np.ndarray) -> np.ndarray:
        post, travel = raw[:, 3:6], raw[:, 6:9]
        return np.sum(post * travel, axis=1) / (
            np.linalg.norm(post, axis=1) * np.linalg.norm(travel, axis=1) + 1e-8
        )

    return cos_angle(collection.real_raw), cos_angle(collection.gen_raw)


def plot_direction_alignment(collection: SampleCollection, bins: int = 50):
    real_cos, gen_cos = direction_alignment(collection)
    fig, ax = plt.subplots(figsize=(5, 4))
    edges = np.linspace(-1, 1, bins + 1)
    ax.hist(real_cos, bins=edges, density=True, histtype="step", label="real")
    ax.hist(gen_cos, bins=edges, density=True, histtype="step", 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 constraint_report(
    collection: SampleCollection, norm_tol: float = 0.05
) -> pd.DataFrame:
    """Rate of physical-constraint violations in the generated raw-space samples.

    The model is an unconstrained MLP, so nothing forces post_dir_local /
    travel_dir_local to stay unit-norm or step_length/delta_e/edep to stay
    non-negative — both hold by construction in the real data, so any
    violation rate here is purely a generation artifact.
    """
    gen = collection.gen_raw
    post_norm = np.linalg.norm(gen[:, 3:6], axis=1)
    travel_norm = np.linalg.norm(gen[:, 6:9], axis=1)

    rows = [
        {
            "check": "post_dir unit norm",
            "violation_rate": float(np.mean(np.abs(post_norm - 1) > norm_tol)),
            "mean_abs_error": float(np.mean(np.abs(post_norm - 1))),
        },
        {
            "check": "travel_dir unit norm",
            "violation_rate": float(np.mean(np.abs(travel_norm - 1) > norm_tol)),
            "mean_abs_error": float(np.mean(np.abs(travel_norm - 1))),
        },
    ]
    for j, name in enumerate(RAW_TARGET_NAMES[:_N_SCALAR_DIMS]):
        rows.append(
            {
                "check": f"{name} >= 0",
                "violation_rate": float(np.mean(gen[:, j] < 0)),
                "mean_abs_error": float(np.mean(np.clip(-gen[:, j], 0, None))),
            }
        )
    return pd.DataFrame(rows)


def constraint_report_pl(
    source: str | Path | pl.LazyFrame, norm_tol: float = 0.05
) -> pl.DataFrame:
    """Polars duplicate of `constraint_report`, reading straight from a predict parquet.

    Computed as a single set of lazy aggregation expressions over the
    `pred_*` columns — never materializes `gen_raw`. See `marginal_table_pl`
    for the `source` argument and why this stays lazy.
    """
    lf = _scan_predicted_local(source)
    pred_cols = [f"pred_{name}" for name in LOCAL_TARGET_NAMES]

    post_norm = sum(pl.col(pred_cols[k]) ** 2 for k in range(3, 6)).sqrt()
    travel_norm = sum(pl.col(pred_cols[k]) ** 2 for k in range(6, 9)).sqrt()
    raw_log_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_log_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_log_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(collection: SampleCollection):
    """Histograms backing `constraint_report`: direction norms and sign of the log-dims."""
    gen = collection.gen_raw
    post_norm = np.linalg.norm(gen[:, 3:6], axis=1)
    travel_norm = np.linalg.norm(gen[:, 6:9], axis=1)

    n_panels = 2 + _N_SCALAR_DIMS
    fig, axes = plt.subplots(1, n_panels, figsize=(4 * n_panels, 3.5))
    for ax, norm, title in [
        (axes[0], post_norm, "||post_dir||"),
        (axes[1], travel_norm, "||travel_dir||"),
    ]:
        ax.hist(norm, bins=_hist_edges(norm, bins=50), histtype="step")
        ax.set_yscale("log")
        ax.axvline(1.0, color="k", linestyle="--", linewidth=1)
        ax.set_title(title)
    for k, name in enumerate(RAW_TARGET_NAMES[:_N_SCALAR_DIMS]):
        ax = axes[2 + k]
        ax.hist(gen[:, k], bins=_hist_edges(gen[:, k], bins=50), histtype="step")
        ax.set_yscale("log")
        ax.axvline(0.0, color="k", linestyle="--", linewidth=1)
        ax.set_title(f"generated {name}")
    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 into a SampleCollection: per-event sums
# (total deposited energy, etc.) would be silently corrupted by the row
# subsampling `load_predicted_local(sample_frac=...)` uses to keep large files
# in memory, and the real files here 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


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

    Bin edges are sized from `pre_pos` alone (no post_pos reconstruction
    needed) — pre-step positions already trace the shower's extent closely
    enough to pick a sensible range, which avoids a second full streaming pass
    just to size the bins.
    """
    narrow = lf.select(["event_id", *_PRE_COLS])
    entry = narrow.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"),
    )

    joined = narrow.join(entry, on="event_id")
    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()

    stats = (
        joined.select(depth.alias("depth_proxy"), transverse.alias("transverse_proxy"))
        .select(
            pl.col("depth_proxy").quantile(0.001).alias("depth_lo"),
            pl.col("depth_proxy").quantile(0.999).alias("depth_hi"),
            pl.col("transverse_proxy").quantile(0.999).alias("transverse_hi"),
        )
        .collect(engine="streaming")
        .row(0, named=True)
    )

    entry_df = entry.collect(engine="streaming").sort("event_id")

    depth_lo, depth_hi = stats["depth_lo"], stats["depth_hi"]
    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(stats["transverse_hi"], 1e-6)
    transverse_edges = np.linspace(0.0, transverse_hi, transverse_bins + 1)
    return entry_df, depth_edges, transverse_edges


def _iter_predicted_local_batches(
    source: str | Path | pl.LazyFrame, columns: list[str], batch_size: int
):
    """Yield the needed columns in bounded-memory chunks.

    A path is streamed via pyarrow's row-batch reader so the file is never
    fully materialized; a `pl.LazyFrame` (the in-memory test-fixture case) is
    just collected once, since that data is small by construction.
    """
    if isinstance(source, pl.LazyFrame):
        yield source.select(columns).collect()
        return
    pf = pq.ParquetFile(Path(source))
    for batch in pf.iter_batches(batch_size=batch_size, columns=columns):
        yield pl.from_arrow(batch)


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

    For each event, the highest-`pre_E` row is taken as the primary's entry
    step (secondaries always carry less energy than their parent), fixing a
    shower axis/entry point shared by both real and generated rows (both are
    conditioned on the same real pre-step state). Every row's `post_pos`/`edep`
    is reconstructed into the world frame in physical units (mm, MeV) via
    `giant.data.transforms.reconstruct_post_pos`/`inv_log_transform` — the same
    functions `giant predict --coord global` uses — and projected onto
    depth-along-axis / transverse-distance-from-axis.

    `event_table` also carries `real_total_length`/`gen_total_length` —
    `sum(step_length)` per event, the total path length traveled by every
    track in the shower (not the same as the depth of any single point).

    Runs in two passes: a cheap pure-polars pass over `pre_*` columns only
    (shower axis + bin-edge sizing), then one streaming pass over the full
    file accumulating per-event and per-bin sums in numpy. Never materializes
    the file as a `SampleCollection` — see the module-level note above.
    """
    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()
    n_events = len(event_ids)
    entry_pos = (
        entry_df.select(["entry_x", "entry_y", "entry_z"]).to_numpy().astype(np.float32)
    )
    axis_dir = (
        entry_df.select(["axis_x", "axis_y", "axis_z"]).to_numpy().astype(np.float32)
    )

    n_steps = np.zeros(n_events, dtype=np.int64)
    real_total_edep = np.zeros(n_events, dtype=np.float64)
    gen_total_edep = np.zeros(n_events, dtype=np.float64)
    real_total_length = np.zeros(n_events, dtype=np.float64)
    gen_total_length = np.zeros(n_events, dtype=np.float64)
    real_sum_edep_depth = np.zeros(n_events, dtype=np.float64)
    gen_sum_edep_depth = np.zeros(n_events, dtype=np.float64)
    real_sum_edep_transverse2 = np.zeros(n_events, dtype=np.float64)
    gen_sum_edep_transverse2 = np.zeros(n_events, dtype=np.float64)
    real_depth_bin_edep = np.zeros((n_events, depth_bins), dtype=np.float64)
    gen_depth_bin_edep = np.zeros((n_events, depth_bins), dtype=np.float64)
    real_transverse_bin_edep = np.zeros((n_events, transverse_bins), dtype=np.float64)
    gen_transverse_bin_edep = np.zeros((n_events, transverse_bins), dtype=np.float64)

    pred_cols = [f"pred_{name}" for name in LOCAL_TARGET_NAMES]
    true_cols = [f"true_{name}" for name in LOCAL_TARGET_NAMES]
    needed_cols = [
        "event_id",
        "pre_x",
        "pre_y",
        "pre_z",
        "pre_E",  # needed to decode the energy simplex into physical edep
        "pre_dx",
        "pre_dy",
        "pre_dz",
        *pred_cols,
        *true_cols,
    ]

    for batch_df in _iter_predicted_local_batches(source, needed_cols, batch_size):
        eid = batch_df["event_id"].to_numpy()
        idx = np.searchsorted(event_ids, eid)

        pre_pos = (
            batch_df.select(["pre_x", "pre_y", "pre_z"]).to_numpy().astype(np.float32)
        )
        pre_dir = (
            batch_df.select(["pre_dx", "pre_dy", "pre_dz"])
            .to_numpy()
            .astype(np.float32)
        )

        pre_E = batch_df["pre_E"].to_numpy().astype(np.float32)

        def _reconstruct(cols: list[str]) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
            raw = batch_df.select(cols).to_numpy().astype(np.float32)
            step_length = inv_log_transform(raw[:, 0])
            edep, _e_sec, _post_E, _delta_e = energy_simplex_decode(raw[:, 1:3], pre_E)
            travel_dir_local = raw[:, 6:9]
            post_pos = reconstruct_post_pos(
                pre_pos, pre_dir, step_length, travel_dir_local
            )
            return post_pos, edep, step_length

        real_post_pos, real_edep, real_step_length = _reconstruct(true_cols)
        gen_post_pos, gen_edep, gen_step_length = _reconstruct(pred_cols)

        e_pos = entry_pos[idx]
        a_dir = axis_dir[idx]

        def _depth_transverse(post_pos: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
            disp = post_pos - e_pos
            depth = np.sum(disp * a_dir, axis=1)
            perp = disp - depth[:, None] * a_dir
            transverse = np.linalg.norm(perp, axis=1)
            return depth, transverse

        real_depth, real_transverse = _depth_transverse(real_post_pos)
        gen_depth, gen_transverse = _depth_transverse(gen_post_pos)

        real_depth_bin = np.digitize(real_depth, depth_edges[1:-1])
        gen_depth_bin = np.digitize(gen_depth, depth_edges[1:-1])
        real_transverse_bin = np.digitize(real_transverse, transverse_edges[1:-1])
        gen_transverse_bin = np.digitize(gen_transverse, transverse_edges[1:-1])

        # np.add.at is an unbuffered ufunc method — a well-known numpy slow path
        # for scatter-add (44% of this function's runtime on a 114M-row profile).
        # np.bincount does the same accumulation with a single optimized pass;
        # the 2D (per-event, per-bin) accumulators flatten (idx, bin) into one
        # bincount index and reshape back, since bincount only scatters into 1D.
        n_steps += np.bincount(idx, minlength=n_events)
        real_total_edep += np.bincount(idx, weights=real_edep, minlength=n_events)
        gen_total_edep += np.bincount(idx, weights=gen_edep, minlength=n_events)
        real_total_length += np.bincount(idx, weights=real_step_length, minlength=n_events)
        gen_total_length += np.bincount(idx, weights=gen_step_length, minlength=n_events)
        real_sum_edep_depth += np.bincount(
            idx, weights=real_edep * real_depth, minlength=n_events
        )
        gen_sum_edep_depth += np.bincount(
            idx, weights=gen_edep * gen_depth, minlength=n_events
        )
        real_sum_edep_transverse2 += np.bincount(
            idx, weights=real_edep * real_transverse**2, minlength=n_events
        )
        gen_sum_edep_transverse2 += np.bincount(
            idx, weights=gen_edep * gen_transverse**2, minlength=n_events
        )
        real_depth_bin_edep += np.bincount(
            idx * depth_bins + real_depth_bin,
            weights=real_edep,
            minlength=n_events * depth_bins,
        ).reshape(n_events, depth_bins)
        gen_depth_bin_edep += np.bincount(
            idx * depth_bins + gen_depth_bin,
            weights=gen_edep,
            minlength=n_events * depth_bins,
        ).reshape(n_events, depth_bins)
        real_transverse_bin_edep += np.bincount(
            idx * transverse_bins + real_transverse_bin,
            weights=real_edep,
            minlength=n_events * transverse_bins,
        ).reshape(n_events, transverse_bins)
        gen_transverse_bin_edep += np.bincount(
            idx * transverse_bins + gen_transverse_bin,
            weights=gen_edep,
            minlength=n_events * transverse_bins,
        ).reshape(n_events, transverse_bins)

    safe_real_total = np.where(real_total_edep > 0, real_total_edep, 1.0)
    safe_gen_total = np.where(gen_total_edep > 0, gen_total_edep, 1.0)
    real_centroid_depth = real_sum_edep_depth / safe_real_total
    gen_centroid_depth = gen_sum_edep_depth / safe_gen_total
    real_transverse_rms = np.sqrt(real_sum_edep_transverse2 / safe_real_total)
    gen_transverse_rms = np.sqrt(gen_sum_edep_transverse2 / safe_gen_total)

    depth_centers = 0.5 * (depth_edges[:-1] + depth_edges[1:])
    real_max_depth = depth_centers[np.argmax(real_depth_bin_edep, axis=1)]
    gen_max_depth = depth_centers[np.argmax(gen_depth_bin_edep, axis=1)]

    safe_n_steps = np.where(n_steps > 0, n_steps, 1)
    real_mean_edep = real_total_edep / safe_n_steps
    gen_mean_edep = gen_total_edep / safe_n_steps
    real_mean_length = real_total_length / safe_n_steps
    gen_mean_length = gen_total_length / safe_n_steps

    # The inverse transform must be applied *before* the median, not after:
    # median commutes with inv_log_transform only for odd-length groups. For an
    # even number of steps polars' .median() averages the two central order
    # statistics, and that averaging does not commute with the nonlinear exp
    # (it would yield their geometric mean rather than the true median). So take
    # the median on the raw (exp-transformed) values directly. No post_pos
    # reconstruction is needed — only the magnitudes matter.
    medians = (
        lf.select(
            "event_id",
            "true_edep_logit",
            "true_sec_logit",
            "pred_edep_logit",
            "pred_sec_logit",
            "pre_E",
            "true_log_step_length",
            "pred_log_step_length",
        )
        .group_by("event_id")
        .agg(
            _edep_pl("true").median().alias("real_median_edep"),
            _edep_pl("pred").median().alias("gen_median_edep"),
            (pl.col("true_log_step_length").exp() - _LOG_EPS)
            .median()
            .alias("real_median_length"),
            (pl.col("pred_log_step_length").exp() - _LOG_EPS)
            .median()
            .alias("gen_median_length"),
        )
        .sort("event_id")
        .collect(engine="streaming")
    )
    real_median_edep = medians["real_median_edep"].to_numpy()
    gen_median_edep = medians["gen_median_edep"].to_numpy()
    real_median_length = medians["real_median_length"].to_numpy()
    gen_median_length = medians["gen_median_length"].to_numpy()

    event_table = pl.DataFrame(
        {
            "event_id": event_ids,
            "n_steps": 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": real_median_edep,
            "gen_median_edep": gen_median_edep,
            "real_median_length": real_median_length,
            "gen_median_length": 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_bin_edep.mean(axis=0),
        gen_depth_profile=gen_depth_bin_edep.mean(axis=0),
        real_depth_profile_std=real_depth_bin_edep.std(axis=0),
        gen_depth_profile_std=gen_depth_bin_edep.std(axis=0),
        real_transverse_profile=real_transverse_bin_edep.mean(axis=0),
        gen_transverse_profile=gen_transverse_bin_edep.mean(axis=0),
        real_transverse_profile_std=real_transverse_bin_edep.std(axis=0),
        gen_transverse_profile_std=gen_transverse_bin_edep.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)
    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)
    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 (rather than overlaying on one axis) keeps each comparison
    legible; if a tail-compression bias is the explanation for a low generated
    mean, the median panel should overlap far more closely than the mean panel.

    Each panel gets its own bin edges (`bins` for the mean, `median_bins` for
    the median): the median is far less spread than the mean, so reusing the
    mean's range would crush it into a few bins — `median_bins` is larger to
    resolve that tighter distribution.
    """
    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. Mean and median
    are drawn in separate panels, each with its own binning (`bins` /
    `median_bins`); see `_plot_mean_median_per_step`.
    """
    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. Mean and median
    are drawn in separate panels, each with its own binning (`bins` /
    `median_bins`); see `_plot_mean_median_per_step`.
    """
    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)
    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. Unlike the
# event-level checks above, this 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,
    )
