La receta de entrenamiento

"""Learning-rate schedule: a linear warmup followed by a cosine decay to a floor."""

from __future__ import annotations

import math
from typing import TYPE_CHECKING

if TYPE_CHECKING:  # pragma: no cover - only for type checkers
    from rukh.train.loop import TrainConfig

def lr_at(step: int, cfg: TrainConfig) -> float:
    """Learning rate for ``step`` (0-based).

    It rises linearly from 0 to ``cfg.lr`` over the first ``cfg.warmup`` steps, then follows a
    cosine down to ``cfg.min_lr_ratio * cfg.lr`` at ``cfg.max_steps`` and stays there for any
    later step (so a resumed run past ``max_steps`` never gets a negative or rising rate).
    """
    if step < 0:
        raise ValueError(f"step must be non-negative, got {step}")
    floor = cfg.lr * cfg.min_lr_ratio
    if step < cfg.warmup:
        return cfg.lr * step / cfg.warmup
    span = cfg.max_steps - cfg.warmup
    if span <= 0:
        return floor
    progress = min(1.0, (step - cfg.warmup) / span)
    return floor + 0.5 * (1.0 + math.cos(math.pi * progress)) * (cfg.lr - floor)

"""Tests for rukh.train.schedule: linear warmup then cosine decay to a floor."""

from __future__ import annotations

import pytest

from rukh.train import TrainConfig, lr_at

pytestmark = pytest.mark.unit

CFG = TrainConfig(lr=6e-4, min_lr_ratio=0.1, warmup=100, max_steps=1000)


def test_warmup_is_linear_and_reaches_the_peak() -> None:
    assert lr_at(0, CFG) == 0.0
    assert lr_at(50, CFG) == pytest.approx(CFG.lr / 2)
    assert lr_at(100, CFG) == pytest.approx(CFG.lr)
    rising = [lr_at(step, CFG) for step in range(101)]
    assert rising == sorted(rising)
    deltas = {round(b - a, 12) for a, b in zip(rising, rising[1:], strict=False)}
    assert len(deltas) == 1

def test_cosine_decays_to_the_floor_and_stays_there() -> None:
    floor = CFG.lr * CFG.min_lr_ratio
    assert lr_at(1000, CFG) == pytest.approx(floor)
    assert lr_at(5000, CFG) == pytest.approx(floor)
    falling = [lr_at(step, CFG) for step in range(100, 1001, 10)]
    assert falling == sorted(falling, reverse=True)
    assert all(value >= floor - 1e-12 for value in falling)
    halfway = lr_at(550, CFG)
    assert halfway == pytest.approx(floor + 0.5 * (CFG.lr - floor))


def test_without_warmup_the_first_step_is_the_peak() -> None:
    cfg = CFG.model_copy(update={"warmup": 0})
    assert lr_at(0, cfg) == pytest.approx(cfg.lr)


def test_degenerate_schedule_returns_the_floor() -> None:
    cfg = CFG.model_copy(update={"warmup": 1000, "max_steps": 1000})
    assert lr_at(1000, cfg) == pytest.approx(cfg.lr * cfg.min_lr_ratio)
    with pytest.raises(ValueError, match="non-negative"):
        lr_at(-1, CFG)


def test_config_rejects_unknown_keys() -> None:
    with pytest.raises(ValueError):
        TrainConfig(learning_rate=1e-3)  # type: ignore[call-arg]

"""The training loop: AdamW with a cosine schedule, bf16, checkpoints and MLflow logging.

The recipe follows ``docs/spec/02`` (component 1): AdamW (0.9/0.95, weight decay 0.1 applied
only to matrices), learning rate 6e-4 with 1 000 warmup steps and a cosine decay, an effective
batch of ``batch_size * grad_accum`` sequences, bf16 autocast on CUDA, gradient clipping at 1.0
and optional ``torch.compile``. Everything the run needs to be reproducible (config, seed, git
SHA, vocabulary hash, data manifest hash) goes to MLflow and into every checkpoint.

Resuming is meant to be indistinguishable from never having stopped: the batch stream is wound
forward past the windows the first half of the run already saw, and the MLflow run id travels in
the checkpoint so the curve carries on in the same run instead of starting a second one.

Two throughput numbers are logged because they answer different questions: ``tokens_per_s``
counts every position in the window (what the GPU actually processed, comparable across runs)
and ``real_tokens_per_s`` counts only the non-``<pad>`` targets (what the model learned from).
The training loss is accumulated token-weighted rather than as a mean of means, so a
micro-batch with fewer real tokens does not count as much as a full one.
"""

