import matplotlib

matplotlib.use("Agg")  # no display needed for plot smoke tests

import numpy as np
import pandas as pd
import polars as pl
import pyarrow as pa
import pyarrow.parquet as pq
import pytest

from giant.analysis import (
    RAW_TARGET_NAMES,
    RolloutVsTruth,
    compute_event_observables_pl,
    compute_rollout_vs_truth_observables_pl,
    constraint_report_pl,
    correlation_matrices_pl,
    marginal_table_pl,
    pdg_contribution_table_pl,
    plot_constraint_violations,
    plot_correlation_matrices,
    plot_direction_alignment,
    plot_kl_bars_pl,
    plot_longitudinal_profile,
    plot_marginals,
    plot_mean_energy_per_step,
    plot_mean_length_per_step,
    plot_pairwise,
    plot_pdg_energy_share,
    plot_pdg_length_share,
    plot_shower_max_depth,
    plot_total_energy,
    plot_total_length,
    plot_transverse_profile,
)
from giant.constants import (
    LOCAL_TARGET_NAMES,
    PREDICT_COORD_METADATA_KEY,
    PREDICT_SCHEMA_VERSION,
    PREDICT_SCHEMA_VERSION_KEY,
    ROLLOUT_COORD_VALUE,
    TERM_ENERGY_CUTOFF,
    TERM_NATURAL_END,
)
from giant.data.transforms import (
    energy_simplex_decode,
    inv_log_transform,
    local_frame_rotation,
    log_transform,
    reconstruct_post_pos,
    travel_direction,
)

# ---------------------------------------------------------------------------
# Fixtures: predict `--coord local` parquet writers
# ---------------------------------------------------------------------------


def _unit_vectors(rng, n):
    v = rng.standard_normal((n, 3)).astype(np.float32)
    return v / np.linalg.norm(v, axis=1, keepdims=True)


def _write_predicted_local_parquet(path, n=50, metadata=None, rng=None):
    """Mimic `giant predict --coord local`'s output schema for the loader tests.

    Column 0 is a log-scaled step_length; columns 1–2 are the deposit/secondary
    ALR energy logits (unconstrained reals, decoded against pre_E); columns 3–8
    are direction components.
    """
    rng = rng or np.random.default_rng(0)
    true_log_local = rng.standard_normal((n, 9)).astype(np.float32)
    true_log_local[:, 0] = log_transform(rng.uniform(0.1, 5.0, n).astype(np.float32))
    pred_log_local = true_log_local + rng.normal(0, 0.01, (n, 9)).astype(np.float32)
    pre_E = rng.uniform(1.0, 100.0, n).astype(np.float32)

    table = pa.table(
        {
            "event_id": rng.integers(0, 10, n),
            "pdg": rng.choice([11, -11, 22], n),
            "pre_x": rng.standard_normal(n).astype(np.float32),
            "pre_y": rng.standard_normal(n).astype(np.float32),
            "pre_z": rng.standard_normal(n).astype(np.float32),
            "pre_E": pre_E,
            "pre_dx": rng.standard_normal(n).astype(np.float32),
            "pre_dy": rng.standard_normal(n).astype(np.float32),
            "pre_dz": rng.standard_normal(n).astype(np.float32),
            "material": rng.choice(["W", "Pb"], n),
            "layer_id": rng.integers(0, 10, n).astype(np.int32),
            "n_sec": rng.integers(0, 3, n).astype(np.int32),
            **{
                f"pred_{name}": pred_log_local[:, j]
                for j, name in enumerate(LOCAL_TARGET_NAMES)
            },
            **{
                f"true_{name}": true_log_local[:, j]
                for j, name in enumerate(LOCAL_TARGET_NAMES)
            },
        }
    )
    if metadata is not None:
        table = table.replace_schema_metadata(metadata)
    pq.write_table(table, path)
    return true_log_local, pred_log_local, pre_E


def _predicted_local_path(tmp_path, n=200, seed=0):
    path = tmp_path / "predicted_local.parquet"
    _write_predicted_local_parquet(
        path,
        n=n,
        rng=np.random.default_rng(seed),
        metadata={
            PREDICT_COORD_METADATA_KEY: "local",
            PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
        },
    )
    return path


# ---------------------------------------------------------------------------
# Inline numpy oracle: the ground truth the streaming functions are checked
# against. This is the raw-space decoding the old `SampleCollection` path did,
# recomputed directly from the small fixture rather than in the module.
# ---------------------------------------------------------------------------


def _raw_from_parquet(path):
    """(real_raw, gen_raw, pdg, material, pre_E) in physical units, via numpy."""
    df = pl.read_parquet(path)
    pre_E = df["pre_E"].to_numpy().astype(np.float32)

    def decode(prefix):
        log_local = np.column_stack(
            [df[f"{prefix}_{name}"].to_numpy() for name in LOCAL_TARGET_NAMES]
        ).astype(np.float32)
        raw = log_local.copy()
        raw[:, 0] = inv_log_transform(log_local[:, 0])
        edep, _e_sec, _post_E, delta_e = energy_simplex_decode(log_local[:, 1:3], pre_E)
        raw[:, 1] = delta_e
        raw[:, 2] = edep
        return raw

    return (
        decode("true"),
        decode("pred"),
        df["pdg"].to_numpy(),
        df["material"].to_numpy(),
        pre_E,
    )


