"""Model tests: parameter count, output shapes, init loss, and causality.

All forward-pass only (no training). Run: `.venv/bin/python tests/test_model.py`
"""

import math
import sys
from pathlib import Path

import torch

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
from v1 import ModelConfig  # noqa: E402


def test_param_count_matches_hand_estimate():
    model = GPT(ModelConfig())
    # tok_emb 4096*256 + pos_emb 256*256
    # + 4 blocks * (qkv 256*768 + attn_proj 256*256 + 2 norms*256
    #               + mlp_fc 256*1024 + mlp_proj 1024*256)
    # + final norm 256   (head is weight-tied, so no extra params)
    expected = (
        4096 * 256 + 256 * 256
        + 4 * (256 * 768 + 256 * 256 + 2 * 256 + 256 * 1024 + 1024 * 256)
        + 256
    )
    assert expected == 4_262_144
    assert model.num_params() == expected


def test_weight_tying():
    model = GPT(ModelConfig())
    # Head and token embedding must be the very same parameter tensor.
    assert model.head.weight is model.tok_emb.weight


def test_forward_shapes():
    cfg = ModelConfig()
    model = GPT(cfg)
    idx = torch.randint(0, cfg.vocab_size, (2, 16))
    logits, loss = model(idx)
    assert logits.shape == (2, 16, cfg.vocab_size)
    assert loss is None
    _, loss = model(idx, targets=idx)
    assert loss.ndim == 0  # scalar


def test_init_loss_near_uniform():
    # A well-initialised model predicts ~uniformly at first, so the
    # cross-entropy loss should be close to ln(vocab_size). Targets are drawn
    # independently of the input: with weight tying, using the input as its own
    # target would leak (each position's residual already holds its embedding).
    cfg = ModelConfig()
    torch.manual_seed(0)
    model = GPT(cfg)
    idx = torch.randint(0, cfg.vocab_size, (8, 64))
    targets = torch.randint(0, cfg.vocab_size, (8, 64))
    _, loss = model(idx, targets=targets)
    assert abs(loss.item() - math.log(cfg.vocab_size)) < 0.5, loss.item()


def test_causality():
    # Changing a token at position t must not affect logits at positions < t.
    cfg = ModelConfig()
    torch.manual_seed(0)
    model = GPT(cfg)
    model.eval()
    idx = torch.randint(0, cfg.vocab_size, (1, 32))
    with torch.no_grad():
        base, _ = model(idx)
        changed = idx.clone()
        changed[0, 20] = (changed[0, 20] + 1) % cfg.vocab_size
        after, _ = model(changed)
    # Positions before the edit are identical; the edited position differs.
    assert torch.allclose(base[0, :20], after[0, :20], atol=1e-5)
    assert not torch.allclose(base[0, 20], after[0, 20])


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()