from __future__ import annotations

import logging
import math
import os
import time
from collections.abc import Iterator
from contextlib import nullcontext
from datetime import UTC, datetime
from pathlib import Path
from typing import Any, Literal

import torch
from torch import nn
from torch.utils.data import DataLoader

from rukh import paths
from rukh.config import BaseConfig
from rukh.models import DecoderConfig, MoveDecoder, preset
from rukh.tokenize.loader import IGNORE_INDEX, PackedDataset, make_loader
from rukh.train.checkpoint import (
    BEST_NAME,
    load_checkpoint,
    read_manifest_sha,
    read_vocab_hash,
    restore,
    save_checkpoint,
    step_name,
)
from rukh.train.schedule import lr_at

log = logging.getLogger(__name__)

Batch = tuple[torch.Tensor, torch.Tensor]

class TrainConfig(BaseConfig):
    """Everything one training run needs; unknown keys in the YAML are an error."""

    preset: Literal["tiny", "small", "medium"] = "small"
    model: DecoderConfig | None = None  # overrides the preset when given
    tokens_dir: str = "data/tokens/uci"
    block: int = 200
    batch_size: int = 64
    grad_accum: int = 4  # 256 effective sequences
    lr: float = 6e-4
    min_lr_ratio: float = 0.1
    warmup: int = 1000
    max_steps: int = 20000
    weight_decay: float = 0.1
    betas: tuple[float, float] = (0.9, 0.95)
    grad_clip: float = 1.0
    precision: Literal["bf16", "fp32"] = "bf16"
    compile: bool = True
    eval_every: int = 500
    eval_batches: int = 50
    ckpt_every: int = 1000
    out_dir: str = "checkpoints"
    seed: int = 42
    run_name: str | None = None
    unique_run_name: bool = True
    """Append a timestamp to ``run_name``: a second run must not overwrite the ``step-*.pt``
    series the ``TrainingReplay`` of the course reads."""
    workers: int = 0  # DataLoader workers; 0 keeps everything in the main process
    log_every: int = 10  # optimizer steps between training metrics

    def decoder(self) -> DecoderConfig:
        """The decoder config of this run: the preset (or ``model``) with ``block`` applied."""
        base = self.model if self.model is not None else preset(self.preset)
        return base.model_copy(update={"block": self.block})

def pick_device() -> str:
    """``RUKH_DEVICE`` if set, else CUDA when available, else CPU."""
    env = os.environ.get("RUKH_DEVICE")
    if env:
        return env
    return "cuda" if torch.cuda.is_available() else "cpu"


def param_groups(model: nn.Module, weight_decay: float) -> list[dict[str, Any]]:
    """Decoupled weight decay on matrices only: norms and biases are left alone."""
    decay = [p for p in model.parameters() if p.requires_grad and p.dim() >= 2]
    no_decay = [p for p in model.parameters() if p.requires_grad and p.dim() < 2]
    return [
        {"params": decay, "weight_decay": weight_decay},
        {"params": no_decay, "weight_decay": 0.0},
    ]

def forever(loader: DataLoader[Batch]) -> Iterator[Batch]:
    """Repeat a loader for as many steps as the schedule asks for."""
    while True:
        yield from loader

def maybe_compile(
    model: MoveDecoder, enabled: bool, sample: torch.Tensor | None = None
) -> nn.Module:
    """``torch.compile`` the model when asked; a failure is a warning, never a stopped run.

    ``torch.compile`` is lazy: on Windows without MSVC it only fails when the first forward
    reaches Inductor. ``sample`` (one batch of the training shape) forces that compilation here,
    where it can still fall back to eager, and Dynamo's own error suppression covers any later
    recompilation for a different shape.
    """
    if not enabled:
        return model
    try:
        import torch._dynamo as dynamo

        dynamo.config.suppress_errors = True
        compiled = torch.compile(model)
        if sample is not None:
            _, loss = compiled(sample, sample)
            if loss is not None:
                loss.backward()
            model.zero_grad(set_to_none=True)
        return compiled
    except Exception as exc:  # noqa: BLE001 - compilation backends fail in many ways
        log.warning("torch.compile is unavailable, training eagerly: %s", exc)
        model.zero_grad(set_to_none=True)
        return model

