arborist/tests/test_t3_bound_calculator.py

"""Tests for the T3 per-window covert-channel bound calculator
(`bench/scripts/t3_bound_calculator.py`, ticket #000036 §11).

Validates the closed-form bound from `docs/soft-hash-channel-t3-
bound.md` §6 against:

  - mathematical identities (B1 / B2 / B3 isolation, monotonicity in
    each input)
  - input-validation hard checks
  - operator-guidance text mode transitions (≤0 / <256 / ≥256 bits)
  - CLI surface (argparse + JSON output, error path)
  - the §11 worked-example numbers in the ticket doc

Pure stdlib — same dep profile as the calculator itself.
"""

from __future__ import annotations

import json
import math
import pathlib
import subprocess
import sys

import pytest

from bench.scripts.t3_bound_calculator import (
    CALCULATOR_VERSION,
    t3_bound_bits,
)


# --- baseline (§11 worked example) -----------------------------------


_BASELINE = dict(
    gradient_fraction=0.05,
    gradient_norm_max=1.0,
    gradient_noise_stddev=0.1,
    lr_decision_interval=100,
    lr_grid_size=8,
    window_length=10000,
    batches_per_epoch=1024,
    steps_per_epoch=1024,
)


def test_baseline_matches_section_11_doc():
    """The §11 worked example baseline. Calculator's actual numbers:

        I_window_bits_upper_bound: 625.8716
        B1_contribution:           292.4813
        B2_contribution:           300.0
        B3_contribution:            33.3904
        snr_grad:                    0.5

    Doc §11's quoted 622.7 / 290.0 / 32.7 are first-cut rounded
    estimates; calculator's actual closed-form values differ by
    ~3 bits total due to the doc's 290 vs actual 292.5 in B1 and
    32.7 vs actual 33.4 in B3. This test pins the closed-form
    truth (the doc text needs a calibration pass; tracked in this
    test's docstring as a follow-up).
    """
    r = t3_bound_bits(**_BASELINE)
    # Total: ⟨B1+B2+B3⟩ ≈ 292.48 + 300.0 + 33.39 = 625.87
    assert r["I_window_bits_upper_bound"] == pytest.approx(625.87, abs=0.05)
    # B2 is exact: ⌈10000/100⌉ · log₂(8) = 100 · 3 = 300
    assert r["B2_contribution"] == pytest.approx(300.0, abs=1e-3)
    # B3: epochs = ⌈10000/1024⌉ = 10; per-epoch factor = 1024·0.1/1.0 = 102.4;
    # log₂(102.4)/2 ≈ 3.339; total ≈ 33.39 bits.
    assert r["B3_contribution"] == pytest.approx(33.39, abs=0.05)
    # B1: C_B1 · g · W · log₂(SNR+1) where SNR = 0.05·1.0/0.1 = 0.5
    # → log₂(1.5) ≈ 0.585; total ≈ 1.0 · 0.05 · 10000 · 0.585 = 292.48
    assert r["B1_contribution"] == pytest.approx(292.48, abs=0.05)


def test_returns_calculator_version_token():
    r = t3_bound_bits(**_BASELINE)
    assert r["calculator_version"] == CALCULATOR_VERSION
    # versioned-default discipline: matches the v1-bottou-refinement token
    assert "v1" in CALCULATOR_VERSION


def test_b2_isolated_zero_when_decisions_zero():
    """If lr_decision_interval > window_length, decisions_in_window = 1
    (one initial decision); the bound floor is log₂ R, not zero.
    Exercises the ceil semantics on small W / large K.
    """
    r = t3_bound_bits(**{**_BASELINE, "window_length": 1, "lr_decision_interval": 100})
    # ceil(1/100) = 1, so B2 = log₂ 8 = 3.0
    assert r["B2_contribution"] == pytest.approx(3.0, abs=1e-6)


def test_snr_grad_formula():
    """SNR_grad = g · ‖∇L_max‖ / σ_grad. Check directly."""
    r = t3_bound_bits(**_BASELINE)
    expected_snr = 0.05 * 1.0 / 0.1  # 0.5
    assert r["snr_grad"] == pytest.approx(expected_snr, abs=1e-6)


# --- monotonicity / bound shape --------------------------------------


def test_monotone_in_window_length():
    """Doubling W roughly doubles each contribution (B1 linear in W,
    B2 + B3 grow with ⌈W/K⌉ + ⌈W/E⌉ which scale linearly with W on
    aligned periods)."""
    base = t3_bound_bits(**_BASELINE)
    long_window = t3_bound_bits(**{**_BASELINE, "window_length": 20000})
    # Total grows monotonically.
    assert long_window["I_window_bits_upper_bound"] > base["I_window_bits_upper_bound"]
    # B1 is exactly linear in W, so doubling W doubles B1.
    assert long_window["B1_contribution"] == pytest.approx(
        2 * base["B1_contribution"], rel=1e-6
    )


