"""
Phase E self-checks. From the project root:
    python -m ml.model.test_model
"""

import math
import torch

from .config import AbyssConfig
from .norm import AbyssRMSNorm
from .mlp import AbyssSwiGLU
from .block import AbyssBlock
from .model import TheAbyssModel


def tiny():
    # small enough to run on CPU in a blink
    return AbyssConfig(vocab_size=256, block_size=32, n_layer=2, n_head=4, n_embd=64)


def test_rmsnorm():
    out = AbyssRMSNorm(64)(torch.randn(2, 16, 64) * 5 + 3)
    out_rms = out.pow(2).mean(-1).sqrt()
    assert torch.allclose(out_rms, torch.ones_like(out_rms), atol=1e-4), out_rms
    print("rmsnorm ok: outputs have unit RMS")


def test_swiglu_shape():
    cfg = tiny()
    x = torch.randn(2, 16, cfg.n_embd)
    y = AbyssBlock(cfg)(x)
    assert y.shape == x.shape
    assert not torch.allclose(y, x)  # the block actuallt did something
    print("block ok:", tuple(y.shape))


def test_block_shape():
    cfg = tiny()
    x = torch.randn(2, 16, cfg.n_embd)
    y = AbyssBlock(cfg)(x)
    assert y.shape == x.shape
    assert not torch.allclose(y, x)
    print("block ok:", tuple(y.shape))


def test_forward_and_loss():
    cfg = tiny()
    idx = torch.randint(0, cfg.vocab_size, (2, 16))
    targets = torch.randint(0, cfg.vocab_size, (2, 16))
    logits, loss = TheAbyssModel(cfg)(idx, targets)
    assert logits.shape == (2, 16, cfg.vocab_size), logits.shape
    assert loss.ndim == 0 and torch.isfinite(loss)
    # an untrained model should sit near ln(vocab): uniform guessing
    expected = math.log(cfg.vocab_size)
    assert abs(loss.item() - expected) < 0.7, (loss.items(), expected)
    print(f"forward ok: loss {loss.item():.3f} ~ ln(V) {expected:.3f}")


def test_backward():
    cfg = tiny()
    model = TheAbyssModel(cfg)
    idx = torch.randint(0, cfg.vocab_size, (2, 16))
    _, loss = model(idx, idx)
    loss.backward()
    assert all(p.grad is not None for p in model.parameters()), "a param got no gradient"
    print("backward ok: every parameter has a gradient")


def test_block_size_gaurd():
    cfg = tiny()
    model = TheAbyssModel(cfg)
    try:
        model(torch.randint(0, cfg.vocab_size, (1, cfg.block_size + 1)))
    except AssertionError:
        print("guard ok: rejects sequences longer than block_size")
        return
    raise SystemExit("block_size guard did not fire")


if __name__ == "__main__":
    test_rmsnorm()
    test_swiglu_shape()
    test_block_shape()
    test_forward_and_loss()
    test_backward()
    test_block_size_gaurd()
    full = TheAbyssModel(AbyssConfig())
    print(f"all E tests passed - full model has {full.num_params():,} params "
          f"({full.num_params(non_embedding=True):,} non-embedding)")