def _histogram_kl(p_samples, q_samples, bins=50, eps=1e-8):
    """KL(P || Q) between two 1D samples via a shared histogram (numpy oracle)."""
    lo = min(p_samples.min(), q_samples.min())
    hi = max(p_samples.max(), q_samples.max())
    if hi <= lo:
        return 0.0
    edges = np.linspace(lo, hi, bins + 1)
    p_hist, _ = np.histogram(p_samples, bins=edges)
    q_hist, _ = np.histogram(q_samples, bins=edges)
    p = p_hist.astype(np.float64) + eps
    q = q_hist.astype(np.float64) + eps
    p /= p.sum()
    q /= q.sum()
    return float(np.sum(p * np.log(p / q)))


def _group_masks(pdg, material, pre_E, group_by, n_energy_bins=4):
    """Replicate the module's `_group` labelling so oracle labels line up."""
    n = len(pdg)
    if group_by is None:
        return [("all", np.ones(n, dtype=bool))]
    if group_by == "pdg":
        return [(f"pdg={int(v)}", pdg == v) for v in np.unique(pdg)]
    if group_by == "material":
        return [(f"material={v}", material == v) for v in np.unique(material)]
    if group_by == "energy":
        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(group_by)


def _numpy_marginal_table(real, gen, pdg, material, pre_E, group_by, bins=50):
    rows = []
    for label, mask in _group_masks(pdg, material, pre_E, group_by):
        if mask.sum() < 2:
            continue
        r, g = real[mask], gen[mask]
        for j, name in enumerate(RAW_TARGET_NAMES):
            rows.append(
                {
                    "group": label,
                    "dim": name,
                    "n": int(mask.sum()),
                    "real_mean": r[:, j].mean(),
                    "gen_mean": g[:, j].mean(),
                    "real_std": r[:, j].std(),
                    "gen_std": g[:, j].std(),
                    "kl_real_gen": _histogram_kl(r[:, j], g[:, j], bins=bins),
                }
            )
    return pd.DataFrame(rows).sort_values(["group", "dim"]).reset_index(drop=True)


def _numpy_constraint_report(gen, norm_tol=0.05):
    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[:3]):
        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)


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


@pytest.mark.parametrize("group_by", [None, "pdg", "material", "energy"])
def test_marginal_table_pl_matches_numpy_oracle(tmp_path, group_by):
    path = _predicted_local_path(tmp_path)
    real, gen, pdg, material, pre_E = _raw_from_parquet(path)

    expected = _numpy_marginal_table(real, gen, pdg, material, pre_E, group_by)
    actual = (
        marginal_table_pl(path, group_by=group_by).sort(["group", "dim"]).to_pandas()
    )

    assert list(expected["group"]) == list(actual["group"])
    assert list(expected["n"]) == list(actual["n"])
    for col in ["real_mean", "gen_mean", "real_std", "gen_std"]:
        np.testing.assert_allclose(
            expected[col].to_numpy(), actual[col].to_numpy(), atol=1e-4, rtol=1e-4
        )
    # KL uses np.histogram (oracle) vs polars binning (lazy path); the two bin the
    # boundary (min/max) sample differently, so allow a small absolute discrepancy
    # rather than requiring bit-identical estimates.
    np.testing.assert_allclose(
        expected["kl_real_gen"].to_numpy(), actual["kl_real_gen"].to_numpy(), atol=2e-2
    )


def test_marginal_table_pl_aggregate_has_all_dims(tmp_path):
    table = marginal_table_pl(_predicted_local_path(tmp_path))
    assert set(table["dim"].to_list()) == set(RAW_TARGET_NAMES)
    assert (table["group"] == "all").all()


def test_marginal_table_pl_accepts_lazyframe(tmp_path):
    path = _predicted_local_path(tmp_path)
    from_path = marginal_table_pl(path).sort(["group", "dim"])
    from_lf = marginal_table_pl(pl.scan_parquet(path)).sort(["group", "dim"])
    np.testing.assert_allclose(
        from_lf["kl_real_gen"].to_numpy(), from_path["kl_real_gen"].to_numpy()
    )


def test_marginal_table_pl_rejects_missing_metadata(tmp_path):
    path = tmp_path / "no_metadata.parquet"
    _write_predicted_local_parquet(path, metadata=None)
    with pytest.raises(ValueError, match="no '.*' parquet metadata"):
        marginal_table_pl(path)


def test_marginal_table_pl_rejects_global_coord(tmp_path):
    path = tmp_path / "global.parquet"
    _write_predicted_local_parquet(
        path,
        metadata={
            PREDICT_COORD_METADATA_KEY: "global",
            PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
        },
    )
    with pytest.raises(ValueError, match="coord=local"):
        marginal_table_pl(path)


def test_marginal_table_pl_rejects_mismatched_schema_version(tmp_path):
    path = tmp_path / "old_version.parquet"
    _write_predicted_local_parquet(
        path,
        metadata={
            PREDICT_COORD_METADATA_KEY: "local",
            PREDICT_SCHEMA_VERSION_KEY: "999",
        },
    )
    with pytest.raises(ValueError, match="schema version"):
        marginal_table_pl(path)


@pytest.mark.parametrize("group_by", [None, "pdg", "material", "energy"])
def test_plot_kl_bars_pl_runs_without_error(tmp_path, group_by):
    fig = plot_kl_bars_pl(_predicted_local_path(tmp_path), group_by=group_by)
    assert fig is not None


def test_plot_marginals_runs_without_error(tmp_path):
    fig = plot_marginals(_predicted_local_path(tmp_path))
    assert len(fig.axes) == len(RAW_TARGET_NAMES)  # single "all" row


def test_plot_marginals_grouped_runs_without_error(tmp_path):
    fig = plot_marginals(_predicted_local_path(tmp_path), group_by="material")
    assert fig is not None


