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
from giant.analysis.context import Context, build_context
from tests.test_analysis_reduce import _reference_frame, _rollout_frame


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


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


@pytest.fixture(scope="module")
def bundle(ctx: Context) -> Bundle:
    return Bundle.open(_rollout_frame(), _reference_frame(), 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",
            "unavailable",
        }
        assert r.title and r.xlabel
        _validate_payload(r)


def _validate_payload(r) -> None:
    p = r.payload
    if r.kind == "overlay_hist":
        n = len(p["edges"]) - 1
        assert len(p["rollout"]) == n and len(p["reference"]) == n
    elif r.kind == "single_hist":
        assert len(p["rollout"]) == 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 len(g["rollout"]) == n and len(g["reference"]) == n
    elif r.kind == "profile":
        n = len(p["edges"]) - 1
        for k in ("rollout_mean", "rollout_std", "reference_mean", "reference_std"):
            assert len(p[k]) == n
    elif r.kind == "bar":
        assert len(p["labels"]) == len(p["rollout"]) == len(p["reference"])
    elif r.kind == "unavailable":
        assert p["note"]
    elif r.kind == "router_gating":
        for side in ("rollout", "reference"):
            if side in p:
                assert len(p[side]["centers"]) == len(p[side]["means"])
    elif r.kind == "router_share":
        for cat in p["categories"]:
            for side in ("rollout", "reference"):
                if side in p:
                    assert cat in p[side]


# ---------------------------------------------------------------------------
# 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), and a chunkable=False passthrough (router_gating).
_CHUNK_EQUIVALENCE_IDS = [
    "marginal_edep",
    "species_edep_share",
    "event_total_edep",
    "shower_longitudinal",
    "leakage_fraction",
    "sec_count_per_species",
    "router_gating",
]


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(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)."""
    spec: PlotSpec = get_spec(spec_id)
    r, t = _rollout_frame(), _reference_frame()

    unchunked_bundle = Bundle.open(r, t, ctx)
    unchunked = spec.finalize([spec.compute_partial(unchunked_bundle)], 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(r, t, ctx, chunk=(k, n_chunks)))
        for k in range(n_chunks)
    ]
    chunked = spec.finalize(parts, ctx)

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