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,
    _integer_confusion,
    _ks_statistic,
)
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",
            "router_specialization",
            "heatmap",
            "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]
    elif r.kind == "router_specialization":
        for side in ("rollout", "reference"):
            if side in p:
                assert len(p[side]["centers"]) == len(p[side]["score"])
    elif r.kind == "heatmap":
        assert len(p["matrix"]) == len(p["row_labels"])
        for row in p["matrix"]:
            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),
# nested sum-merge into a scorecard (marginal_distance_summary), concat-then-
# event-id-join (n_sec_confusion), 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",
    "router_gating",
    "marginal_distance_summary",
    "n_sec_confusion",
    "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(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)


# ---------------------------------------------------------------------------
# new (gitea #76) reductions: KS distance, confusion matrix, 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_integer_confusion_matches_event_pairing():
    # true (reference) n_sec = [1, 1]; predicted (rollout) n_sec = [1, 0]
    labels, mat = _integer_confusion(np.array([1, 1]), np.array([1, 0]))
    assert labels == ["0", "1+"]
    assert mat.tolist() == [[0, 0], [1, 1]]  # row=true, col=pred


def test_integer_confusion_caps_pathological_outliers():
    labels, mat = _integer_confusion(np.array([0, 500]), np.array([0, 0]), max_bins=5)
    assert labels[-1] == "4+"
    assert mat.shape == (5, 5)
    assert mat.sum() == 2


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_n_sec_confusion_spec(bundle):
    spec = get_spec("n_sec_confusion")
    r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
    assert r.payload["row_labels"] == r.payload["col_labels"] == ["0", "1+"]
    assert r.payload["matrix"] == [[0, 0], [1, 1]]