def test_plot_marginals_caps_groups(tmp_path):
    path = _predicted_local_path(tmp_path, n=400)
    df = pl.read_parquet(path).with_columns(
        pl.Series("pdg", np.arange(400) % 8, dtype=pl.Int64)
    )
    fig = plot_marginals(df.lazy(), group_by="pdg", max_groups=3)
    n_rows = len(fig.axes) // len(RAW_TARGET_NAMES)
    assert n_rows <= 3


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


def test_correlation_matrices_pl_matches_numpy(tmp_path):
    path = _predicted_local_path(tmp_path)
    real, gen, *_ = _raw_from_parquet(path)

    real_corr, gen_corr = correlation_matrices_pl(path)

    np.testing.assert_allclose(
        real_corr, np.corrcoef(real, rowvar=False), atol=1e-4, rtol=1e-4
    )
    np.testing.assert_allclose(
        gen_corr, np.corrcoef(gen, rowvar=False), atol=1e-4, rtol=1e-4
    )
    for corr in (real_corr, gen_corr):
        np.testing.assert_allclose(np.diag(corr), 1.0, atol=1e-6)
        np.testing.assert_allclose(corr, corr.T, atol=1e-6)


def test_plot_correlation_matrices_runs_without_error(tmp_path):
    fig = plot_correlation_matrices(_predicted_local_path(tmp_path))
    assert len(fig.axes) >= 3  # real, generated, difference (+ colorbars)


def test_plot_pairwise_runs_without_error(tmp_path):
    fig = plot_pairwise(_predicted_local_path(tmp_path), n_sample=100)
    assert len(fig.axes) == 6  # 2 rows (real/gen) × 3 default pairs


def test_plot_direction_alignment_runs_without_error(tmp_path):
    fig = plot_direction_alignment(_predicted_local_path(tmp_path))
    assert fig is not None


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


def test_constraint_report_pl_matches_numpy_oracle(tmp_path):
    path = _predicted_local_path(tmp_path)
    _real, gen, *_ = _raw_from_parquet(path)

    expected = _numpy_constraint_report(gen)
    actual = constraint_report_pl(path).to_pandas()

    assert list(expected["check"]) == list(actual["check"])
    np.testing.assert_allclose(
        expected["violation_rate"].to_numpy(),
        actual["violation_rate"].to_numpy(),
        atol=1e-6,
    )
    np.testing.assert_allclose(
        expected["mean_abs_error"].to_numpy(),
        actual["mean_abs_error"].to_numpy(),
        atol=1e-4,
    )


def test_constraint_report_pl_flags_bad_direction_norms(tmp_path):
    path = _predicted_local_path(tmp_path)
    # Force the generated post_dir off the unit sphere for every row: the fixed
    # (1, 1, 1) vector has norm √3 ≈ 1.73, well outside the tolerance.
    df = pl.read_parquet(path).with_columns(
        pl.lit(1.0).alias("pred_post_dx"),
        pl.lit(1.0).alias("pred_post_dy"),
        pl.lit(1.0).alias("pred_post_dz"),
    )
    report = constraint_report_pl(df.lazy()).to_pandas()
    rate = dict(zip(report["check"], report["violation_rate"]))
    assert rate["post_dir unit norm"] == 1.0


def test_plot_constraint_violations_runs_without_error(tmp_path):
    fig = plot_constraint_violations(_predicted_local_path(tmp_path))
    assert len(fig.axes) == 2 + 3  # 2 direction norms + 3 scalar dims


# ---------------------------------------------------------------------------
# Tier 4: event-level (shower) observables
# ---------------------------------------------------------------------------


def _make_event_level_arrays(rng):
    """3 events (3/2/4 steps), each with an unambiguous highest-pre_E row.

    The forced max-pre_E rows (indices 1, 3, 7) fix a known shower axis/entry
    point per event, so the expected event_table can be re-derived independently
    in the test without depending on compute_event_observables_pl.
    """
    event_id = np.array([0, 0, 0, 1, 1, 2, 2, 2, 2], dtype=np.int64)
    n = len(event_id)
    pre_pos = rng.uniform(-5.0, 5.0, (n, 3)).astype(np.float32)
    pre_dir = _unit_vectors(rng, n)
    pre_E = rng.uniform(1.0, 50.0, n).astype(np.float32)
    pre_E[1] = 100.0  # event 0's entry step
    pre_E[3] = 100.0  # event 1's entry step
    pre_E[7] = 100.0  # event 2's entry step

    def _local_block():
        block = rng.standard_normal((n, 9)).astype(np.float32)
        block[:, 0] = log_transform(rng.uniform(0.1, 5.0, n).astype(np.float32))
        # cols 1–2 stay as random ALR energy logits (decoded against pre_E)
        block[:, 3:6] = _unit_vectors(rng, n)
        block[:, 6:9] = _unit_vectors(rng, n)
        return block

    true_log_local = _local_block()
    pred_log_local = _local_block()
    return event_id, pre_pos, pre_dir, pre_E, true_log_local, pred_log_local


def _write_event_level_parquet(path, rng=None):
    rng = rng or np.random.default_rng(7)
    event_id, pre_pos, pre_dir, pre_E, true_log_local, pred_log_local = (
        _make_event_level_arrays(rng)
    )
    n = len(event_id)
    pdg = rng.choice([11, -11, 22], n)
    table = pa.table(
        {
            "event_id": event_id,
            "pdg": pdg,
            "pre_x": pre_pos[:, 0],
            "pre_y": pre_pos[:, 1],
            "pre_z": pre_pos[:, 2],
            "pre_E": pre_E,
            "pre_dx": pre_dir[:, 0],
            "pre_dy": pre_dir[:, 1],
            "pre_dz": pre_dir[:, 2],
            "material": rng.choice(["W", "Pb"], n),
            "layer_id": rng.integers(0, 10, n).astype(np.int32),
            "n_sec": rng.integers(0, 3, n).astype(np.int32),
            **{
                f"pred_{name}": pred_log_local[:, j]
                for j, name in enumerate(LOCAL_TARGET_NAMES)
            },
            **{
                f"true_{name}": true_log_local[:, j]
                for j, name in enumerate(LOCAL_TARGET_NAMES)
            },
        }
    )
    table = table.replace_schema_metadata(
        {
            PREDICT_COORD_METADATA_KEY: "local",
            PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
        }
    )
    pq.write_table(table, path)
    return event_id, pre_pos, pre_dir, pre_E, true_log_local, pred_log_local, pdg


