tests/test_spec_diagnostics.py

"""Diagnostics closed-form spec tests (step 8).

Sources: the ablation spec §3 (CKA per Kornblith 2019
eqs (4)-(5); PCGrad surgery metrics per Yu 2020; ESS per Kish 1965;
§3.4 exact permutation test + bootstrap CI; §3.5 JS divergence).
Expected values are hand-computed in the docstrings. The craftax twin
file carries the same CKA/ESS/surgery and significance assertions --
`write_significance_test` is byte-identical across the repos -- while
the action-distribution and merge diagnostics are minihack-specific
(spec-ablations §3.5).
"""

from __future__ import annotations

import json
import math
import re

import numpy as np
import pytest
import torch
from scipy import stats as scipy_stats

from experiments.rl_finetuning.ablations.registry import REGISTRY
from experiments.rl_finetuning.ablations.training import (
    AblationHistory,
    _effective_batch_size,
)
from experiments.rl_finetuning.analysis.action_distribution import (
    compute_all_metrics,
    compute_entropy,
    compute_js,
    compute_kl,
    run_statistical_tests,
)
from experiments.rl_finetuning.analysis.plots import _ema
from experiments.rl_finetuning.analysis.report import (
    _HYPOTHESIS_GROUPS,
    _score_hypothesis,
)
from experiments.rl_finetuning.analysis.tables import (
    _macro_name,
    baseline_rl_score_of,
    make_forgetting_analysis_table,
    make_per_env_table,
    metric_scale,
    verdict,
    write_significance_test,
    write_tex_macros,
)
from experiments.rl_finetuning.diagnostics.gradient import compute_surgery_metrics
from experiments.rl_finetuning.diagnostics.representation import _linear_cka
from experiments.rl_finetuning.run_ablations import _merge_result_files


def test_linear_cka_is_one_for_identical_and_corr_squared_for_1d():
    """Linear CKA (Kornblith 2019 eqs (4)-(5)): CKA(X, X) = 1 and for
    1-D features CKA = corr^2. Same derivation and numbers as the
    craftax twin: x=[1,2,3,4], y=[1,3,2,4] -> CKA = 0.64.
    """
    x = torch.tensor([[1.0, 0.5], [2.0, -1.0], [-0.5, 0.25], [0.0, 3.0]])
    assert _linear_cka(x, x) == pytest.approx(1.0, abs=1e-5)
    x1 = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
    y1 = torch.tensor([[1.0], [3.0], [2.0], [4.0]])
    assert _linear_cka(x1, y1) == pytest.approx(0.64, abs=1e-5)


def test_linear_cka_is_invariant_to_scaling_and_orthogonal_maps():
    """CKA(X, c X Q) = 1 for isotropic c and orthogonal Q
    (Kornblith 2019 §2.3)."""
    x = torch.tensor([[1.0, 0.5], [2.0, -1.0], [-0.5, 0.25], [0.0, 3.0]])
    theta = 0.3
    q = torch.tensor(
        [
            [math.cos(theta), -math.sin(theta)],
            [math.sin(theta), math.cos(theta)],
        ]
    )
    assert _linear_cka(x, 2.5 * (x @ q)) == pytest.approx(1.0, abs=1e-5)


def test_effective_sample_size_closed_form():
    """ESS = (sum w)^2 / sum w^2 (Kish 1965): w=[1,1,2] -> 16/6;
    uniform weights give N. Same numbers as the craftax twin."""
    assert _effective_batch_size(torch.tensor([1.0, 1.0, 2.0])) == pytest.approx(
        16 / 6, rel=1e-6
    )
    assert _effective_batch_size(torch.ones(7)) == pytest.approx(7.0, rel=1e-6)


def test_surgery_metrics_measure_removed_gradient_mass():
    """Same derivation as the craftax twin: leaf a [2,0]->[1,0], leaf b
    unchanged -> fraction 3/29, one conflicting tensor."""
    before = {"a": torch.tensor([2.0, 0.0]), "b": torch.tensor([3.0, 4.0])}
    after = {"a": torch.tensor([1.0, 0.0]), "b": torch.tensor([3.0, 4.0])}
    frac, n_conf = compute_surgery_metrics(before, after)
    assert frac == pytest.approx(3 / 29, rel=1e-5)
    assert n_conf == 1


def test_kl_and_js_closed_forms():
    """KL and JS on hand-computable distributions (spec-ablations §3.5:
    JS(p,q) = KL(p||m)/2 + KL(q||m)/2, m = (p+q)/2, natural log).

    Derivation: p=[1,0], q=[0,1] -> m=[0.5,0.5], KL(p||m) = ln 2 ->
    JS = ln 2 (the eps=1e-10 smoothing perturbs this below 1e-4).
    KL(p,p) = JS(p,p) = 0; JS is symmetric.
    """
    p = np.array([1.0, 0.0])
    q = np.array([0.0, 1.0])
    assert compute_kl(p, p) == pytest.approx(0.0, abs=1e-8)
    assert compute_js(p, p) == pytest.approx(0.0, abs=1e-8)
    assert compute_js(p, q) == pytest.approx(math.log(2), abs=1e-4)
    assert compute_js(p, q) == pytest.approx(compute_js(q, p), abs=1e-12)