def test_monotone_in_gradient_fraction():
    """Increasing g raises the bound (more adversarial control)."""
    base = t3_bound_bits(**_BASELINE)
    high_g = t3_bound_bits(**{**_BASELINE, "gradient_fraction": 0.5})
    assert high_g["I_window_bits_upper_bound"] > base["I_window_bits_upper_bound"]
    # B2 + B3 do not depend on g, so the increase is in B1 only.
    assert high_g["B2_contribution"] == base["B2_contribution"]
    assert high_g["B3_contribution"] == base["B3_contribution"]


def test_b3_floor_at_zero_when_factor_below_one():
    """When N_b · σ_grad / ‖∇L_max‖ ≤ 1, log₂ goes ≤ 0; the bound
    floors at 0 — adversary can't do worse than random shuffle in
    this regime. Verify the floor."""
    # σ_grad = 0.001, gradient_norm_max = 100, N_b = 4 → 0.00004 < 1
    args = {**_BASELINE, "gradient_noise_stddev": 0.001,
            "gradient_norm_max": 100.0, "batches_per_epoch": 4,
            "steps_per_epoch": 4}
    r = t3_bound_bits(**args)
    assert r["B3_contribution"] == pytest.approx(0.0, abs=1e-9)


def test_constant_C_B1_scales_b1():
    """C_B1 · B1 — halving C_B1 halves B1 (calculator rounds output
    to 4 dp; tolerate rounding)."""
    base = t3_bound_bits(**_BASELINE)
    half_c = t3_bound_bits(**_BASELINE, c_b1=0.5)
    assert half_c["B1_contribution"] == pytest.approx(
        0.5 * base["B1_contribution"], abs=1e-3
    )
    # B2 / B3 unchanged.
    assert half_c["B2_contribution"] == base["B2_contribution"]
    assert half_c["B3_contribution"] == base["B3_contribution"]


def test_constants_recorded():
    r = t3_bound_bits(**_BASELINE, c_b1=0.4, c_b2=0.3, c_b3=0.2)
    assert r["constants"] == {"C_B1": 0.4, "C_B2": 0.3, "C_B3": 0.2}


# --- input validation ------------------------------------------------


@pytest.mark.parametrize("g", [0.0, -0.1, 1.5, 2.0])
def test_invalid_gradient_fraction_rejected(g):
    args = {**_BASELINE, "gradient_fraction": g}
    with pytest.raises(ValueError, match="gradient_fraction"):
        t3_bound_bits(**args)


@pytest.mark.parametrize("name,value", [
    ("gradient_norm_max", 0.0),
    ("gradient_norm_max", -1.0),
    ("gradient_noise_stddev", 0.0),
    ("gradient_noise_stddev", -0.1),
])
def test_invalid_positive_floats_rejected(name, value):
    args = {**_BASELINE, name: value}
    with pytest.raises(ValueError, match=name):
        t3_bound_bits(**args)


@pytest.mark.parametrize("name", [
    "lr_decision_interval", "lr_grid_size", "window_length",
    "batches_per_epoch", "steps_per_epoch",
])
@pytest.mark.parametrize("bad", [0, -1, 1.5, "a"])
def test_invalid_positive_int_rejected(name, bad):
    args = {**_BASELINE, name: bad}
    with pytest.raises(ValueError, match=name):
        t3_bound_bits(**args)


@pytest.mark.parametrize("c_name,c_val", [
    ("c_b1", -0.1), ("c_b1", 1.5),
    ("c_b2", -0.5), ("c_b2", 2.0),
    ("c_b3", -1.0), ("c_b3", 1.001),
])
def test_constants_clamped_to_unit_interval(c_name, c_val):
    args = {**_BASELINE, c_name: c_val}
    with pytest.raises(ValueError, match=c_name):
        t3_bound_bits(**args)


# --- recommendation text mode transitions ----------------------------


def test_recommendation_total_le_zero():
    """When the bound floors at 0 across all three terms (extreme
    benign config), recommendation reads as 'vacuously safe'."""
    args = {
        **_BASELINE,
        "gradient_fraction": 1e-9,           # near-zero g → tiny B1
        "lr_decision_interval": 1_000_000,    # huge K → ⌈W/K⌉ = 1, but log₂ R = log₂ 1 = 0 below
        "lr_grid_size": 1,                    # log₂(1) = 0 → B2 = 0
        "gradient_noise_stddev": 0.001,
        "gradient_norm_max": 100.0,
        "batches_per_epoch": 4,
        "steps_per_epoch": 4,                 # b3_factor = 4·0.001/100 = 4e-5 < 1 → B3 floor
    }
    r = t3_bound_bits(**args)
    assert r["I_window_bits_upper_bound"] >= 0
    # B1 is tiny but not exactly 0; total floor regime depends on
    # whether the sum reaches positive territory. Assert at least
    # B2 + B3 are zero on this config.
    assert r["B2_contribution"] == pytest.approx(0.0, abs=1e-9)
    assert r["B3_contribution"] == pytest.approx(0.0, abs=1e-9)