def _expected_event_table(
    event_id, pre_pos, pre_dir, pre_E, true_log_local, pred_log_local
):
    """Independent re-derivation of total/centroid/RMS per event, for comparison."""
    expected = {}
    for e in sorted(np.unique(event_id).tolist()):
        mask = event_id == e
        entry_idx = np.where(mask)[0][np.argmax(pre_E[mask])]
        entry_pos = pre_pos[entry_idx]
        axis_dir = pre_dir[entry_idx]

        def agg(log_local, mask=mask, entry_pos=entry_pos, axis_dir=axis_dir):
            step_length = inv_log_transform(log_local[mask, 0])
            edep, _e_sec, _post_E, _delta_e = energy_simplex_decode(
                log_local[mask, 1:3], pre_E[mask]
            )
            travel_dir_local = log_local[mask, 6:9]
            post_pos = reconstruct_post_pos(
                pre_pos[mask], pre_dir[mask], step_length, travel_dir_local
            )
            disp = post_pos - entry_pos
            depth = disp @ axis_dir
            transverse = np.linalg.norm(disp - depth[:, None] * axis_dir, axis=1)
            total_edep = float(edep.sum())
            total_length = float(step_length.sum())
            centroid = float((edep * depth).sum() / total_edep)
            rms = float(np.sqrt((edep * transverse**2).sum() / total_edep))
            return total_edep, total_length, centroid, rms

        real_total_edep, real_total_length, real_centroid, real_rms = agg(
            true_log_local
        )
        gen_total_edep, gen_total_length, gen_centroid, gen_rms = agg(pred_log_local)
        expected[e] = (
            real_total_edep,
            gen_total_edep,
            real_total_length,
            gen_total_length,
            real_centroid,
            gen_centroid,
            real_rms,
            gen_rms,
        )
    return expected


def test_compute_event_observables_pl_matches_manual_reconstruction(tmp_path):
    path = tmp_path / "event_level.parquet"
    event_id, pre_pos, pre_dir, pre_E, true_log_local, pred_log_local, _pdg = (
        _write_event_level_parquet(path)
    )
    expected = _expected_event_table(
        event_id, pre_pos, pre_dir, pre_E, true_log_local, pred_log_local
    )

    obs = compute_event_observables_pl(path, depth_bins=5, transverse_bins=5)
    real_table = obs.real_table.sort("event_id")
    gen_table = obs.gen_table.sort("event_id")

    for i, eid in enumerate(real_table["event_id"].to_list()):
        (
            real_total_edep,
            gen_total_edep,
            real_total_length,
            gen_total_length,
            real_centroid,
            gen_centroid,
            real_rms,
            gen_rms,
        ) = expected[eid]
        np.testing.assert_allclose(
            real_table["total_edep"][i], real_total_edep, rtol=1e-4
        )
        np.testing.assert_allclose(
            gen_table["total_edep"][i], gen_total_edep, rtol=1e-4
        )
        np.testing.assert_allclose(
            real_table["total_length"][i], real_total_length, rtol=1e-4
        )
        np.testing.assert_allclose(
            gen_table["total_length"][i], gen_total_length, rtol=1e-4
        )
        np.testing.assert_allclose(
            real_table["centroid_depth"][i], real_centroid, rtol=1e-3, atol=1e-4
        )
        np.testing.assert_allclose(
            gen_table["centroid_depth"][i], gen_centroid, rtol=1e-3, atol=1e-4
        )
        np.testing.assert_allclose(
            real_table["transverse_rms"][i], real_rms, rtol=1e-3, atol=1e-4
        )
        np.testing.assert_allclose(
            gen_table["transverse_rms"][i], gen_rms, rtol=1e-3, atol=1e-4
        )


def test_compute_event_observables_pl_profile_shapes(tmp_path):
    path = tmp_path / "event_level.parquet"
    _write_event_level_parquet(path)
    obs = compute_event_observables_pl(path, depth_bins=7, transverse_bins=4)

    assert obs.depth_edges.shape == (8,)
    assert obs.transverse_edges.shape == (5,)
    assert obs.real_depth_profile.shape == (7,)
    assert obs.gen_depth_profile.shape == (7,)
    assert obs.real_transverse_profile.shape == (4,)
    assert obs.gen_transverse_profile.shape == (4,)
    assert len(obs.real_table) == 3
    assert len(obs.gen_table) == 3


def test_compute_event_observables_pl_accepts_lazyframe(tmp_path):
    path = tmp_path / "event_level.parquet"
    _write_event_level_parquet(path)

    obs_from_path = compute_event_observables_pl(path, depth_bins=5, transverse_bins=5)
    obs_from_lf = compute_event_observables_pl(
        pl.scan_parquet(path), depth_bins=5, transverse_bins=5
    )

    np.testing.assert_allclose(
        obs_from_lf.real_table.sort("event_id")["total_edep"].to_numpy(),
        obs_from_path.real_table.sort("event_id")["total_edep"].to_numpy(),
    )