def _grad_alignment_setup(tiny_cfg, perturb: float):
    """A model displaced `perturb` from its pretrained reference, and a batch."""
    import copy

    import torch

    from src.diffusion.schedules import get_schedule
    from src.models.denoiser import make_model

    torch.manual_seed(0)
    tiny_cfg._schedule_fn = get_schedule(tiny_cfg.noise_schedule)

    ref_model = make_model(tiny_cfg)
    ref_model.eval()
    for param in ref_model.parameters():
        param.requires_grad = False

    model = copy.deepcopy(ref_model)
    for param in model.parameters():
        param.requires_grad = True
    if perturb:
        with torch.no_grad():
            for param in model.parameters():
                param.add_(torch.randn_like(param) * perturb)

    batch = 8
    local = torch.randint(0, 1000, (batch, tiny_cfg.crop_size, tiny_cfg.crop_size))
    glob = torch.randint(0, 1000, (batch, tiny_cfg.map_h, tiny_cfg.map_w))
    x0 = torch.randint(0, tiny_cfg.action_dim, (batch, tiny_cfg.seq_len))
    return model, ref_model, local.long(), glob.long(), x0.long(), torch.device("cpu")


def test_action_entropy_is_reported_in_nats():
    """Action-distribution entropy is in nats, and every column that carries
    it says so (spec-ablations §3.5; craftax `_compute_metrics`).

    Both repos reported "entropy" under one label with the unit stated
    nowhere: craftax natural log, minihack log base 2, a factor of
    1/ln 2 = 1.442695 apart. Canon is nats, which is what the NELBO and
    cross-entropy figures throughout both suites already use.

    Derivation: for [1/2, 1/4, 1/8, 1/8] the entropy is
    (1/2)ln2 + (1/4)ln4 + 2*(1/8)ln8 = 1.75 ln 2 = 1.2130075656 nats,
    which is 1.75 bits. Uniform over A actions is ln A: over the 8 actions
    below, 2.0794415417 nats against 3 bits.
    """
    probs = np.array([0.5, 0.25, 0.125, 0.125])
    assert compute_entropy(probs) == pytest.approx(1.2130075656, abs=1e-9)
    assert compute_entropy(probs) == pytest.approx(1.75 * math.log(2), abs=1e-12)

    uniform = np.full(8, 1.0 / 8.0)
    assert compute_entropy(uniform) == pytest.approx(math.log(8), abs=1e-12)

    def _entropy_stats():
        return {
            "action_counts": {},
            "episode_returns": np.array([0.0, 1.0]),
            "episode_won": np.array([0.0, 1.0]),
        }

    padded = np.concatenate([probs, np.zeros(4)])
    metrics = compute_all_metrics(
        uniform,
        padded,
        _entropy_stats(),
        _entropy_stats(),
        8,
    )
    assert metrics["Max Possible Entropy (nats)"] == pytest.approx(math.log(8), abs=1e-12)
    assert metrics["Pre-RL Entropy (nats)"] == pytest.approx(math.log(8), abs=1e-12)
    assert metrics["Entropy Change (nats)"] == pytest.approx(
        1.2130075656 - math.log(8), abs=1e-9
    )
    # Normalised entropy is a ratio, so it is unit-free and carries no suffix.
    assert metrics["Pre-RL Normalised Entropy"] == pytest.approx(1.0, abs=1e-12)


def test_the_action_distribution_chi_squared_compares_two_observed_samples():
    """The action-distribution chi-squared is a contingency test on two
    observed count vectors, not a goodness-of-fit test against one of them
    (spec-ablations §3.5).

    Both action count vectors are sampled. Handing one to
    ``scipy.stats.chisquare`` as the expectation asserts it is known
    exactly, which drops half the sampling error from the comparison and
    roughly doubles the statistic; the contingency form estimates the
    shared expectation from both margins instead. Degrees of freedom are
    A - 1 either way -- (2-1)(A-1) for the table -- so the two differ only
    in the expectation, and the statistic ratio below is that difference.

    Derivation of the null rate: with the same distribution generating
    both samples the test should reject at alpha = 0.05 on about 5 % of
    draws. Measured over 400 draws of 2000 actions across 8 actions, the
    goodness-of-fit form rejects on roughly 40 % and the contingency form
    on roughly 5 %.
    """
    action_dim = 8
    rng = np.random.default_rng(0)
    probs = rng.dirichlet(np.ones(action_dim) * 2.0)

    def _stats(counts):
        return {
            "action_counts": {i: int(c) for i, c in enumerate(counts)},
            "episode_returns": [0.0, 1.0, 2.0],
        }

    def _goodness_of_fit_p(pre, post):
        """The form this replaced: post rescaled and used as the expectation."""
        p = pre + 1.0
        q = post + 1.0
        return scipy_stats.chisquare(p, q * (p.sum() / q.sum()))[1]

    trials = 400
    gof_hits = contingency_hits = 0
    for _ in range(trials):
        pre = rng.multinomial(2000, probs).astype(float)
        post = rng.multinomial(2000, probs).astype(float)
        contingency_hits += run_statistical_tests(
            _stats(pre), _stats(post), action_dim
        )["chi2_significant"]
        gof_hits += _goodness_of_fit_p(pre, post) < 0.05

    assert gof_hits / trials > 0.25
    assert 0.01 < contingency_hits / trials < 0.10

    # A real shift is still detected.
    shifted = probs * 0.5
    shifted[0] += 0.5
    out = run_statistical_tests(
        _stats(rng.multinomial(5000, probs).astype(float)),
        _stats(rng.multinomial(5000, shifted).astype(float)),
        action_dim,
    )
    assert out["chi2_significant"]
    assert out["chi2_p"] < 1e-6


