"""Tests for training helpers that don't require an optimization run.

Covers the LR schedule (a pure function) and the optimizer param grouping.
Run: `.venv/bin/python tests/test_train.py`
"""

import sys
from dataclasses import dataclass
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 train import configure_optimizer, cosine_lr  # noqa: E402
from v1 import ModelConfig  # noqa: E402


@dataclass(frozen=True)
class TC:
    learning_rate: float = 1e-3
    min_lr: float = 1e-4
    warmup_iters: int = 100
    max_iters: int = 1000
    weight_decay: float = 0.1
    beta1: float = 0.9
    beta2: float = 0.95


def test_lr_warmup_is_linear():
    tc = TC()
    # Ramps from ~0 up to the peak across warmup_iters.
    assert cosine_lr(0, tc) < cosine_lr(50, tc) < cosine_lr(99, tc)
    assert abs(cosine_lr(tc.warmup_iters - 1, tc) - tc.learning_rate) < 1e-9


def test_lr_peaks_then_decays_to_floor():
    tc = TC()
    peak = cosine_lr(tc.warmup_iters, tc)
    assert abs(peak - tc.learning_rate) < 1e-6
    # Monotonically decreasing through the cosine phase.
    mid = cosine_lr(tc.warmup_iters + (tc.max_iters - tc.warmup_iters) // 2, tc)
    assert tc.min_lr < mid < peak
    # Bottoms out at the floor.
    assert abs(cosine_lr(tc.max_iters, tc) - tc.min_lr) < 1e-9
    assert abs(cosine_lr(tc.max_iters + 500, tc) - tc.min_lr) < 1e-9


def test_optimizer_grouping():
    model = GPT(ModelConfig())
    opt = configure_optimizer(model, TC())
    decay_group, no_decay_group = opt.param_groups
    assert decay_group["weight_decay"] == 0.1
    assert no_decay_group["weight_decay"] == 0.0
    # Norm/bias params (1-D) must be in the no-decay group only.
    assert all(p.dim() >= 2 for p in decay_group["params"])
    assert all(p.dim() < 2 for p in no_decay_group["params"])


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