def test_compute_event_observables_pl_approx_median(tmp_path):
    """The streaming log-bin median approximates the true per-event median."""
    rng = np.random.default_rng(11)
    n = 4000  # one event, many steps → a well-defined median
    true_log = rng.standard_normal((n, 9)).astype(np.float32)
    true_log[:, 0] = log_transform(rng.uniform(0.1, 5.0, n).astype(np.float32))
    for s in (3, 6):
        v = rng.standard_normal((n, 3)).astype(np.float32)
        true_log[:, s : s + 3] = v / np.linalg.norm(v, axis=1, keepdims=True)
    pre_E = rng.uniform(1.0, 100.0, n).astype(np.float32)
    pre = rng.standard_normal((n, 3)).astype(np.float32)
    pdir = _unit_vectors(rng, n)
    pre_E[0] = 1000.0  # entry step

    path = tmp_path / "one_event.parquet"
    table = pa.table(
        {
            "event_id": np.zeros(n, dtype=np.int64),
            "pdg": rng.choice([11, 22], n),
            "pre_x": pre[:, 0],
            "pre_y": pre[:, 1],
            "pre_z": pre[:, 2],
            "pre_E": pre_E,
            "pre_dx": pdir[:, 0],
            "pre_dy": pdir[:, 1],
            "pre_dz": pdir[:, 2],
            "material": rng.choice(["W", "Pb"], n),
            "layer_id": rng.integers(0, 10, n).astype(np.int32),
            "n_sec": rng.integers(0, 3, n).astype(np.int32),
            **{f"pred_{nm}": true_log[:, j] for j, nm in enumerate(LOCAL_TARGET_NAMES)},
            **{f"true_{nm}": true_log[:, j] for j, nm in enumerate(LOCAL_TARGET_NAMES)},
        }
    ).replace_schema_metadata(
        {
            PREDICT_COORD_METADATA_KEY: "local",
            PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
        }
    )
    pq.write_table(table, path)

    true_edep = energy_simplex_decode(true_log[:, 1:3], pre_E)[0]
    true_length = inv_log_transform(true_log[:, 0])
    exact_edep_median = np.median(true_edep)
    exact_length_median = np.median(true_length)

    obs = compute_event_observables_pl(path)
    approx_edep = obs.real_table["median_edep"][0]
    approx_length = obs.real_table["median_length"][0]

    # log-bin interpolation: expect within a few percent of the true median.
    np.testing.assert_allclose(approx_edep, exact_edep_median, rtol=0.05)
    np.testing.assert_allclose(approx_length, exact_length_median, rtol=0.05)


def test_event_level_plots_run_without_error(tmp_path):
    path = tmp_path / "event_level.parquet"
    _write_event_level_parquet(path)
    obs = compute_event_observables_pl(path, depth_bins=5, transverse_bins=5)

    assert plot_total_energy(obs) is not None
    assert plot_total_length(obs) is not None
    assert plot_longitudinal_profile(obs) is not None
    assert plot_transverse_profile(obs) is not None
    assert plot_shower_max_depth(obs) is not None


# ---------------------------------------------------------------------------
# Particle-species (pdg) contribution shares
# ---------------------------------------------------------------------------


def test_pdg_contribution_table_pl_matches_manual_sums(tmp_path):
    path = tmp_path / "event_level.parquet"
    _, _, _, pre_E, true_log_local, pred_log_local, pdg = _write_event_level_parquet(
        path
    )

    real_edep = energy_simplex_decode(true_log_local[:, 1:3], pre_E)[0]
    gen_edep = energy_simplex_decode(pred_log_local[:, 1:3], pre_E)[0]
    real_length = inv_log_transform(true_log_local[:, 0])
    gen_length = inv_log_transform(pred_log_local[:, 0])

    expected = {}
    for p in np.unique(pdg):
        mask = pdg == p
        expected[int(p)] = (
            float(real_edep[mask].sum()),
            float(gen_edep[mask].sum()),
            float(real_length[mask].sum()),
            float(gen_length[mask].sum()),
        )

    table = pdg_contribution_table_pl(path).sort("pdg")
    for i, p in enumerate(table["pdg"].to_list()):
        real_e, gen_e, real_l, gen_l = expected[int(p)]
        np.testing.assert_allclose(table["real_total_edep"][i], real_e, rtol=1e-4)
        np.testing.assert_allclose(table["gen_total_edep"][i], gen_e, rtol=1e-4)
        np.testing.assert_allclose(table["real_total_length"][i], real_l, rtol=1e-4)
        np.testing.assert_allclose(table["gen_total_length"][i], gen_l, rtol=1e-4)


def test_pdg_contribution_table_pl_accepts_lazyframe(tmp_path):
    path = tmp_path / "event_level.parquet"
    _write_event_level_parquet(path)

    from_path = pdg_contribution_table_pl(path).sort("pdg")
    from_lf = pdg_contribution_table_pl(pl.scan_parquet(path)).sort("pdg")

    np.testing.assert_allclose(
        from_lf["real_total_edep"].to_numpy(), from_path["real_total_edep"].to_numpy()
    )


def test_pdg_pie_plots_run_without_error(tmp_path):
    path = tmp_path / "event_level.parquet"
    _write_event_level_parquet(path)
    table = pdg_contribution_table_pl(path)

    assert plot_pdg_energy_share(table) is not None
    assert plot_pdg_length_share(table) is not None


def test_plot_pdg_energy_share_caps_slices():
    table = pl.DataFrame(
        {
            "pdg": list(range(10)),
            "real_total_edep": [float(10 - i) for i in range(10)],
            "gen_total_edep": [float(10 - i) for i in range(10)],
            "real_total_length": [float(10 - i) for i in range(10)],
            "gen_total_length": [float(10 - i) for i in range(10)],
        }
    )
    fig = plot_pdg_energy_share(table, max_slices=4)
    for ax in fig.axes:
        assert len(ax.patches) == 4