@torch.no_grad()
def evaluate(
    model: nn.Module,
    loader: DataLoader[Batch],
    batches: int,
    device: torch.device,
    autocast: Any = None,
) -> tuple[float, float]:
    """Validation loss and top-1 next-token accuracy over at most ``batches`` batches.

    The loss is token-weighted: each batch's mean is weighted by the number of non-``<pad>``
    targets it had, so a short last batch does not count as much as a full one.
    """
    was_training = model.training
    model.eval()
    loss_sum = 0.0
    weighted = 0
    hits = 0
    counted = 0
    for index, (x, y) in enumerate(loader):
        if index >= batches:
            break
        x, y = x.to(device), y.to(device)
        with autocast if autocast is not None else nullcontext():
            logits, loss = model(x, y)
        mask = y != IGNORE_INDEX
        tokens = int(mask.sum())
        if loss is not None and torch.isfinite(loss) and tokens:
            loss_sum += loss.float().item() * tokens
            weighted += tokens
        hits += int((logits.argmax(dim=-1) == y)[mask].sum())
        counted += tokens
    if was_training:
        model.train()
    return (loss_sum / weighted if weighted else math.nan, hits / counted if counted else 0.0)

def run_dir(cfg: TrainConfig, resume: Path | None = None) -> Path:
    """Where this run writes: the resumed run's folder, or ``out_dir/<run name>``.

    The name carries a timestamp unless ``unique_run_name`` is off, because two runs of the same
    config would otherwise write the same ``step-*.pt`` files and the second would quietly
    overwrite the checkpoint series of the first.
    """
    if resume is not None:
        return Path(resume).resolve().parent
    name = cfg.run_name or cfg.preset
    if cfg.unique_run_name:
        name = f"{name}-{datetime.now(UTC):%Y%m%d-%H%M%S}"
    return paths.resolve(cfg.out_dir) / name


def skip_batches(batches: Iterator[Batch], count: int) -> int:
    """Wind the batch stream forward ``count`` batches and return how many were skipped.

    A resumed run must not start again at the first window of the first epoch: it would train
    twice on the same games while the schedule believes it is halfway. Windows are memmap slices,
    so winding forward is cheap compared with a step, and it is logged because it is not free.
    """
    if count <= 0:
        return 0
    started = time.perf_counter()
    for _ in range(count):
        next(batches)
    log.info("skipped %d batches in %.1f s to resume", count, time.perf_counter() - started)
    return count

