"""Sampling tests: generation shape, early stop, greedy determinism.

Uses a tiny untrained model (output is gibberish — we only check mechanics).
Run: `.venv/bin/python tests/test_sample.py`
"""

import sys
from pathlib import Path
from types import SimpleNamespace

import torch

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

from model import GPT  # noqa: E402
from sample import generate, sample_text  # noqa: E402
from tokenizer import BPETokenizer  # noqa: E402

EOT = "<|endoftext|>"


def _tiny_model():
    cfg = SimpleNamespace(vocab_size=300, context_length=32, n_layers=2,
                          n_heads=2, d_model=32, d_ff=64)
    torch.manual_seed(0)
    return GPT(cfg).eval()


def _tiny_tokenizer():
    corpus = ("The cat sat. Tom ran. Lily played happily.\n" + EOT + "\n") * 30
    return BPETokenizer.train(corpus, vocab_size=300, special_tokens=[EOT])


def test_generate_appends_expected_length():
    model = _tiny_model()
    idx = torch.zeros((1, 3), dtype=torch.long)
    out = generate(model, idx, max_new_tokens=10, temperature=0.8, top_k=50)
    # Without an early stop, output grows by exactly max_new_tokens.
    assert out.shape[1] == 3 + 10


def test_generate_respects_context_window():
    model = _tiny_model()  # context_length = 32
    idx = torch.zeros((1, 40), dtype=torch.long)  # longer than the window
    out = generate(model, idx, max_new_tokens=5, temperature=0.8, top_k=50)
    assert out.shape[1] == 40 + 5  # must not error on over-long context


def test_generate_stops_at_stop_id():
    model = _tiny_model()
    idx = torch.zeros((1, 2), dtype=torch.long)
    # Greedy with stop_id equal to whatever greedy picks first -> stops immediately.
    with torch.no_grad():
        first = model(idx)[0][:, -1, :].argmax().item()
    out = generate(model, idx, max_new_tokens=20, temperature=0.0, stop_id=first)
    assert out.shape[1] == 3  # original 2 + the single stop token


def test_greedy_is_deterministic():
    model = _tiny_model()
    idx = torch.zeros((1, 4), dtype=torch.long)
    a = generate(model, idx, max_new_tokens=8, temperature=0.0)
    b = generate(model, idx, max_new_tokens=8, temperature=0.0)
    assert torch.equal(a, b)


def test_sample_text_returns_str_without_eot():
    model = _tiny_model()
    tok = _tiny_tokenizer()
    text = sample_text(model, tok, "The cat", "cpu",
                       max_new_tokens=15, temperature=0.8, top_k=50)
    assert isinstance(text, str)
    assert EOT not in text  # the special token is stripped from output


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