# ---------------------------------------------------------------------------
# Rollout vs. held-out truth: unpaired `RolloutVsTruth` source (Tier 1-3) and
# `compute_rollout_vs_truth_observables_pl` (Tier 4)
# ---------------------------------------------------------------------------

_WORLD_FRAME_COLS = [
    "event_id",
    "pdg",
    "material",
    "layer_id",
    "pre_x",
    "pre_y",
    "pre_z",
    "pre_E",
    "pre_dx",
    "pre_dy",
    "pre_dz",
    "post_x",
    "post_y",
    "post_z",
    "post_E",
    "post_dx",
    "post_dy",
    "post_dz",
    "edep",
    "step_length",
]


def _write_world_frame_parquet(
    path, n=200, rng=None, rollout=True, termination_reasons=None
):
    """A `giant rollout` output or truth-schema steps file, in raw world-frame units.

    `rollout=True` adds `track_id`/`termination_reason`/`n_sec_pred` (rollout-
    only columns `RolloutVsTruth` reads) and tags `ROLLOUT_COORD_VALUE`
    metadata; `rollout=False` is a plain truth-schema file (no metadata tag,
    matching a real held-out/val parquet). `termination_reasons`, if given,
    overrides the (rollout-only) `termination_reason` column — pass a fixed
    array to test synthetic-row dropping deterministically.
    """
    rng = rng or np.random.default_rng(0)
    event_id = rng.integers(0, 5, n)
    pre_pos = rng.standard_normal((n, 3)).astype(np.float32) * 5
    pre_dir = _unit_vectors(rng, n)
    pre_E = rng.uniform(1.0, 100.0, n).astype(np.float32)
    post_dir = _unit_vectors(rng, n)
    step_length = rng.uniform(0.1, 5.0, n).astype(np.float32)
    travel_dir_world = _unit_vectors(rng, n)
    post_pos = pre_pos + step_length[:, None] * travel_dir_world
    post_E = (pre_E * rng.uniform(0.3, 0.99, n)).astype(np.float32)
    edep = rng.uniform(0, 1, n).astype(np.float32)

    columns = {
        "event_id": event_id,
        "pdg": rng.choice([11, -11, 22], n),
        "material": rng.choice(["W", "Pb"], n),
        "layer_id": rng.integers(0, 10, n).astype(np.int32),
        "pre_x": pre_pos[:, 0],
        "pre_y": pre_pos[:, 1],
        "pre_z": pre_pos[:, 2],
        "pre_E": pre_E,
        "pre_dx": pre_dir[:, 0],
        "pre_dy": pre_dir[:, 1],
        "pre_dz": pre_dir[:, 2],
        "post_x": post_pos[:, 0],
        "post_y": post_pos[:, 1],
        "post_z": post_pos[:, 2],
        "post_E": post_E,
        "post_dx": post_dir[:, 0],
        "post_dy": post_dir[:, 1],
        "post_dz": post_dir[:, 2],
        "edep": edep,
        "step_length": step_length,
    }
    metadata = None
    if rollout:
        columns["track_id"] = rng.integers(0, 3, n).astype(np.int64)
        columns["termination_reason"] = (
            termination_reasons
            if termination_reasons is not None
            else rng.choice(["", TERM_NATURAL_END], n)
        )
        columns["n_sec_pred"] = rng.integers(0, 3, n).astype(np.int32)
        metadata = {PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE}

    table = pa.table(columns)
    if metadata is not None:
        table = table.replace_schema_metadata(metadata)
    pq.write_table(table, path)
    return columns


def _raw_from_world_frame(columns):
    """(N, 9) `RAW_TARGET_NAMES` oracle for a world-frame steps dict, via numpy.

    Independent re-derivation of `_world_frame_local_exprs`'s forward Rodrigues
    rotation, using `giant.data.transforms.local_frame_rotation`/
    `travel_direction` directly rather than the module's polars expressions.
    """
    pre_pos = np.column_stack(
        [columns["pre_x"], columns["pre_y"], columns["pre_z"]]
    ).astype(np.float32)
    pre_dir = np.column_stack(
        [columns["pre_dx"], columns["pre_dy"], columns["pre_dz"]]
    ).astype(np.float32)
    post_pos = np.column_stack(
        [columns["post_x"], columns["post_y"], columns["post_z"]]
    ).astype(np.float32)
    post_dir = np.column_stack(
        [columns["post_dx"], columns["post_dy"], columns["post_dz"]]
    ).astype(np.float32)
    pre_E = np.asarray(columns["pre_E"], dtype=np.float32)
    post_E = np.asarray(columns["post_E"], dtype=np.float32)

    post_dir_local = local_frame_rotation(pre_dir, post_dir)
    travel_dir_local = local_frame_rotation(
        pre_dir, travel_direction(pre_pos, post_pos)
    )
    return np.column_stack(
        [
            np.asarray(columns["step_length"], dtype=np.float32),
            pre_E - post_E,
            np.asarray(columns["edep"], dtype=np.float32),
            post_dir_local,
            travel_dir_local,
        ]
    ).astype(np.float32)


