"""Unit tests for multi-token prediction (Meta MTP) support.

Run: .venv/bin/python tests/test_mtp.py
"""
import subprocess
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import torch
import torch.nn.functional as F

from model.config import TinyLiquidConfig, CONFIGS
from model.tiny_liquid import TinyLiquid

ROOT = Path(__file__).resolve().parents[1]


def make_model(vocab=512, mtp=2):
    cfg = TinyLiquidConfig(vocab_size=vocab, mtp_heads=mtp, **CONFIGS["micro6m"])
    return TinyLiquid(cfg)


def test_forward_mtp_shapes():
    m = make_model()
    ids = torch.randint(0, 512, (2, 16))
    logits, aux = m.forward_mtp(ids)
    assert tuple(logits.shape) == (2, 16, 512)
    assert len(aux) == 2 and all(tuple(a.shape) == (2, 16, 512) for a in aux)
    # main forward unchanged
    assert tuple(m(ids).shape) == (2, 16, 512)


def test_mtp_loss_backward():
    m = make_model()
    ids = torch.randint(0, 512, (2, 24))
    logits, aux = m.forward_mtp(ids)
    loss = F.cross_entropy(logits.reshape(-1, 512), ids.reshape(-1))
    for k, a in enumerate(aux):
        off = k + 2
        loss = loss + 0.1 * F.cross_entropy(
            a[:, :-off].reshape(-1, 512), ids[:, off:].reshape(-1))
    loss.backward()
    assert m.mtp_heads[0][0].weight.grad is not None
    assert m.tok_emb.weight.grad is not None
    assert torch.isfinite(loss)


def test_mtp_checkpoint_roundtrip():
    m = make_model()
    sd = {"config": m.cfg.__dict__, "model": m.state_dict()}
    m2 = TinyLiquid(TinyLiquidConfig(**sd["config"]))
    missing, unexpected = m2.load_state_dict(sd["model"], strict=True)
    assert not missing and not unexpected


def test_train_lm_mtp_smoke(tmp=None):
    tmp = Path(tmp or (ROOT / "data" / "_mtp_smoke"))
    tmp.mkdir(parents=True, exist_ok=True)
    import numpy as np
    rng = np.random.default_rng(0)
    (tmp / "train.bin").write_bytes(rng.integers(1, 500, size=20000, dtype=np.uint16).tobytes())
    (tmp / "valid.bin").write_bytes(rng.integers(1, 500, size=5000, dtype=np.uint16).tobytes())
    ckpt = tmp / "ckpt"
    cmd = [
        sys.executable, "-u", "train/train_lm.py",
        "--data", str(tmp / "train.bin"), "--val", str(tmp / "valid.bin"),
        "--tok", "data/tokenizer.json", "--config", "micro6m",
        "--ckpt", str(ckpt), "--batch", "2", "--seq", "32",
        "--lr", "1e-4", "--warmup", "0", "--steps", "3",
        "--eval-every", "2", "--save-every", "2", "--threads", "2",
        "--mtp", "2", "--log-every", "1",
    ]
    env = {"PYTHONPATH": str(ROOT)}
    r = subprocess.run(cmd, capture_output=True, text=True, cwd=ROOT, env=env,
                       timeout=300)
    assert r.returncode == 0, r.stderr[-1500:]
    saved = sorted((ckpt).glob("*.pt"))
    assert saved, "no checkpoint saved"
    import torch as T
    sd = T.load(saved[-1], map_location="cpu", weights_only=False)
    assert sd["config"]["mtp_heads"] == 2
    assert any("mtp_heads" in k for k in sd["model"])
    for f in ["train.bin", "valid.bin"]:
        (tmp / f).unlink()
    for f in ckpt.glob("*.pt"):
        f.unlink()
    ckpt.rmdir()
    tmp.rmdir()


if __name__ == "__main__":
    fns = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
    for fn in fns:
        fn()
        print(f"ok {fn.__name__}")
    print(f"\n{len(fns)} mtp tests passed")