giant/tests/test_catalog.py

"""Tests for the plot catalog: id uniqueness + every spec computes a valid Reduced."""

from __future__ import annotations

import numpy as np
import pytest

from giant.analysis import build_catalog, catalog_ids, get_spec
from giant.analysis.catalog import (
    Bundle,
    PlotSpec,
    _containment_depths,
    _ks_statistic,
)
from giant.analysis.context import Context, build_context
from giant.analysis.grouping import pdg_label
from giant.analysis.sources import RolloutSpec
from tests.test_analysis_reduce import _reference_frame, _rollout_frame


def _build_ctx() -> Context:
    r, t = _rollout_frame(), _reference_frame()
    return build_context(
        [RolloutSpec("rollout", r)], t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000
    )


def _two_rollout_specs() -> list[RolloutSpec]:
    # Two distinct rollout sources so multi-series merging/finalize code is
    # exercised even though the underlying frame is the same fixture.
    return [RolloutSpec("flow", _rollout_frame()), RolloutSpec("wgan", _rollout_frame())]


@pytest.fixture(scope="module")
def ctx() -> Context:
    return _build_ctx()


@pytest.fixture(scope="module")
def two_ctx() -> Context:
    t = _reference_frame()
    return build_context(_two_rollout_specs(), t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000)


@pytest.fixture(scope="module")
def bundle(ctx: Context) -> Bundle:
    return Bundle.open([RolloutSpec("rollout", _rollout_frame())], _reference_frame(), ctx)


@pytest.fixture(scope="module")
def two_bundle(two_ctx: Context) -> Bundle:
    return Bundle.open(_two_rollout_specs(), _reference_frame(), two_ctx)


def test_catalog_ids_unique_and_nonempty():
    ids = catalog_ids()
    assert ids and len(ids) == len(set(ids))
    # the required families are all present
    fams = {s.family for s in build_catalog()}
    assert {"marginals", "event", "shower", "species", "secondaries"} <= fams


def test_get_spec_roundtrip_and_unknown():
    spec = get_spec("marginal_edep")
    assert spec.id == "marginal_edep" and spec.family == "marginals"
    with pytest.raises(KeyError):
        get_spec("does_not_exist")


def test_every_spec_computes_valid_reduced(bundle: Bundle):
    for spec in build_catalog():
        r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
        assert r.id == spec.id
        assert r.kind in {
            "overlay_hist",
            "grouped_hist",
            "profile",
            "bar",
            "single_hist",
            "router_gating",
            "router_share",
            "router_specialization",
            "heatmap",
            "unavailable",
        }
        assert r.title and r.xlabel
        _validate_payload(r, ["rollout"])


def test_every_spec_computes_valid_reduced_with_two_rollouts(two_bundle: Bundle):
    for spec in build_catalog():
        r = spec.finalize([spec.compute_partial(two_bundle)], two_bundle.ctx)
        assert r.id == spec.id
        _validate_payload(r, ["flow", "wgan"])


def _validate_payload(r, names: list[str]) -> None:
    p = r.payload
    if r.kind == "overlay_hist":
        n = len(p["edges"]) - 1
        assert list(p["series"]) == names
        for v in p["series"].values():
            assert len(v) == n
        assert len(p["reference"]) == n
    elif r.kind == "single_hist":
        assert list(p["series"]) == names
        for v in p["series"].values():
            assert len(v) == len(p["edges"]) - 1
    elif r.kind == "grouped_hist":
        n = len(p["edges"]) - 1
        assert p["groups"], "grouped hist must have at least one group"
        for g in p["groups"].values():
            assert list(g["series"]) == names
            for v in g["series"].values():
                assert len(v) == n
            assert len(g["reference"]) == n
    elif r.kind == "profile":
        n = len(p["edges"]) - 1
        assert list(p["series"]) == names
        for side in p["series"].values():
            assert len(side["mean"]) == n and len(side["std"]) == n
        assert len(p["reference"]["mean"]) == n and len(p["reference"]["std"]) == n
    elif r.kind == "bar":
        assert list(p["series"]) == names
        for v in p["series"].values():
            assert len(p["labels"]) == len(v)
        assert len(p["labels"]) == len(p["reference"])
    elif r.kind == "unavailable":
        assert p["note"]
    elif r.kind == "router_gating":
        for entry in p["series"].values():
            for side in ("rollout", "reference"):
                if side in entry:
                    assert len(entry[side]["centers"]) == len(entry[side]["means"])
    elif r.kind == "router_share":
        for entry in p["series"].values():
            for cat in entry["categories"]:
                for side in ("rollout", "reference"):
                    if side in entry:
                        assert cat in entry[side]
    elif r.kind == "router_specialization":
        for entry in p["series"].values():
            for side in ("rollout", "reference"):
                if side in entry:
                    assert len(entry[side]["centers"]) == len(entry[side]["score"])
    elif r.kind == "heatmap":
        assert list(p["series"]) == names
        for mat in p["series"].values():
            assert len(mat) == len(p["row_labels"])
            for row in mat:
                assert len(row) == len(p["col_labels"])


# ---------------------------------------------------------------------------
# chunked (compute_partial x N -> finalize) must match the unchunked (N=1) result
# ---------------------------------------------------------------------------