def test_rollout_vs_truth_decodes_raw_targets_correctly(tmp_path):
    rng = np.random.default_rng(1)
    rollout_path = tmp_path / "rollout.parquet"
    truth_path = tmp_path / "truth.parquet"
    rollout_cols = _write_world_frame_parquet(
        rollout_path, n=150, rng=rng, rollout=True
    )
    truth_cols = _write_world_frame_parquet(truth_path, n=200, rng=rng, rollout=False)

    expected_gen = _raw_from_world_frame(rollout_cols)
    expected_real = _raw_from_world_frame(truth_cols)

    source = RolloutVsTruth(rollout=rollout_path, truth=truth_path)
    table = marginal_table_pl(source).to_pandas().set_index("dim")

    assert (table["n"] == len(truth_cols["pdg"])).all()
    assert (table["n_gen"] == len(rollout_cols["pdg"])).all()
    for j, name in enumerate(RAW_TARGET_NAMES):
        np.testing.assert_allclose(
            table.loc[name, "real_mean"], expected_real[:, j].mean(), atol=1e-3
        )
        np.testing.assert_allclose(
            table.loc[name, "gen_mean"], expected_gen[:, j].mean(), atol=1e-3
        )


def test_rollout_vs_truth_allows_unpaired_lengths(tmp_path):
    rng = np.random.default_rng(2)
    rollout_path = tmp_path / "rollout.parquet"
    truth_path = tmp_path / "truth.parquet"
    _write_world_frame_parquet(rollout_path, n=30, rng=rng, rollout=True)
    _write_world_frame_parquet(truth_path, n=500, rng=rng, rollout=False)

    source = RolloutVsTruth(rollout=rollout_path, truth=truth_path)
    table = marginal_table_pl(source)
    assert (table["n"] == 500).all()
    assert (table["n_gen"] == 30).all()


def test_rollout_vs_truth_drops_synthetic_termination_rows(tmp_path):
    rng = np.random.default_rng(3)
    rollout_path = tmp_path / "rollout.parquet"
    truth_path = tmp_path / "truth.parquet"
    n = 100
    # Half the rollout rows are synthetic bookkeeping (dropped), half are real.
    reasons = np.where(np.arange(n) % 2 == 0, TERM_ENERGY_CUTOFF, "")
    rollout_cols = _write_world_frame_parquet(
        rollout_path, n=n, rng=rng, rollout=True, termination_reasons=reasons
    )
    _write_world_frame_parquet(truth_path, n=50, rng=rng, rollout=False)

    kept = reasons != TERM_ENERGY_CUTOFF
    expected_gen = _raw_from_world_frame(
        {k: np.asarray(v)[kept] for k, v in rollout_cols.items()}
    )

    source = RolloutVsTruth(rollout=rollout_path, truth=truth_path)
    table = marginal_table_pl(source)
    assert (table["n_gen"] == kept.sum()).all()
    step_length_row = table.filter(pl.col("dim") == "step_length")
    np.testing.assert_allclose(
        step_length_row["gen_mean"][0], expected_gen[:, 0].mean(), atol=1e-3
    )


def test_rollout_vs_truth_rejects_wrong_coord_metadata(tmp_path):
    rollout_path = tmp_path / "rollout.parquet"
    truth_path = tmp_path / "truth.parquet"
    _write_world_frame_parquet(rollout_path, rollout=True)
    _write_world_frame_parquet(truth_path, rollout=False)
    # Overwrite with a mismatched coord tag (predict-schema "local").
    table = pq.read_table(rollout_path).replace_schema_metadata(
        {PREDICT_COORD_METADATA_KEY: "local"}
    )
    pq.write_table(table, rollout_path)

    source = RolloutVsTruth(rollout=rollout_path, truth=truth_path)
    with pytest.raises(ValueError, match="not a rollout file"):
        marginal_table_pl(source)


@pytest.mark.parametrize("group_by", [None, "pdg", "material", "energy"])
def test_rollout_vs_truth_downstream_plots_run_without_error(tmp_path, group_by):
    rollout_path = tmp_path / "rollout.parquet"
    truth_path = tmp_path / "truth.parquet"
    _write_world_frame_parquet(rollout_path, n=150, rollout=True)
    _write_world_frame_parquet(truth_path, n=150, rollout=False)
    source = RolloutVsTruth(rollout=rollout_path, truth=truth_path)

    assert plot_marginals(source, group_by=group_by) is not None
    assert plot_kl_bars_pl(source, group_by=group_by) is not None


def test_rollout_vs_truth_joint_and_constraint_checks_run(tmp_path):
    rollout_path = tmp_path / "rollout.parquet"
    truth_path = tmp_path / "truth.parquet"
    _write_world_frame_parquet(rollout_path, n=150, rollout=True)
    _write_world_frame_parquet(truth_path, n=150, rollout=False)
    source = RolloutVsTruth(rollout=rollout_path, truth=truth_path)

    assert plot_correlation_matrices(source) is not None
    assert plot_pairwise(source, n_sample=50) is not None
    assert plot_direction_alignment(source) is not None
    assert plot_constraint_violations(source) is not None
    report = constraint_report_pl(source).to_pandas()
    assert set(report["check"]) == {
        "post_dir unit norm",
        "travel_dir unit norm",
        "step_length >= 0",
        "delta_e >= 0",
        "edep >= 0",
    }


