"""Guardrail: each degradation config changes exactly one thing from v1.

Enforces the controlled-experiment discipline — if a variant accidentally
drifts from the baseline in an unintended field, this fails. Forward-only.
Run: `.venv/bin/python tests/test_configs.py`
"""

import importlib
import sys
from dataclasses import asdict
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "src"))
sys.path.insert(0, str(ROOT / "configs"))

from model import GPT  # noqa: E402

# For each variant: which model fields and which train fields may differ from v1.
EXPECTED_DIFFS = {
    "small": ({"n_layers", "d_model", "n_heads", "d_ff"}, set()),
    "short": (set(), {"max_iters"}),
    "ctx": ({"context_length"}, set()),
    "data": (set(), {"max_train_tokens"}),
}


def _diff_keys(a, b):
    return {k for k in a if a[k] != b[k]}


def test_each_variant_changes_only_intended_fields():
    v1 = importlib.import_module("v1")
    base_model, base_train = asdict(v1.model), asdict(v1.train)
    for name, (exp_model, exp_train) in EXPECTED_DIFFS.items():
        cfg = importlib.import_module(name)
        model_diffs = _diff_keys(asdict(cfg.model), base_model)
        train_diffs = _diff_keys(asdict(cfg.train), base_train)
        assert model_diffs == exp_model, f"{name}: model diffs {model_diffs} != {exp_model}"
        assert train_diffs == exp_train, f"{name}: train diffs {train_diffs} != {exp_train}"


def test_each_variant_builds_a_model():
    for name in EXPECTED_DIFFS:
        cfg = importlib.import_module(name)
        model = GPT(cfg.model)
        assert model.num_params() > 0


def test_small_param_count():
    cfg = importlib.import_module("small")
    model = GPT(cfg.model)
    # tok_emb 4096*64 + pos_emb 256*64
    # + 2*(qkv 64*192 + proj 64*64 + 2 norms*64 + mlp_fc 64*256 + mlp_proj 256*64)
    # + final norm 64
    expected = 4096 * 64 + 256 * 64 + 2 * (64 * 192 + 64 * 64 + 2 * 64 + 64 * 256 + 256 * 64) + 64
    assert expected == 377_152
    assert model.num_params() == expected
    # The embedding table dominates: non-embedding capacity is tiny.
    assert model.num_params(non_embedding=True) == 98_624


def main() -> None:
    tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
    for t in tests:
        t()
        print(f"ok  {t.__name__}")
    print(f"\n{len(tests)} passed")


if __name__ == "__main__":
    main()