# One representative id per merge shape: sum-mergeable (marginal_edep,
# sec_count_per_species via pdg-keyed sums), concat-then-finalize with
# data-dependent edges (event_total_edep), concat-then-mean/std (shower_
# longitudinal), concat-then-max-edge (leakage_fraction), pdg-keyed sum with a
# ratio (species_edep_share), a chunkable=False passthrough (router_gating),
# sum-mergeable-with-a-zero-fill-denominator (sec_count_per_step{,_by_species}),
# nested sum-merge into a scorecard (marginal_distance_summary), and
# concat-then-per-event-derived-quantity
# (shower_containment_depth_90, reusing the profile matrix's own merge shape).
_CHUNK_EQUIVALENCE_IDS = [
    "marginal_edep",
    "species_edep_share",
    "event_total_edep",
    "shower_longitudinal",
    "leakage_fraction",
    "sec_count_per_species",
    "sec_count_per_step",
    "sec_count_per_step_by_species",
    "router_gating",
    "marginal_distance_summary",
    "shower_containment_depth_90",
]


def _assert_payload_close(a, b, path: str = "payload") -> None:
    """Recursively compare two JSON-shaped payloads (float-tolerant)."""
    assert type(a) is type(b), f"{path}: {type(a)} != {type(b)}"
    if isinstance(a, dict):
        assert set(a) == set(b), f"{path}: key mismatch {set(a)} != {set(b)}"
        for k in a:
            _assert_payload_close(a[k], b[k], f"{path}.{k}")
    elif isinstance(a, list):
        assert len(a) == len(b), f"{path}: length mismatch"
        for i, (x, y) in enumerate(zip(a, b)):
            _assert_payload_close(x, y, f"{path}[{i}]")
    elif isinstance(a, float):
        assert np.isclose(a, b, atol=1e-9), f"{path}: {a} != {b}"
    else:
        assert a == b, f"{path}: {a} != {b}"


@pytest.mark.parametrize("spec_id", _CHUNK_EQUIVALENCE_IDS)
def test_chunked_matches_unchunked(two_ctx: Context, spec_id: str):
    """A plot computed over N event-disjoint chunks then merged must equal the
    same plot computed in one unchunked pass — the core chunking correctness
    guarantee (see the analysis-rollout-plots chunking plan). Exercised with
    two rollout series so the per-rollout merge path is covered too."""
    spec: PlotSpec = get_spec(spec_id)
    rollouts, t = _two_rollout_specs(), _reference_frame()

    unchunked_bundle = Bundle.open(rollouts, t, two_ctx)
    unchunked = spec.finalize([spec.compute_partial(unchunked_bundle)], two_ctx)

    # 4 chunks over only 2 distinct event_ids also exercises empty chunks.
    n_chunks = 4 if spec.chunkable else 1
    parts = [spec.compute_partial(Bundle.open(rollouts, t, two_ctx, chunk=(k, n_chunks))) for k in range(n_chunks)]
    chunked = spec.finalize(parts, two_ctx)

    assert chunked.id == unchunked.id
    assert chunked.kind == unchunked.kind
    _assert_payload_close(unchunked.payload, chunked.payload)


# ---------------------------------------------------------------------------
# new (gitea #76) reductions: KS distance and containment depth
# ---------------------------------------------------------------------------


def test_ks_statistic():
    assert _ks_statistic([10, 10], [10, 10]) == 0.0  # identical shape -> 0
    assert _ks_statistic([10, 0], [0, 10]) == 1.0  # fully disjoint -> 1
    assert _ks_statistic([0, 0], [0, 0]) != _ks_statistic([0, 0], [0, 0])  # nan (no data either side)
    assert _ks_statistic([10, 0], [0, 0]) == 1.0  # one side empty, other isn't -> maximal mismatch


def test_containment_depths_simple_ramp():
    # one event, edep concentrated in the first bin -> 90%/95% containment
    # depth is the first bin's right edge; a zero-energy event is dropped.
    mat = np.array([[9.0, 1.0, 0.0], [0.0, 0.0, 0.0]])
    edges = np.array([0.0, 1.0, 2.0, 3.0])
    depths = _containment_depths(mat, edges, 0.90)
    assert depths.tolist() == [1.0]


def test_sec_count_per_step_counts_empty_steps(bundle):
    spec = get_spec("sec_count_per_step")
    r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
    # reference: 3 steps, two of which emit exactly one secondary
    assert r.payload["reference"][:2] == [1, 2]
    # rollout: 4 physical steps, one of which emits a single secondary
    assert r.payload["series"]["rollout"][:2] == [3, 1]
    assert sum(r.payload["reference"]) == 3


def test_sec_count_per_step_by_species_zero_row_is_per_species(bundle):
    spec = get_spec("sec_count_per_step_by_species")
    r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
    cols = r.payload["col_labels"]
    ref = r.payload["reference"]
    g = cols.index(pdg_label(22))
    # two reference steps emit one photon each; the third emits none
    assert [row[g] for row in ref][:2] == [1, 2]
    # every other species column is "no such secondary" on all 3 steps
    for j, _ in enumerate(cols):
        if j != g:
            assert ref[0][j] == 3 and sum(row[j] for row in ref[1:]) == 0