def test_grad_alignment_shares_one_draw_and_references_the_pretrained_params(tiny_cfg):
    """The RL and BC gradients come from one ``(z_t, t)`` draw, and the BC
    gradient is taken at the pretrained parameters (spec-ablations §3.2; the
    same definition as craftax's `make_grad_alignment_fn`).

    Derivation of the exact case: uniform advantages make the RL loss
    ``(per_sample * 1).mean()`` and the BC loss ``per_sample.mean()`` the
    same expression, so on one draw at one parameter point the two
    gradients are the same vector and the cosine is exactly 1. Anything
    less is the draw differing: at independent draws the metric is a
    Monte-Carlo estimate whose scatter is the size of the quantity, and it
    reports objective disagreement where there is none by construction.

    Displacing the model from the reference then drops the cosine below 1
    while nothing about the objectives has changed, which is what taking
    the BC gradient at a fixed pretrained reference means.
    """
    import torch

    from experiments.rl_finetuning.ablations.losses import _core_loss
    from experiments.rl_finetuning.diagnostics.gradient import (
        _at_reference_parameters,
        _collect_flat_grad,
        compute_grad_alignment,
    )

    model, ref_model, local, glob, x0, device = _grad_alignment_setup(tiny_cfg, 0.0)
    batch = x0.shape[0]
    uniform = torch.ones(batch)

    # One draw, one parameter point, one objective in two spellings.
    cos, rl_norm, bc_norm = compute_grad_alignment(
        model, ref_model, local, glob, x0, uniform, tiny_cfg, device
    )
    assert cos == pytest.approx(1.0, abs=1e-4)
    assert rl_norm == pytest.approx(bc_norm, rel=1e-5)

    def shipped_independent_draws() -> float:
        """What the metric was: a second draw, and the BC gradient at `model`."""
        model.train()
        model.zero_grad()
        _core_loss(model, local, glob, x0, uniform, tiny_cfg, device).backward()
        g_rl = _collect_flat_grad(model)
        model.zero_grad()
        _core_loss(model, local, glob, x0, None, tiny_cfg, device).backward()
        g_bc = _collect_flat_grad(model)
        model.zero_grad()
        return (torch.dot(g_rl, g_bc) / (g_rl.norm() * g_bc.norm() + 1e-10)).item()

    independent = [shipped_independent_draws() for _ in range(5)]
    assert max(independent) < 1.0 - 1e-3
    assert max(independent) - min(independent) > 1e-3

    # The reference is the pretrained point, not wherever the run has got to.
    model, ref_model, local, glob, x0, device = _grad_alignment_setup(tiny_cfg, 0.05)
    displaced, _, _ = compute_grad_alignment(
        model, ref_model, local, glob, x0, uniform, tiny_cfg, device
    )
    assert displaced < 1.0 - 1e-3

    # And the swap that gets it there puts every parameter back.
    before = torch.cat([p.detach().clone().reshape(-1) for p in model.parameters()])
    reference = torch.cat([p.detach().reshape(-1) for p in ref_model.parameters()])
    assert (before - reference).abs().max() > 1e-3
    with _at_reference_parameters(model, ref_model):
        inside = torch.cat([p.detach().reshape(-1) for p in model.parameters()])
        assert (inside - reference).abs().max() == pytest.approx(0.0, abs=1e-12)
    after = torch.cat([p.detach().reshape(-1) for p in model.parameters()])
    assert (after - before).abs().max() == pytest.approx(0.0, abs=1e-12)


def test_the_curve_smoother_leaves_a_gap_where_data_is_missing():
    """A missing evaluation is a hole in the record, not a measurement of
    zero, and `_ema` draws it as a gap (spec-ablations §3.9).

    A NaN metric round-trips through the results JSON as null and comes
    back as None. minihack substituted 0.0 for it and craftax raised a
    TypeError, so the same hole either invented a collapse or lost the
    figure.

    Derivation: with one hole in a flat 0.65 curve, substituting zero gives
    [0.65, 0.455, 0.5135, 0.5544, 0.5831] -- a 30 % drop and a four-point
    recovery that the run never had, on a win-rate axis where that is
    exactly the shape the suite is looking for. Carrying the hole through
    as NaN leaves the curve flat at 0.65 with one point missing, which
    matplotlib renders as a break in the line.

    A hole-free input is unchanged, so no existing figure moves: the
    recursion is the same expression, seeded from the first real value.
    """
    flat = [0.65, 0.65, 0.65, 0.65, 0.65]
    assert _ema(flat) == _ema([0.65, 0.65, 0.65, 0.65, 0.65])
    assert all(v == pytest.approx(0.65) for v in _ema(flat))

    holed = _ema([0.65, None, 0.65, 0.65, 0.65])
    assert math.isnan(holed[1])
    assert [v for i, v in enumerate(holed) if i != 1] == pytest.approx(
        [0.65, 0.65, 0.65, 0.65]
    )
    # The zero substitution this replaced would have produced these.
    assert holed[2] != pytest.approx(0.5135)

    # A NaN that never reached JSON behaves the same as the None it becomes.
    assert math.isnan(_ema([0.65, float("nan"), 0.65])[1])

    # Leading holes stay holes rather than seeding the average.
    leading = _ema([None, 1.0, 1.0])
    assert math.isnan(leading[0])
    assert leading[1:] == pytest.approx([1.0, 1.0])

    assert _ema([]) == []