def train(cfg: TrainConfig, resume: Path | None = None, device: str | None = None) -> Path:
    """Train a ``MoveDecoder`` and return the path of the final checkpoint."""
    if min(cfg.max_steps, cfg.grad_accum, cfg.batch_size) < 1:
        raise ValueError("max_steps, grad_accum and batch_size must be positive")
    torch.manual_seed(cfg.seed)
    where = torch.device(device or pick_device())
    tokens_dir = paths.resolve(cfg.tokens_dir)
    model_cfg = cfg.decoder()

    train_set = PackedDataset(tokens_dir / "train", block=cfg.block)
    val_set = PackedDataset(tokens_dir / "val", block=cfg.block)
    if train_set.info.vocab_size != model_cfg.vocab_size:
        raise ValueError(
            f"{tokens_dir / 'train'} has vocab_size {train_set.info.vocab_size}, "
            f"the model expects {model_cfg.vocab_size}"
        )
    train_loader = make_loader(train_set, cfg.batch_size, seed=cfg.seed, workers=cfg.workers)
    val_loader = make_loader(
        val_set, cfg.batch_size, seed=cfg.seed, workers=0, shuffle=False, drop_last=False
    )
    if not len(train_loader):
        raise ValueError(f"{tokens_dir / 'train'} has fewer than {cfg.batch_size} windows")

    model = MoveDecoder(model_cfg).to(where)
    optimizer = torch.optim.AdamW(param_groups(model, cfg.weight_decay), lr=cfg.lr, betas=cfg.betas)
    start_step = 0
    best_val = math.inf
    run_id: str | None = None
    if resume is not None:
        payload = load_checkpoint(resume, map_location=where)
        start_step = restore(payload, model, optimizer)
        recorded = payload.get("best_val")
        best_val = float(recorded) if isinstance(recorded, int | float) else best_val
        previous = payload.get("run_id")
        run_id = str(previous) if isinstance(previous, str) and previous else None
        log.info("resumed %s at step %d (mlflow run %s)", resume, start_step, run_id or "new")

    use_bf16 = cfg.precision == "bf16" and where.type == "cuda"
    autocast = (
        torch.autocast(device_type="cuda", dtype=torch.bfloat16) if use_bf16 else nullcontext()
    )
    warmup = torch.ones((cfg.batch_size, cfg.block), dtype=torch.long, device=where)
    with autocast:  # compile under the same precision the loop will use
        runnable = maybe_compile(model, cfg.compile, warmup)
    runnable.train()

    out_dir = run_dir(cfg, resume)
    out_dir.mkdir(parents=True, exist_ok=True)
    vocab_hash = read_vocab_hash(tokens_dir / "train")
    manifest_sha = read_manifest_sha(paths.data_dir() / "raw" / "manifest.json")
    tokens_per_step = cfg.batch_size * cfg.grad_accum * cfg.block
    batches = forever(train_loader)
    skip_batches(batches, start_step * cfg.grad_accum)
    final = out_dir / step_name(cfg.max_steps)

    from rukh.tracking import git_sha, start_run

    params = {
        **cfg.model_dump(mode="json"),
        "vocab_hash": vocab_hash,
        "data_manifest_sha": manifest_sha,
        "device": str(where),
        "num_params": model.num_params(),
    }
    with start_run(out_dir.name, params, tags={"preset": cfg.preset}, run_id=run_id) as run:
        log.info("run %s in %s on %s", run.info.run_id, out_dir, where)
        this_run = str(run.info.run_id)

        def save(path: Path, step: int) -> Path:
            return save_checkpoint(
                path,
                step=step,
                model=model,
                optimizer=optimizer,
                cfg=cfg.model_dump(mode="json"),
                model_cfg=model_cfg.model_dump(mode="json"),
                vocab_hash=vocab_hash,
                data_manifest_sha=manifest_sha,
                git_sha=git_sha(),
                best_val=None if math.isinf(best_val) else best_val,
                run_id=this_run,
            )

        clock = time.perf_counter()
        for step in range(start_step, cfg.max_steps):
            lr = lr_at(step, cfg)
            for group in optimizer.param_groups:
                group["lr"] = lr
            optimizer.zero_grad(set_to_none=True)
            # Token-weighted, on the device: one synchronisation per step instead of one per
            # micro-batch, and a micro-batch with fewer real tokens weighs less in the mean.
            loss_sum = torch.zeros((), device=where)
            real_tokens = torch.zeros((), device=where)
            for _ in range(cfg.grad_accum):
                x, y = next(batches)
                x, y = x.to(where), y.to(where)
                with autocast:
                    _, loss = runnable(x, y)
                assert loss is not None
                (loss / cfg.grad_accum).backward()
                tokens = (y != IGNORE_INDEX).sum()
                loss_sum += loss.detach().float() * tokens
                real_tokens += tokens
            grad_norm = float(nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip))
            optimizer.step()

            done = step + 1
            last = done == cfg.max_steps
            if done % cfg.log_every == 0 or last:
                counted = float(real_tokens.item())
                total = float(loss_sum.item()) / counted if counted else math.nan
                elapsed = max(time.perf_counter() - clock, 1e-9)
                steps = min(cfg.log_every, done - start_step)
                log_metrics(
                    {
                        "train/loss": total,
                        "lr": lr,
                        "grad_norm": grad_norm,
                        "tokens_per_s": tokens_per_step * steps / elapsed,
                        "real_tokens_per_s": counted * steps / elapsed,
                    },
                    step=done,
                )
                clock = time.perf_counter()
            if done % cfg.eval_every == 0 or last:
                val_loss, top1 = evaluate(
                    runnable, val_loader, cfg.eval_batches, where, autocast if use_bf16 else None
                )
                log_metrics({"val/loss": val_loss, "val/top1": top1}, step=done)
                log.info("step %d  val/loss %.4f  val/top1 %.4f", done, val_loss, top1)
                if not math.isnan(val_loss) and val_loss < best_val:
                    best_val = val_loss
                    save(out_dir / BEST_NAME, done)
                clock = time.perf_counter()
            if done % cfg.ckpt_every == 0 or last:
                final = save(out_dir / step_name(done), done)
                clock = time.perf_counter()
    return final