def test_recommendation_below_sha256():
    """Baseline (~622 bits) actually exceeds 256 — wrong band for
    this test. Use a config that lands in the (0, 256) range."""
    args = {**_BASELINE, "window_length": 1000}
    r = t3_bound_bits(**args)
    assert 0 < r["I_window_bits_upper_bound"] < 256, (
        f"window_length=1000 should land in (0, 256); got {r['I_window_bits_upper_bound']}"
    )
    assert "windows" in r["recommendation"]
    assert "256" in r["recommendation"]


def test_recommendation_exceeds_sha256():
    """Baseline (W=10000) → 622 bits > 256 — recommendation should
    flag M2's single-window guarantee as broken."""
    r = t3_bound_bits(**_BASELINE)
    assert r["I_window_bits_upper_bound"] > 256
    assert "EXCEEDS" in r["recommendation"]
    assert "256" in r["recommendation"]


# --- CLI surface (subprocess invocation) -----------------------------


def test_cli_baseline_runs_clean(tmp_path):
    """Invoke the calculator CLI with §11 worked-example flags;
    parse the JSON output; verify the I_window_bits_upper_bound
    matches the function call."""
    cmd = [
        sys.executable, "-m", "bench.scripts.t3_bound_calculator",
        "--gradient-fraction", "0.05",
        "--gradient-norm-max", "1.0",
        "--gradient-noise-stddev", "0.1",
        "--lr-decision-interval", "100",
        "--lr-grid-size", "8",
        "--window-length", "10000",
        "--batches-per-epoch", "1024",
        "--steps-per-epoch", "1024",
    ]
    out = subprocess.run(cmd, capture_output=True, text=True, check=True,
                         cwd="/home/fox/git/arborist")
    j = json.loads(out.stdout)
    assert j["calculator_version"] == CALCULATOR_VERSION
    # Closed-form actual; doc §11's 622.7 is approximate.
    assert j["I_window_bits_upper_bound"] == pytest.approx(625.87, abs=0.05)


def test_cli_invalid_input_exits_2():
    """ValueError → exit code 2 + stderr error message."""
    cmd = [
        sys.executable, "-m", "bench.scripts.t3_bound_calculator",
        "--gradient-fraction", "1.5",            # invalid
        "--gradient-norm-max", "1.0",
        "--gradient-noise-stddev", "0.1",
        "--lr-decision-interval", "100",
        "--lr-grid-size", "8",
        "--window-length", "100",
        "--batches-per-epoch", "1024",
        "--steps-per-epoch", "1024",
    ]
    out = subprocess.run(cmd, capture_output=True, text=True,
                         cwd="/home/fox/git/arborist")
    assert out.returncode == 2
    assert "gradient_fraction" in out.stderr


# --- closed-form sanity ---------------------------------------------


def test_b2_exact_formula():
    """B2 = C_B2 · ⌈W/K⌉ · log₂(R). Direct closed-form check
    against several configurations."""
    cases = [
        # (W, K, R) → expected ⌈W/K⌉ · log₂ R
        (1000, 100, 8, 10 * 3.0),
        (1000, 200, 4, 5 * 2.0),
        (999, 100, 16, 10 * 4.0),     # ceil(999/100)=10
        (1, 100, 2, 1 * 1.0),
    ]
    for W, K, R, expected in cases:
        r = t3_bound_bits(
            **{**_BASELINE, "window_length": W,
               "lr_decision_interval": K, "lr_grid_size": R},
        )
        assert r["B2_contribution"] == pytest.approx(expected, abs=1e-9), (
            f"B2 mismatch on (W={W}, K={K}, R={R}): "
            f"{r['B2_contribution']} vs expected {expected}"
        )


def test_b1_exact_formula():
    """B1 = C_B1 · g · W · log₂(SNR + 1) where SNR = g·‖∇L_max‖/σ_grad.
    Hand-compute for one config and check exact agreement."""
    g, gnm, sigma, W = 0.1, 2.0, 0.5, 5000
    snr = g * gnm / sigma  # 0.4
    expected = 1.0 * g * W * math.log2(snr + 1)
    r = t3_bound_bits(**{
        **_BASELINE,
        "gradient_fraction": g,
        "gradient_norm_max": gnm,
        "gradient_noise_stddev": sigma,
        "window_length": W,
    })
    # Output rounded to 4 dp inside the calculator.
    assert r["B1_contribution"] == pytest.approx(expected, abs=1e-3)
    assert r["snr_grad"] == pytest.approx(snr, abs=1e-3)