def test_the_forgetting_table_is_one_definition_across_the_repos():
    """The forgetting table is one function in both repos, and each of the
    five places the two halves had drifted apart resolves the same way
    (spec-ablations §3.8).

    Derivation, boundary: multiplicative, ``pretrained * (1 - 0.1)``. At a
    pretrained score of 1.0 that is 0.9, so an evaluation of 0.85 is a
    collapse and 0.92 is not. The absolute form minihack used --
    ``pretrained - 0.05`` -- puts the boundary at 0.95 instead, which makes
    `dipped_but_not_collapsed` a collapse at iteration 20 rather than an
    arm that never collapsed. The two rules coincide only at a pretrained
    score of 0.5; on a Craftax achievement score the absolute 0.05 is a
    different fraction entirely, which is why the verdict rule was scaled
    to the metric on 2026-08-17.

    Derivation, recovery: `collapsed_then_recovered` drops to 0.85 at
    iteration 20 and climbs to 0.95, so `Recovered` is "Y". `healthy` never
    goes below 0.9, so it is "N/A" -- not recovery, because there was no
    collapse; the rule minihack used, final score at or above the boundary,
    calls it recovered.

    Derivation, recovery score: `score_differs_from_last_eval` has a
    terminal evaluation of 0.42 and a last in-loop evaluation of 0.99. The
    terminal one is what the main results table, the verdict rule and the
    hypothesis table all read, so `Recovery_Score` is 0.42.

    Derivation, empty history: `no_history` still gets a row, with a null
    minimum and no collapse. Dropping it would leave four rows where five
    arms ran, and any count taken over this table would silently change
    denominator.

    Derivation, order: the rows come out in sorted name order, so the CSV
    is byte-reproducible across runs.
    """
    results = {
        "healthy": {
            "history": AblationHistory(eval_iters=[10, 20, 30], eval_score=[1.0, 0.95, 0.98]),
            "score": 0.98,
        },
        "dipped_but_not_collapsed": {
            "history": AblationHistory(eval_iters=[10, 20, 30], eval_score=[1.0, 0.92, 0.97]),
            "score": 0.97,
        },
        "collapsed_then_recovered": {
            "history": AblationHistory(eval_iters=[10, 20, 30], eval_score=[1.0, 0.85, 0.95]),
            "score": 0.95,
        },
        "collapsed_and_stayed": {
            "history": AblationHistory(eval_iters=[10, 20, 30], eval_score=[1.0, 0.85, 0.20]),
            "score": 0.20,
        },
        "score_differs_from_last_eval": {
            "history": AblationHistory(eval_iters=[10, 20], eval_score=[1.0, 0.99]),
            "score": 0.42,
        },
        "no_history": {"history": AblationHistory(), "score": 0.5},
    }
    df = make_forgetting_analysis_table(results, pretrained_score=1.0)
    rows = {r["Method"]: r for r in df.to_dicts()}

    # Every arm gets a row, including the one with nothing to plot.
    assert df.shape[0] == 6
    assert rows["no_history"]["Min_Score"] is None
    assert rows["no_history"]["First_Collapse_Iter"] == "never"
    assert rows["no_history"]["Recovered"] == "N/A"

    # Sorted, so the CSV is reproducible.
    assert df["Method"].to_list() == sorted(results)

    # The boundary is 0.9, not 0.95: a dip to 0.92 is not a collapse.
    assert rows["dipped_but_not_collapsed"]["First_Collapse_Iter"] == "never"
    assert rows["dipped_but_not_collapsed"]["Recovered"] == "N/A"
    assert rows["healthy"]["First_Collapse_Iter"] == "never"
    assert rows["healthy"]["Recovered"] == "N/A"

    # Recovery is judged from the first collapse onward.
    assert rows["collapsed_then_recovered"]["First_Collapse_Iter"] == "20"
    assert rows["collapsed_then_recovered"]["Recovered"] == "Y"
    assert rows["collapsed_and_stayed"]["First_Collapse_Iter"] == "20"
    assert rows["collapsed_and_stayed"]["Recovered"] == "N"

    # The recovery score is the terminal evaluation, not the last in-loop one.
    assert rows["score_differs_from_last_eval"]["Recovery_Score"] == pytest.approx(0.42)
    assert rows["score_differs_from_last_eval"]["Min_Score"] == pytest.approx(0.99)
    assert rows["score_differs_from_last_eval"]["Min_Score_Iter"] == 20