def log_metrics(metrics: dict[str, float], step: int) -> None:
    """Send metrics to the active MLflow run; a tracking failure never stops training."""
    import mlflow

    clean = {key: value for key, value in metrics.items() if not math.isnan(value)}
    if not clean:
        return
    try:
        mlflow.log_metrics(clean, step=step)
    except Exception as exc:  # noqa: BLE001 - tracking is not worth a lost run
        log.warning("could not log metrics at step %d: %s", step, exc)

"""Checkpoints: what a run writes so it can be resumed, evaluated and published.

A payload holds only plain Python values and tensors, so it can be read back with
``torch.load(..., weights_only=True)``. Provenance (``vocab_hash``, ``data_manifest_sha``,
``git_sha``) is recorded on purpose: weights whose data or vocabulary cannot be identified are
not reproducible.
"""

from __future__ import annotations

import hashlib
import json
import random
from collections.abc import Mapping
from pathlib import Path
from typing import Any

import numpy as np
import torch
from torch import nn

from rukh.models import DecoderConfig, MoveDecoder

BEST_NAME = "best.pt"
TIED_HEAD = "lm_head.weight"
"""Tied to ``tokens.weight``; a published state dict leaves it out and it is re-tied on load."""

def step_name(step: int) -> str:
    """File name of the checkpoint written after ``step`` optimizer steps."""
    return f"step-{step}.pt"