# --- integration with §6 framework -----------------------------------


def test_total_equals_sum_of_three_contributions():
    """Closure check: I_window ≡ B1 + B2 + B3 (no missing term, no
    double-counting). Multiple inputs to widen the cone."""
    for window_length, lr_decision_interval in [
        (100, 10),
        (5000, 250),
        (10000, 100),
        (50000, 1000),
    ]:
        args = {**_BASELINE, "window_length": window_length,
                "lr_decision_interval": lr_decision_interval}
        r = t3_bound_bits(**args)
        total = r["B1_contribution"] + r["B2_contribution"] + r["B3_contribution"]
        # Allow small rounding (round to 4 dp on each contribution).
        assert r["I_window_bits_upper_bound"] == pytest.approx(total, abs=1e-3)


# --- B3 hand-computed formula (gap caught by the docs/calculator-test-
# patterns.md checklist, 2026-05-10) ---------------------------------


def test_b3_exact_formula():
    """B3 = C_B3 · ⌈W/E⌉ · log₂(N_b · σ_grad / ‖∇L_max‖) / 2 per
    #000036 §5 (Bottou-Bousquet refinement).

    Pairs the existing test_b1_exact_formula + test_b2_exact_formula
    by hand-computing B3 from the spec formula. Closes the
    calculator-test-patterns.md §3 checklist gap (item 3:
    hand-formula assertions per output field).
    """
    # Pick a config where the B3 factor (N_b · σ / ‖∇L_max‖) > 1
    # so the floor doesn't kick in — that case is covered separately
    # by test_b3_floor_at_zero_when_factor_below_one.
    g, gnm, sigma, W = 0.05, 1.0, 0.5, 4096
    K, R = 100, 8
    N_b, E = 512, 512
    factor = N_b * sigma / gnm  # 256.0 > 1
    epochs = -(-W // E)         # ceil
    expected_b3 = 1.0 * epochs * math.log2(factor) / 2.0
    r = t3_bound_bits(
        gradient_fraction=g, gradient_norm_max=gnm,
        gradient_noise_stddev=sigma, lr_decision_interval=K,
        lr_grid_size=R, window_length=W,
        batches_per_epoch=N_b, steps_per_epoch=E,
    )
    assert r["B3_contribution"] == pytest.approx(expected_b3, abs=1e-3), (
        f"B3 mismatch: function returned {r['B3_contribution']}, "
        f"hand-computed {expected_b3} via "
        f"⌈{W}/{E}⌉ · log₂({factor})/2 = {epochs} · {math.log2(factor):.4f}/2"
    )


# --- KAT regression (gap caught by checklist, 2026-05-10) -----------


KAT_FIXTURE = (
    pathlib.Path(__file__).parent.parent
    / "bench" / "fixtures" / "t3-bound"
    / "known-answer-tests.jsonl"
)


@pytest.mark.skipif(
    not KAT_FIXTURE.exists(),
    reason="KAT fixture not yet generated",
)
def test_t3_bound_known_answer_tests():
    """Replay the pinned t3-bound/known-answer-tests.jsonl fixture;
    every entry's per-contribution + total must match what the
    function produces today.

    Algorithm change MUST bump ``CALCULATOR_VERSION`` and emit a new
    fixture file under ``bench/fixtures/t3-bound/`` — old runs
    replay against old data per the calculator-test-patterns.md §1
    discipline.

    Closes the calculator-test-patterns.md §1 checklist gap (no
    KAT fixture for this calculator).
    """
    n_kats = 0
    for line in KAT_FIXTURE.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        kat = json.loads(line)
        n_kats += 1
        r = t3_bound_bits(**kat["inputs"])
        assert r["calculator_version"] == kat["calculator_version"], (
            f"KAT version mismatch on {kat['label']!r}: "
            f"recorded {kat['calculator_version']}, "
            f"function {r['calculator_version']}"
        )
        for field, key in [
            ("expected_total", "I_window_bits_upper_bound"),
            ("expected_b1", "B1_contribution"),
            ("expected_b2", "B2_contribution"),
            ("expected_b3", "B3_contribution"),
            ("expected_snr_grad", "snr_grad"),
        ]:
            recorded = float(kat[field])
            observed = float(r[key])
            assert observed == pytest.approx(recorded, abs=1e-3), (
                f"KAT {field} mismatch on {kat['label']!r}: "
                f"recorded {recorded}, observed {observed}"
            )
    assert n_kats >= 5, (
        f"KAT fixture seems incomplete: only {n_kats} entries; "
        "calculator-test-patterns.md §1 wants ≥ 5"
    )