def _write_truth_shower(path, event_id, n_steps_per_event, rng):
    """One-track-per-event world-frame shower: highest-`pre_E` row fixes the axis."""
    rows_per_event = n_steps_per_event
    rows = []
    for i, eid in enumerate(event_id):
        n = rows_per_event
        pre_pos = rng.normal(size=(n, 3)).astype(np.float32) * 5
        pre_E = rng.uniform(1.0, 50.0, n).astype(np.float32)
        pre_E[0] = 1000.0  # entry step
        axis = _unit_vectors(rng, 1)[0]
        pre_dir = np.tile(axis, (n, 1)).astype(np.float32)
        step_length = rng.uniform(0.1, 5.0, n).astype(np.float32)
        post_pos = pre_pos + step_length[:, None] * axis
        edep = rng.uniform(0, 1, n).astype(np.float32)
        for j in range(n):
            rows.append(
                {
                    "event_id": eid,
                    "pdg": 11,
                    "material": "Pb",
                    "layer_id": 0,
                    "pre_x": pre_pos[j, 0],
                    "pre_y": pre_pos[j, 1],
                    "pre_z": pre_pos[j, 2],
                    "pre_E": pre_E[j],
                    "pre_dx": pre_dir[j, 0],
                    "pre_dy": pre_dir[j, 1],
                    "pre_dz": pre_dir[j, 2],
                    "post_x": post_pos[j, 0],
                    "post_y": post_pos[j, 1],
                    "post_z": post_pos[j, 2],
                    "edep": edep[j],
                    "step_length": step_length[j],
                    "track_id": i,
                    "termination_reason": TERM_NATURAL_END if j == n - 1 else "",
                    "n_sec_pred": 0,
                }
            )
    return rows


def test_compute_rollout_vs_truth_observables_pl_matches_manual_reconstruction(
    tmp_path,
):
    rng = np.random.default_rng(4)
    rollout_rows = _write_truth_shower(
        tmp_path / "_r", event_id=[0, 1, 2], n_steps_per_event=15, rng=rng
    )
    truth_rows = _write_truth_shower(
        tmp_path / "_t", event_id=[0, 1, 2], n_steps_per_event=20, rng=rng
    )

    rollout_path = tmp_path / "rollout.parquet"
    truth_path = tmp_path / "truth.parquet"
    pq.write_table(
        pa.table(
            {k: [r[k] for r in rollout_rows] for k in rollout_rows[0]}
        ).replace_schema_metadata({PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE}),
        rollout_path,
    )
    truth_cols = [
        "event_id",
        "pdg",
        "material",
        "layer_id",
        "pre_x",
        "pre_y",
        "pre_z",
        "pre_E",
        "pre_dx",
        "pre_dy",
        "pre_dz",
        "post_x",
        "post_y",
        "post_z",
        "edep",
        "step_length",
    ]
    pq.write_table(
        pa.table({k: [r[k] for r in truth_rows] for k in truth_cols}), truth_path
    )

    obs = compute_rollout_vs_truth_observables_pl(
        rollout_path, truth_path, depth_bins=5, transverse_bins=5
    )

    rollout_df = pd.DataFrame(rollout_rows)
    truth_df = pd.DataFrame(truth_rows)
    expected_gen_total = rollout_df.groupby("event_id")["edep"].sum().sort_index()
    expected_real_total = truth_df.groupby("event_id")["edep"].sum().sort_index()

    gen_table = obs.gen_table.sort("event_id")
    real_table = obs.real_table.sort("event_id")
    np.testing.assert_allclose(
        gen_table["total_edep"].to_numpy(), expected_gen_total.to_numpy(), rtol=1e-4
    )
    np.testing.assert_allclose(
        real_table["total_edep"].to_numpy(), expected_real_total.to_numpy(), rtol=1e-4
    )
    assert (gen_table["n_tracks"].to_numpy() == 1).all()
    assert (gen_table["leaked_E"].to_numpy() == 0).all()


def test_compute_rollout_vs_truth_observables_pl_profile_shapes(tmp_path):
    rng = np.random.default_rng(5)
    rollout_rows = _write_truth_shower(
        tmp_path / "_r", event_id=[0, 1], n_steps_per_event=10, rng=rng
    )
    truth_rows = _write_truth_shower(
        tmp_path / "_t", event_id=[0, 1, 2], n_steps_per_event=10, rng=rng
    )
    rollout_path = tmp_path / "rollout.parquet"
    truth_path = tmp_path / "truth.parquet"
    pq.write_table(
        pa.table(
            {k: [r[k] for r in rollout_rows] for k in rollout_rows[0]}
        ).replace_schema_metadata({PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE}),
        rollout_path,
    )
    truth_cols = [
        "event_id",
        "pdg",
        "material",
        "layer_id",
        "pre_x",
        "pre_y",
        "pre_z",
        "pre_E",
        "pre_dx",
        "pre_dy",
        "pre_dz",
        "post_x",
        "post_y",
        "post_z",
        "edep",
        "step_length",
    ]
    pq.write_table(
        pa.table({k: [r[k] for r in truth_rows] for k in truth_cols}), truth_path
    )

    obs = compute_rollout_vs_truth_observables_pl(
        rollout_path, truth_path, depth_bins=6, transverse_bins=4
    )
    assert obs.depth_edges.shape == (7,)
    assert obs.transverse_edges.shape == (5,)
    assert len(obs.gen_table) == 2
    assert len(obs.real_table) == 3

    assert plot_total_energy(obs) is not None
    assert plot_total_length(obs) is not None
    assert plot_mean_energy_per_step(obs) is not None
    assert plot_mean_length_per_step(obs) is not None
    assert plot_longitudinal_profile(obs) is not None
    assert plot_transverse_profile(obs) is not None
    assert plot_shower_max_depth(obs) is not None


def test_compute_rollout_vs_truth_observables_pl_rejects_wrong_coord_metadata(
    tmp_path,
):
    rollout_path = tmp_path / "rollout.parquet"
    truth_path = tmp_path / "truth.parquet"
    _write_world_frame_parquet(rollout_path, rollout=True)
    _write_world_frame_parquet(truth_path, rollout=False)
    table = pq.read_table(rollout_path).replace_schema_metadata(
        {PREDICT_COORD_METADATA_KEY: "local"}
    )
    pq.write_table(table, rollout_path)

    with pytest.raises(ValueError, match="not a rollout file"):
        compute_rollout_vs_truth_observables_pl(rollout_path, truth_path)