def read_vocab_hash(tokens_dir: Path) -> str | None:
    """``vocab_hash`` from a pack ``meta.json``, or None when it cannot be read."""
    try:
        meta = json.loads((Path(tokens_dir) / "meta.json").read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return None
    value = meta.get("vocab_hash")
    return str(value) if isinstance(value, str) else None


def read_manifest_sha(manifest: Path) -> str | None:
    """SHA-256 of a data manifest, or None when the file is absent or unreadable.

    Training must never depend on ``data/`` being present: a checkpoint trained from a pack
    copied elsewhere simply records ``None``.
    """
    try:
        payload = Path(manifest).read_bytes()
    except OSError:
        return None
    return hashlib.sha256(payload).hexdigest()

def rng_state() -> dict[str, Any]:
    """Snapshot of the Python, NumPy and torch CPU generators."""
    name, keys, pos, has_gauss, cached = np.random.get_state(legacy=True)
    return {
        "python": json.dumps(random.getstate()),
        "numpy": json.dumps(
            [name, np.asarray(keys).tolist(), int(pos), int(has_gauss), float(cached)]
        ),
        "torch": torch.get_rng_state(),
    }


def set_rng_state(state: Mapping[str, Any]) -> None:
    """Restore a snapshot taken by ``rng_state`` (missing entries are skipped)."""
    python = state.get("python")
    if isinstance(python, str):
        version, keys, gauss = json.loads(python)
        random.setstate((version, tuple(keys), gauss))
    numpy = state.get("numpy")
    if isinstance(numpy, str):
        name, keys, pos, has_gauss, cached = json.loads(numpy)
        np.random.set_state((name, np.array(keys, dtype=np.uint32), pos, has_gauss, cached))
    torch_state = state.get("torch")
    if isinstance(torch_state, torch.Tensor):
        torch.set_rng_state(torch_state.to(torch.uint8).cpu())

def save_checkpoint(
    path: Path,
    *,
    step: int,
    model: nn.Module,
    optimizer: torch.optim.Optimizer | None,
    cfg: Mapping[str, Any],
    model_cfg: Mapping[str, Any],
    vocab_hash: str | None = None,
    data_manifest_sha: str | None = None,
    git_sha: str | None = None,
    best_val: float | None = None,
    run_id: str | None = None,
) -> Path:
    """Write one checkpoint atomically (temporary file plus replace) and return its path.

    ``run_id`` is the MLflow run that wrote it, so ``--resume`` can carry on logging into the
    same run and the loss curve stays one line instead of two.
    """
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    payload: dict[str, Any] = {
        "step": int(step),
        "model_state": {k: v.detach().cpu() for k, v in model.state_dict().items()},
        "opt_state": optimizer.state_dict() if optimizer is not None else None,
        "cfg": dict(cfg),
        "model_cfg": dict(model_cfg),
        "vocab_hash": vocab_hash,
        "data_manifest_sha": data_manifest_sha,
        "git_sha": git_sha,
        "best_val": best_val,
        "run_id": run_id,
        "rng": rng_state(),
    }
    tmp = path.with_suffix(path.suffix + ".tmp")
    torch.save(payload, tmp)
    tmp.replace(path)
    return path

def load_checkpoint(path: Path, map_location: str | torch.device = "cpu") -> dict[str, Any]:
    """Read a checkpoint written by ``save_checkpoint``."""
    payload = torch.load(Path(path), map_location=map_location, weights_only=True)
    if not isinstance(payload, dict) or "model_state" not in payload:
        raise ValueError(f"{path} is not a rukh checkpoint")
    return payload


def load_state(model: nn.Module, state: Mapping[str, Any]) -> None:
    """Load weights, accepting a state dict whose tied head was left out when it was saved.

    ``MoveDecoder`` ties ``lm_head.weight`` to ``tokens.weight`` in its constructor, so the tied
    head is already correct once the embedding is loaded: a state dict without it (the one
    ``rukh publish`` writes, because two names for one tensor is what makes ``safetensors``
    refuse the file) loads cleanly and nothing else may be missing.
    """
    missing, unexpected = model.load_state_dict(dict(state), strict=False)
    absent = [name for name in missing if name != TIED_HEAD]
    if absent or unexpected:
        raise ValueError(
            f"the weights do not match the model: missing {absent}, unexpected {list(unexpected)}"
        )

def load_model(
    path: Path, map_location: str | torch.device = "cpu"
) -> tuple[MoveDecoder, dict[str, Any]]:
    """Rebuild the decoder a checkpoint describes, in eval mode, plus the whole payload."""
    payload = load_checkpoint(path, map_location=map_location)
    model = MoveDecoder(DecoderConfig.model_validate(payload["model_cfg"]))
    load_state(model, payload["model_state"])
    return model.eval(), payload


def restore(
    payload: Mapping[str, Any],
    model: nn.Module,
    optimizer: torch.optim.Optimizer | None = None,
    with_rng: bool = True,
) -> int:
    """Load weights (and optionally the optimizer and RNG) and return the step reached."""
    load_state(model, payload["model_state"])
    opt_state = payload.get("opt_state")
    if optimizer is not None and opt_state is not None:
        optimizer.load_state_dict(opt_state)
    rng = payload.get("rng")
    if with_rng and isinstance(rng, Mapping):
        set_rng_state(rng)
    return int(payload["step"])

"""Training: the loop, the learning-rate schedule and checkpoint handling."""

from rukh.train.checkpoint import (
    BEST_NAME,
    TIED_HEAD,
    load_checkpoint,
    load_model,
    load_state,
    read_manifest_sha,
    read_vocab_hash,
    restore,
    save_checkpoint,
    step_name,
)
from rukh.train.loop import (
    TrainConfig,
    evaluate,
    param_groups,
    pick_device,
    run_dir,
    skip_batches,
    train,
)
from rukh.train.schedule import lr_at

__all__ = [
    "BEST_NAME",
    "TIED_HEAD",
    "TrainConfig",
    "evaluate",
    "load_checkpoint",
    "load_model",
    "load_state",
    "lr_at",
    "param_groups",
    "pick_device",
    "read_manifest_sha",
    "read_vocab_hash",
    "restore",
    "run_dir",
    "save_checkpoint",
    "skip_batches",
    "step_name",
    "train",
]

"""Tests for rukh.train.checkpoint: round-trip, provenance and RNG restoration."""

from __future__ import annotations

import json
import random
from pathlib import Path

import numpy as np
import pytest
import torch

from rukh.models import DecoderConfig, MoveDecoder
from rukh.train import (
    TrainConfig,
    load_checkpoint,
    read_manifest_sha,
    read_vocab_hash,
    restore,
    save_checkpoint,
    step_name,
)

pytestmark = pytest.mark.unit

TOY = DecoderConfig(vocab_size=32, n_layer=2, n_head=2, d_model=32, block=16)


def make_pair() -> tuple[MoveDecoder, torch.optim.Optimizer]:
    torch.manual_seed(0)
    model = MoveDecoder(TOY)
    return model, torch.optim.AdamW(model.parameters(), lr=1e-3)

def test_save_and_load_reproduce_weights_and_step(tmp_path: Path) -> None:
    model, optimizer = make_pair()
    idx = torch.randint(1, TOY.vocab_size, (2, 8))
    _, loss = model(idx, idx)
    assert loss is not None
    loss.backward()
    optimizer.step()

    path = save_checkpoint(
        tmp_path / step_name(7),
        step=7,
        model=model,
        optimizer=optimizer,
        cfg=TrainConfig(max_steps=7).model_dump(mode="json"),
        model_cfg=TOY.model_dump(mode="json"),
        vocab_hash="abc",
        data_manifest_sha=None,
        git_sha="deadbeef",
        best_val=1.25,
    )
    assert path.name == "step-7.pt"
    assert not list(tmp_path.glob("*.tmp"))

    payload = load_checkpoint(path)
    assert payload["step"] == 7
    assert payload["vocab_hash"] == "abc"
    assert payload["data_manifest_sha"] is None
    assert payload["git_sha"] == "deadbeef"
    assert payload["best_val"] == pytest.approx(1.25)
    assert payload["cfg"]["max_steps"] == 7
    assert payload["model_cfg"]["d_model"] == 32

    fresh = MoveDecoder(TOY)
    fresh_opt = torch.optim.AdamW(fresh.parameters(), lr=1e-3)
    assert restore(payload, fresh, fresh_opt) == 7
    for a, b in zip(model.state_dict().values(), fresh.state_dict().values(), strict=True):
        assert torch.equal(a, b)
    with torch.no_grad():
        assert torch.allclose(model.eval()(idx)[0], fresh.eval()(idx)[0], atol=1e-6)

def test_restoring_the_rng_repeats_the_same_draws(tmp_path: Path) -> None:
    model, optimizer = make_pair()
    random.seed(11)
    np.random.seed(11)
    torch.manual_seed(11)
    path = save_checkpoint(
        tmp_path / "rng.pt",
        step=0,
        model=model,
        optimizer=optimizer,
        cfg={},
        model_cfg=TOY.model_dump(mode="json"),
    )
    expected = (random.random(), float(np.random.random()), torch.rand(3).tolist())

    random.seed(99)
    np.random.seed(99)
    torch.manual_seed(99)
    restore(load_checkpoint(path), MoveDecoder(TOY))
    assert (random.random(), float(np.random.random()), torch.rand(3).tolist()) == expected

def test_vocab_hash_and_manifest_sha_are_defensive(tmp_path: Path) -> None:
    assert read_vocab_hash(tmp_path) is None
    (tmp_path / "meta.json").write_text(json.dumps({"n_games": 1}), encoding="utf-8")
    assert read_vocab_hash(tmp_path) is None
    (tmp_path / "meta.json").write_text(json.dumps({"vocab_hash": "cafe"}), encoding="utf-8")
    assert read_vocab_hash(tmp_path) == "cafe"

    assert read_manifest_sha(tmp_path / "does-not-exist" / "manifest.json") is None
    manifest = tmp_path / "manifest.json"
    manifest.write_bytes(b"{}")
    sha = read_manifest_sha(manifest)
    assert sha is not None and len(sha) == 64


def test_a_foreign_file_is_not_a_checkpoint(tmp_path: Path) -> None:
    path = tmp_path / "other.pt"
    torch.save({"weights": torch.zeros(2)}, path)
    with pytest.raises(ValueError, match="not a rukh checkpoint"):
        load_checkpoint(path)

def start_run(
    name: str,
    config: Mapping[str, Any],
    tags: Mapping[str, str] | None = None,
    run_id: str | None = None,
) -> Iterator[Any]:
    """Start an MLflow run in the local store with ``config`` logged as flattened params.

    Tags always include ``rukh_version`` and, when the project lives in a git repo, ``git_sha``.
    With ``run_id`` the existing run is reopened instead of a new one being created, so a
    training run that is resumed keeps one curve rather than starting a second one; a run id
    that cannot be reopened (deleted, or from another store) falls back to a fresh run.
    Yields the active ``mlflow.ActiveRun``.
    """

    started = None
    if run_id:
        try:
            started = mlflow.start_run(run_id=run_id, tags=run_tags)
        except Exception as exc:  # noqa: BLE001 - a lost run must never stop a training run
            log.warning("could not reopen the MLflow run %s (%s); starting a new one", run_id, exc)
    with started or mlflow.start_run(run_name=name, tags=run_tags) as run:
        params = flatten(config)
        if params:
            # A resumed run already carries these; re-logging a changed value is an error there.
            try:
                mlflow.log_params(params)
            except Exception as exc:  # noqa: BLE001 - the params are already recorded
                log.warning("could not log the run parameters: %s", exc)
        yield run

# `small` (12 layers, d=512, 38,971,392 parameters): the course model. Measured at 440k tokens/s
# eager on a 5090 (D-024), so the 20,000 steps below are about 40 minutes, not the 4-8 hours the
# spec estimated before anyone timed it.
preset: small
tokens_dir: data/tokens/uci
block: 200
batch_size: 64
grad_accum: 4 # 256 effective sequences of 200 tokens
lr: 6.0e-4
min_lr_ratio: 0.1
warmup: 1000
max_steps: 20000
weight_decay: 0.1
betas: [0.9, 0.95]
grad_clip: 1.0
precision: bf16
compile: true
eval_every: 500
eval_batches: 50
ckpt_every: 1000
out_dir: checkpoints
seed: 42
run_name: small
# A timestamp is appended: a second run must not overwrite the step-*.pt series.
unique_run_name: true
workers: 4
log_every: 10

# `tiny` (6 layers, d=256, ~5M parameters): the model used to iterate, under 20 minutes on a 5090.
preset: tiny
tokens_dir: data/tokens/uci
block: 200
batch_size: 128
grad_accum: 2 # 256 effective sequences
lr: 1.0e-3
min_lr_ratio: 0.1
warmup: 500
max_steps: 6000
weight_decay: 0.1
betas: [0.9, 0.95]
grad_clip: 1.0
precision: bf16
compile: true
eval_every: 250
eval_batches: 50
ckpt_every: 1000 # every checkpoint feeds the TrainingReplay island of M2
out_dir: checkpoints
seed: 42
run_name: tiny
# A timestamp is appended: a second run must not overwrite the step-*.pt series.
unique_run_name: true
workers: 4
log_every: 10

@app.command("train")
def train_cmd(
    config: Annotated[
        Path,
        typer.Option(
            "--config", exists=True, dir_okay=False, readable=True, help="Training YAML config."
        ),
    ],
    model_preset: Annotated[
        str | None, typer.Option("--preset", help="Override the preset: tiny, small or medium.")
    ] = None,
    resume: Annotated[
        Path | None,
        typer.Option(
            "--resume", exists=True, dir_okay=False, readable=True, help="Checkpoint to continue."
        ),
    ] = None,
    max_steps: Annotated[
        int | None, typer.Option("--max-steps", help="Override max_steps from the config.")
    ] = None,
) -> None:
    """Train a MoveDecoder from a packed token stream, logging the run to MLflow."""
    from rukh.config import load_yaml
    from rukh.models import PRESETS
    from rukh.train import TrainConfig, train

    cfg = load_yaml(config, TrainConfig)
    if model_preset is not None:
        if model_preset not in PRESETS:
            typer.echo(f"error: --preset must be one of {', '.join(PRESETS)}", err=True)
            raise typer.Exit(code=2)
        cfg = cfg.model_copy(update={"preset": model_preset, "model": None})
    if max_steps is not None:
        cfg = cfg.model_copy(update={"max_steps": max_steps})
    logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(message)s")
    try:
        checkpoint = train(cfg, resume=resume)
    except (FileNotFoundError, ValueError) as exc:
        typer.echo(f"error: {exc}", err=True)
        raise typer.Exit(code=1) from exc
    typer.echo(f"preset:     {cfg.preset}")
    typer.echo(f"steps:      {cfg.max_steps}")
    typer.echo(f"checkpoint: {checkpoint}")