def test_the_significance_test_states_its_floor_and_corrects_for_selection(tmp_path):
    """The significance test is exact over all C(n_a+n_b, n_b) relabellings,
    reports the floor that enumeration imposes, and draws its null
    distribution over every candidate arm rather than over the one it picked
    (spec-ablations §3.4; both repos' experiments/README tables).

    Derivation, floor: every relabelling's complement negates each mean
    difference and so ties the statistic, which makes the count at least two
    -- p >= 2/C(6,3) = 0.100 at three seeds a side, for any data whatsoever.
    Baseline [0,0,0] against [1e6,1e6,1e6] therefore reports p = 0.100, and
    0.100 has to be reported as the floor rather than left to read as
    marginal significance.

    Derivation, selection: baseline [0,1,2,3] against [4,5,6,7] has an
    observed difference of 4, which only the two extreme partitions of the
    70 relabellings reach -- p = 2/70 = 0.029 while that arm is the only
    candidate. The null arm [-6,-2,2,6] scores no better than baseline but
    is spread widely enough that its own relabellings reach a statistic of 4
    another twelve times, and it is a candidate the maximum must range over,
    so p becomes 14/70 = 0.200. Selecting the arm from the same scores and
    then testing it uncorrected reports 0.029 either way.
    """
    write_significance_test(
        {
            "baseline_rl": {"all_scores": [0.0, 0.0, 0.0]},
            "kl_penalty": {"all_scores": [1e6, 1e6, 1e6]},
        },
        tmp_path,
    )
    text = (tmp_path / "significance_test.txt").read_text()
    assert "20 relabellings" in text
    assert "p = 0.100" in text
    assert "minimum attainable p at 3 baseline and 3 condition seeds: 0.100" in text
    assert "AT the floor" in text

    alone = tmp_path / "alone"
    write_significance_test(
        {
            "baseline_rl": {"all_scores": [0.0, 1.0, 2.0, 3.0]},
            "kl_penalty": {"all_scores": [4.0, 5.0, 6.0, 7.0]},
        },
        alone,
    )
    text = (alone / "significance_test.txt").read_text()
    assert "1 candidate arm " in text
    assert "p = 0.029" in text
    ci_line = next(line for line in text.splitlines() if "bootstrap" in line)
    assert float(ci_line.split("[")[1].split(",")[0]) > 0.0

    with_null_arm = tmp_path / "with_null_arm"
    write_significance_test(
        {
            "baseline_rl": {"all_scores": [0.0, 1.0, 2.0, 3.0]},
            "kl_penalty": {"all_scores": [4.0, 5.0, 6.0, 7.0]},
            "ewc": {"all_scores": [-6.0, -2.0, 2.0, 6.0]},
        },
        with_null_arm,
    )
    text = (with_null_arm / "significance_test.txt").read_text()
    assert "2 candidate arms" in text
    assert "p = 0.200" in text


def test_merge_concatenates_scores_and_recomputes_over_the_union(tmp_path):
    """--merge concatenates per-seed scores for the same ablation and
    recomputes score/score_std over the union; the merged
    pretrained_score is the mean of the inputs (spec-ablations §1.3).

    Derivation: files with all_scores [1,2] and [3] merge to [1,2,3]:
    score = 2.0, score_std = population std = sqrt(2/3) = 0.8165;
    pretrained (0.4, 0.6) -> 0.5.
    """
    def _file(name, scores, pretrained):
        payload = {
            "pretrained_score": pretrained,
            "config": {"batch_size": 1},
            "ablations": {
                "baseline_rl": {
                    "score": float(np.mean(scores)),
                    "score_std": float(np.std(scores)),
                    "all_scores": scores,
                    "history": {},
                }
            },
        }
        path = tmp_path / name
        path.write_text(json.dumps(payload))
        return str(path)

    merged, pretrained, _ = _merge_result_files(
        [_file("a.json", [1.0, 2.0], 0.4), _file("b.json", [3.0], 0.6)]
    )
    assert merged["baseline_rl"]["all_scores"] == [1.0, 2.0, 3.0]
    assert merged["baseline_rl"]["score"] == pytest.approx(2.0)
    assert merged["baseline_rl"]["score_std"] == pytest.approx(
        math.sqrt(2 / 3), rel=1e-6
    )
    assert pretrained == pytest.approx(0.5)


def test_the_tex_macros_carry_the_numbers_the_manuscript_prints(tmp_path):
    """`results.tex` is how a generated number reaches the draft, so each
    macro must be a usable control sequence holding the value at the
    precision the manuscript prints it: win rates in percentage points,
    CV_A = sqrt(B / ESS - 1) averaged over iterations.

    Derivation: score 0.4375 -> 43.75. With B = 4608 and ESS 4608/2 and
    4608/5, CV_A = (sqrt(1) + sqrt(4)) / 2 = 1.50. Pooled seed sd over the
    one condition carrying seeds is its own sample sd: scores 0.3875 and
    0.4875 give sqrt(0.005) = 0.0707 -> 7.07 points.
    """
    history = AblationHistory(effective_batch_size=[4608 / 2, 4608 / 5])
    results = {
        "baseline_rl": {
            "score": 0.4375,
            "score_std": 0.0612,
            "all_scores": [0.3875, 0.4875],
            "history": history,
        },
        "layer_ablation_top1": {
            "score": 0.4125,
            "score_std": 0.01,
            "history": AblationHistory(),
        },
    }
    path = write_tex_macros(
        results, 0.475, tmp_path / "results.tex", {"batch_size": 4608}
    )
    text = path.read_text()

    names = re.findall(r"\\newcommand\{\\([A-Za-z]*)\}", text)
    for name in names:
        assert name.isalpha() and name.startswith("mh")
    # definitions only, one per line, and no name defined twice
    assert len(re.findall(r"\\newcommand", text)) == len(names)
    assert len(set(names)) == len(names)

    assert "\\newcommand{\\mhPretrainedScore}{47.50}" in text
    assert "\\newcommand{\\mhBatchSize}{4608}" in text
    assert "\\newcommand{\\mhScoreBaselineRl}{43.75}" in text
    assert "\\newcommand{\\mhScoreSdBaselineRl}{6.12}" in text
    assert "\\newcommand{\\mhCvABaselineRl}{1.50}" in text
    assert "\\newcommand{\\mhEssBaselineRl}{1613}" in text
    # deltas are magnitudes; the manuscript carries the sign
    assert "\\newcommand{\\mhDeltaPretrainedBaselineRl}{3.75}" in text
    assert "\\newcommand{\\mhDeltaBaselineBaselineRl}{0.00}" in text
    assert "\\newcommand{\\mhPooledSeedSd}{7.07}" in text
    # a digit in the condition name is spelled out, or the macro is unusable
    assert "\\newcommand{\\mhScoreLayerAblationTopOne}{41.25}" in text
    # no ESS recorded -> no CV_A macro invented for it
    assert f"mhCvA{_macro_name('layer_ablation_top1')}" not in text
    # group means come from the same table the CSV does
    assert "\\newcommand{\\mhGroupMeanBaseline}{43.75}" in text
    # and the file says which evaluation its per-layout numbers came from
    assert "% Per-layout macros come from the merged single-run history." in text


def test_the_macro_mangling_rule_matches_the_sibling_suite(tmp_path):
    """The manuscript inputs both suites' `results.tex`, so a condition must
    mangle to the same tag in both repositories and the prefixes must be the
    only thing separating them. `_macro_name` is the sibling's rule verbatim:
    `-`/`_` are word boundaries, digits are spelled out, every word is
    capitalised.
    """
    assert _macro_name("advantage_clip") == "AdvantageClip"
    assert _macro_name("layer_ablation_top1") == "LayerAblationTopOne"
    assert _macro_name("normalized_adv") == "NormalizedAdv"
    assert _macro_name("bc_wins") == "BcWins"
    # only the first character of each word is raised, digits and all
    assert _macro_name("Room-5x5") == "RoomFivexFive"
    assert _macro_name("Corridor-R3") == "CorridorRThree"
    assert _macro_name("group", "Baseline") == "GroupBaseline"

    # A collision silently redefines an earlier number, so it must raise
    # rather than write.
    results = {
        "baseline_rl": {"score": 0.4, "history": AblationHistory()},
        "baseline-rl": {"score": 0.9, "history": AblationHistory()},
    }
    with pytest.raises(ValueError, match="Duplicate"):
        write_tex_macros(results, 0.475, tmp_path / "results.tex")


def test_the_per_env_table_reads_the_evaluation_the_score_comes_from():
    """`score` is the mean of the post-training evaluation, so the
    per-environment table must be that same evaluation's detail
    (`per_seed_final_evals`), not the last in-loop one (`per_seed_finals`).

    Two draws of the same 80 episodes differ by several points, which is what
    made `tab:group_summary` fail to reconcile with `tab:per-env`. The
    in-loop record stays as a fallback for results files written before the
    final evaluation's detail was kept.
    """
    final = {"MiniHack-Room-v0": 0.5, "MiniHack-Corridor-v0": 0.7}
    in_loop = {"MiniHack-Room-v0": 0.1, "MiniHack-Corridor-v0": 0.3}
    history = AblationHistory(per_env_win_rates=[in_loop])

    both = {
        "baseline_rl": {
            "score": 0.6,
            "history": history,
            "per_seed_final_evals": [{"per_env_win_rates": final}],
            "per_seed_finals": [{"per_env_win_rates": in_loop}],
        }
    }
    row = make_per_env_table(both).to_dicts()[0]
    assert row["MiniHack-Room-v0"] == pytest.approx(0.5)
    assert row["MiniHack-Corridor-v0"] == pytest.approx(0.7)
    # the mean of the row is the score it is tabulated beside
    assert np.mean([row[k] for k in final]) == pytest.approx(0.6)

    legacy = {"baseline_rl": {k: v for k, v in both["baseline_rl"].items()
                              if k != "per_seed_final_evals"}}
    assert make_per_env_table(legacy).to_dicts()[0]["MiniHack-Room-v0"] == (
        pytest.approx(0.1)
    )


# ---------------------------------------------------------------------------
# Ablation-suite verdict rule (shared with the sibling repo, character for
# character; cross-implementation note open question resolved 2026-08-17)
# ---------------------------------------------------------------------------


def test_verdict_labels_against_baseline_rl_at_metric_scale():
    """Labels are taken against `baseline_rl`, with thresholds that are
    fractions of the metric's magnitude: IMPROVEMENT above +5%, COLLAPSE
    below -10%, NEUTRAL between.

    Derivation at scale 10 (`baseline_rl` 10.0, pretrained 8.0, the order
    of magnitude of a Craftax episode-weighted mean return): the
    improvement bar is +0.5 and the collapse bar -1.0, so 10.6 improves,
    10.4 does not, 9.1 holds and 8.9 collapses.

    The last case is the one the absolute rule got wrong. Constructed to
    the recorded shape: an arm sitting 1.911 below `baseline_rl` read
    IMPROVEMENT under the old craftax rule, because +0.089 against
    pretrained cleared an absolute +0.05 bar.
    """
    assert verdict(10.6, 10.0, 8.0) == "IMPROVEMENT"
    assert verdict(10.4, 10.0, 8.0) == "NEUTRAL"
    assert verdict(9.1, 10.0, 8.0) == "NEUTRAL"
    assert verdict(8.9, 10.0, 8.0) == "COLLAPSE"
    assert verdict(10.0 - 1.911, 10.0, 8.0) == "COLLAPSE"


def test_verdict_reduces_to_the_absolute_rule_at_a_metric_scale_of_one():
    """At scale 1.0 the fractions are the absolute +0.05 / -0.10 they
    replace, and both comparisons are strict.

    This is the anchor for a bounded metric: a MiniHack win rate lives in
    [0, 1], so the rule that governed it is unchanged in form. With
    `baseline_rl` 0.0 and pretrained 1.0 the scale is exactly 1.0 and the
    delta is the score itself, so the boundaries are exact in float.
    """
    assert verdict(0.05, 0.0, 1.0) == "NEUTRAL"
    assert verdict(0.06, 0.0, 1.0) == "IMPROVEMENT"
    assert verdict(-0.10, 0.0, 1.0) == "NEUTRAL"
    assert verdict(-0.11, 0.0, 1.0) == "COLLAPSE"


def test_verdict_scale_is_the_larger_reference_and_one_is_required():
    """The scale is the larger reference score in absolute value, so a
    `baseline_rl` near zero cannot shrink the threshold to nothing; with
    both references at zero there is no scale and no label is defensible.

    Derivation: `baseline_rl` 0.0 with pretrained 8.0 gives scale 8.0, so
    the bars are +0.4 and -0.8, not +0.0 and -0.0.
    """
    assert metric_scale(0.0, 8.0) == 8.0
    assert metric_scale(10.0, 8.0) == 10.0
    assert metric_scale(-3.0, 1.0) == 3.0

    assert verdict(0.39, 0.0, 8.0) == "NEUTRAL"
    assert verdict(0.41, 0.0, 8.0) == "IMPROVEMENT"
    assert verdict(-0.79, 0.0, 8.0) == "NEUTRAL"
    assert verdict(-0.81, 0.0, 8.0) == "COLLAPSE"

    assert verdict(0.0, 0.0, 0.0) == "NEUTRAL"
    assert verdict(1.0, 0.0, 0.0) == "NEUTRAL"


def test_the_reference_arm_falls_back_to_the_pretrained_score():
    """A suite run without `baseline_rl` has no reference arm, so the
    pretrained score stands in and every delta is measured from it."""
    assert baseline_rl_score_of({"baseline_rl": {"score": 0.7}}, 0.5) == 0.7
    assert baseline_rl_score_of({"kl_penalty": {"score": 0.6}}, 0.5) == 0.5


# ---------------------------------------------------------------------------
# Hypothesis attribution: the evidence set and the recommendation must agree
# (shared with the sibling repo, character for character; S7-9, decided
# 2026-08-18)
# ---------------------------------------------------------------------------

# The six groups and their evidence sets, pinned to the same literal in both
# repos. `analysis/report.py` carried no test at all until now, and the two
# dicts drifting apart would put different numbers under one heading in a
# cross-repo table. Update both files together or not at all.
_EXPECTED_EVIDENCE_SETS = {
    "Catastrophic Forgetting": [
        "ewc", "frozen_backbone", "head_only", "kl_penalty", "llrd", "lora",
    ],
    "Gradient Conflict": ["gradient_surgery", "kl_penalty", "low_t"],
    "Signal Sparsity": [
        "bc_wins", "reward_filtering", "reward_model", "running_stats",
    ],
    "Distributional Shift": ["action_diversity", "mixed_replay"],
    "Mode Collapse": ["advantage_clip", "entropy_bonus", "normalized_adv"],
    "t-Bias": ["low_t", "t_curriculum"],
}


def _named_in(text: str, arm: str) -> bool:
    """Does *text* name *arm* by its registry name, in prose?

    Registry keys are snake_case and the recommendations write them as prose,
    so the separator is relaxed to space, underscore or hyphen: `low_t`
    appears as "low-t" and `entropy_bonus` as "entropy bonus". The word
    bounds are what keep this honest -- a bare substring test would find
    `ewc` inside any word containing those letters, and matching on the
    relaxed separator alone would miss the hyphenated forms entirely.
    """
    pattern = r"\b" + r"[ _\-]".join(re.escape(part) for part in arm.split("_")) + r"\b"
    return re.search(pattern, text, re.IGNORECASE) is not None


def test_every_arm_a_recommendation_names_is_in_its_own_evidence_set():
    """A hypothesis may not recommend an intervention whose ablation it
    excludes from the evidence that scores it.

    `Catastrophic Forgetting` recommended LoRA -- "or use LoRA to restrict
    the parameter update space" -- while omitting the `lora` arm from its
    `supporting_ablations`, in both repos identically. Not cosmetic:
    `_score_hypothesis` computes `evidence_score = n_supporting /
    max(n_tested, 1)` over that list, so the omission changes the ranking
    `diagnosis.md` and the hypothesis-verdict tables print. Author decision
    2026-08-18: drift, not scoping.

    Only recommendations that name a registered arm are constrained. That
    eight of the 25 arms are cited by no hypothesis at all is a separate,
    deliberately open question and is not asserted here.
    """
    offenders = {
        name: sorted(
            arm
            for arm in REGISTRY
            if _named_in(info["recommendation"], arm)
            and arm not in info["supporting_ablations"]
        )
        for name, info in _HYPOTHESIS_GROUPS.items()
    }
    offenders = {name: arms for name, arms in offenders.items() if arms}

    assert not offenders, (
        "hypotheses recommending an intervention whose arm they leave out of "
        f"their own evidence set: {offenders}"
    )


def test_the_hypothesis_evidence_sets_are_the_pinned_shared_ones():
    """The groups and their membership are identical across the two repos.

    Nothing else pins `_HYPOTHESIS_GROUPS`, and it is the input to every
    number in `diagnosis.md`'s hypothesis ranking, so silent drift here is
    invisible until two repos disagree in one table.
    """
    actual = {
        name: sorted(info["supporting_ablations"])
        for name, info in _HYPOTHESIS_GROUPS.items()
    }

    assert actual == _EXPECTED_EVIDENCE_SETS


def test_the_evidence_margin_is_a_fraction_of_the_metric_scale():
    """An arm counts as evidence for a hypothesis when it clears the larger
    reference by a fraction of the metric scale, not by a flat 0.01
    (spec-ablations §3.7; the same scaling as :func:`verdict`).

    Both repos demanded an absolute +0.01 over `max(pretrained, baseline)`
    on metrics of different magnitude: 0.4 % of a 2.6 Craftax achievement
    score against 1.5 % of a 0.65 MiniHack win rate, so the same margin
    asked nearly four times the relative improvement on one side. That is
    the defect the verdict rule had, and it is fixed the same way.

    Derivation. At a baseline of 2.0 the scale is 2.0, so the margin is
    0.02 and the threshold 2.02: an arm at 2.015 is a 0.75 % gain and is
    not evidence, though the flat rule's 2.01 threshold would have counted
    it. At a baseline of 0.2 the scale is 0.2, the margin 0.002 and the
    threshold 0.202: an arm at 0.205 is a 2.5 % gain and is evidence,
    though the flat rule's 0.21 threshold would have refused it. The two
    cases move in opposite directions, which is what an unscaled margin on
    two different metrics does.

    With no scale to measure against, nothing supports anything, which is
    what the verdict rule calls NEUTRAL.
    """
    hyp = {
        "supporting_ablations": ["kl_penalty", "ewc"],
        "description": "d",
        "recommendation": "r",
    }

    def _n(baseline, arm_scores):
        results = {"baseline_rl": {"score": baseline}}
        for name, s in zip(("kl_penalty", "ewc"), arm_scores, strict=True):
            results[name] = {"score": s}
        return _score_hypothesis("h", hyp, results, pretrained_score=baseline)

    # Scale 2.0 -> threshold 2.02. The flat rule's threshold was 2.01.
    assert _n(2.0, [2.015, 2.025])["n_supporting"] == 1
    assert _n(2.0, [2.025, 2.5])["n_supporting"] == 2

    # Scale 0.2 -> threshold 0.202. The flat rule's threshold was 0.21.
    assert _n(0.2, [0.205, 0.2015])["n_supporting"] == 1
    assert _n(0.2, [0.205, 0.203])["n_supporting"] == 2

    # No scale, no evidence.
    assert _n(0.0, [1.0, 1.0])["n_supporting"] == 0
    assert _n(0.0, [1.0, 1.0])["evidence_score"] == 0.0
    assert _n(0.0, [1.0, 1.0])["n_tested"] == 2


def test_the_evidence_score_is_the_raw_quotient_in_both_repos():
    """`evidence_score` is the unrounded fraction; rounding is for display.

    minihack returned `round(evidence, 3)` and craftax the raw quotient, so
    the same inputs gave 0.3330 and 0.3333 under one field name. Every
    consumer already formats at the point of use -- `:.0%` in the report
    tables and `int(score * 5)` for the star rating -- so rounding inside the
    scorer bought nothing and cost cross-repo agreement.
    """
    results = {
        "baseline_rl": {"score": 10.0},
        "kl_penalty": {"score": 10.5},
        "ewc": {"score": 10.5},
        "llrd": {"score": 9.0},
        "lora": {"score": 9.0},
        "frozen_backbone": {"score": 9.0},
        "head_only": {"score": 9.0},
    }
    scored = _score_hypothesis(
        "Catastrophic Forgetting",
        _HYPOTHESIS_GROUPS["Catastrophic Forgetting"],
        results,
        8.0,
    )

    assert scored["n_supporting"] == 2
    assert scored["n_tested"] == 6
    assert scored["evidence_score"] == 2 / 6


def test_an_unregistered_supporting_arm_is_an_error():
    """A typo'd or retired arm name must not be scored as a smaller sample.

    `_score_hypothesis` skips arms absent from `results`, which is correct for
    a run that did not include them -- and indistinguishable from a name that
    can never appear. Left unguarded, renaming an arm silently lowers
    `n_tested` and moves every evidence score that cites it.
    """
    broken = dict(_HYPOTHESIS_GROUPS["Catastrophic Forgetting"])
    broken["supporting_ablations"] = ["ewc", "not_an_ablation"]

    with pytest.raises(KeyError, match="not_an_ablation"):
        _score_hypothesis("Catastrophic Forgetting", broken, {}, 0.0)