int8 y el jugador ONNX: un grafo exportado en la escalera

"""An ONNX export plays through the same player as the checkpoint it came from."""

from __future__ import annotations

from pathlib import Path
from typing import Any

import chess
import numpy as np
import pytest
import torch

from rukh.export.onnx import INPUT_NAME, LastStepLogits
from rukh.infer import DecoderPlayer, SampleConfig
from rukh.infer import onnx_player as module
from rukh.infer.onnx_player import OnnxDecoder
from rukh.models.decoder import DecoderConfig, MoveDecoder
from rukh.tokenize.uci_vocab import UciTokenizer

pytestmark = pytest.mark.unit

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


class FakeSession:
    """Stands in for ``onnxruntime.InferenceSession`` with the checkpoint's own logits."""

    def __init__(self, model: MoveDecoder, block: int | None = TOY.block) -> None:
        self.wrapper = LastStepLogits(model).eval()
        self.block = block
        self.seen: list[int] = []

    def get_modelmeta(self) -> Any:
        meta = {} if self.block is None else {module.METADATA_BLOCK: str(self.block)}
        return type("Meta", (), {"custom_metadata_map": meta})()

    def run(self, _outputs: Any, feed: dict[str, np.ndarray]) -> list[np.ndarray]:
        idx = feed[INPUT_NAME]
        assert idx.dtype == np.int64
        self.seen.append(idx.shape[1])
        with torch.no_grad():
            return [self.wrapper(torch.from_numpy(idx)).numpy()]

def test_the_export_picks_the_checkpoints_move_greedily(
    model: MoveDecoder, monkeypatch: pytest.MonkeyPatch
) -> None:
    session = FakeSession(model)
    monkeypatch.setattr(module, "_session", lambda _path: session)
    tok = UciTokenizer()
    cfg = SampleConfig(temperature=0, top_k=None, mask_illegal=False, seed=0)

    reference = DecoderPlayer(model, tok, cfg)
    exported = DecoderPlayer(OnnxDecoder(Path("model.onnx")), tok, cfg)  # type: ignore[arg-type]
    assert exported.limit == reference.limit

    board = chess.Board()
    for player in (reference, exported):
        player.start(1800, 1800)
    for _ in range(6):
        a, legal_a = reference.choose(board)
        b, legal_b = exported.choose(board)
        assert a == b and legal_a == legal_b
        if a is None:
            break
        board.push(a)
        reference.observe(a)
        exported.observe(a)
    assert session.seen and max(session.seen) <= TOY.block

def test_the_block_comes_from_the_metadata_or_the_caller(
    model: MoveDecoder, monkeypatch: pytest.MonkeyPatch
) -> None:
    monkeypatch.setattr(module, "_session", lambda _path: FakeSession(model, block=None))
    with pytest.raises(ValueError, match="rukh_block"):
        OnnxDecoder(Path("model.onnx"))
    assert OnnxDecoder(Path("model.onnx"), block=32).cfg.block == 32
    monkeypatch.setattr(module, "_session", lambda _path: FakeSession(model, block=48))
    assert OnnxDecoder(Path("model.onnx")).cfg.block == 48

RUN_STAMP = re.compile(r"^(?P<name>.+)-(?P<stamp>\d{8}-\d{6})$")
"""What ``unique_run_name`` appends to a run folder: ``medium-v4`` becomes
``medium-v4-20260919-174623``. The stamp sorts as text, so the newest run of a name is the last."""


def resolve_run(path: str | Path) -> Path:
    """The path as given when it exists; otherwise the same thing inside the newest stamped run.

    Configs and lessons name a checkpoint by its run: ``checkpoints/medium-v4/best.pt``. That is
    the spelling ``rukh pull`` writes, and it is not the one a reader who trained the run has,
    because ``unique_run_name`` stamps every folder (``medium-v4-20260919-174623``). Both mean
    "the run called medium-v4", so when the stable folder is missing this looks for the stamped
    ones beside it and takes the newest **by name**, never by modification time: copying a
    folder must not change which run a config points at. A folder spec (``checkpoints/lora-e4``)
    resolves the same way. Anything with no candidate comes back unchanged, and whoever opens it
    reports the missing file with the path the user wrote.
    """
    path = Path(path)
    if path.exists():
        return path
    is_file = bool(path.suffix)
    run_dir = path.parent if is_file else path
    parent, name = run_dir.parent, run_dir.name
    if not parent.is_dir():
        return path
    stamped = sorted(
        candidate
        for candidate in parent.iterdir()
        if candidate.is_dir()
        and (match := RUN_STAMP.match(candidate.name)) is not None
        and match.group("name") == name
        and (not is_file or (candidate / path.name).is_file())
    )
    if not stamped:
        return path
    found = stamped[-1] / path.name if is_file else stamped[-1]
    log.info("%s is not there; using the newest run of that name, %s", path, found)
    return found

def load_checkpoint(path: Path, map_location: str | torch.device = "cpu") -> dict[str, Any]:
    """Read a checkpoint written by ``save_checkpoint`` (see ``resolve_run`` for the path)."""
    payload = torch.load(resolve_run(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

    if init_from is None:
        return None
    path = resolve_run(paths.resolve(init_from))
    if not path.is_file():
        raise FileNotFoundError(f"init_from checkpoint not found: {path}")

    read_vocab_hash,
    resolve_run,
    restore,

def _run_path(value: Path | None) -> Path | None:
    """A checkpoint or run folder; a missing stable name resolves to the newest stamped run.

    A path is validated here instead of with ``exists=True`` so that a command written against
    the stable spelling of a run (``checkpoints/small/best.pt``, what ``rukh pull`` writes) also
    works for the reader whose own training left ``checkpoints/small-20260919-062911/`` behind.
    """
    if value is None:
        return None
    from rukh.train.checkpoint import resolve_run

    found = resolve_run(value)
    if not found.exists():
        run = value.parent.name if value.suffix else value.name
        raise typer.BadParameter(f"{value} does not exist, and no stamped run `{run}-*` holds it")
    return found


def _run_exists(spec: str) -> bool:
    """Whether ``spec`` names a checkpoint on disk, directly or through ``resolve_run``."""
    from rukh.train.checkpoint import resolve_run

    return resolve_run(Path(spec)).is_file()

def test_a_stable_run_name_resolves_to_the_newest_stamped_run(tmp_path: Path) -> None:
    """``checkpoints/small/best.pt`` finds ``checkpoints/small-<newest stamp>/best.pt``.

    This is what lets a config or a lesson name a run without its date: the reader who trained
    it has the stamped folder, the reader who pulled it from the Hub has the stable one, and the
    same spelling serves both.
    """
    from rukh.train.checkpoint import resolve_run

    root = tmp_path / "checkpoints"
    older = root / "small-20260919-062911"
    newer = root / "small-20260920-150000"
    decoy = root / "small-v3-20260921-000000"  # another run whose name merely starts the same
    for folder in (older, newer, decoy):
        folder.mkdir(parents=True)
        (folder / "best.pt").write_bytes(b"x")
    (root / "small-20260921-000000").mkdir()  # newest stamp, but it holds no best.pt yet

    assert resolve_run(root / "small" / "best.pt") == newer / "best.pt"
    assert resolve_run(root / "small") == root / "small-20260921-000000"
    assert resolve_run(root / "small-v3" / "best.pt") == decoy / "best.pt"
    # A stable folder wins over any stamped one, and an unknown name comes back untouched.
    (root / "small").mkdir()
    (root / "small" / "best.pt").write_bytes(b"y")
    assert resolve_run(root / "small" / "best.pt") == root / "small" / "best.pt"
    assert resolve_run(root / "medium" / "best.pt") == root / "medium" / "best.pt"
    assert resolve_run(tmp_path / "nowhere" / "run" / "best.pt") == (
        tmp_path / "nowhere" / "run" / "best.pt"
    )


def test_load_checkpoint_reads_through_the_stable_name(tmp_path: Path) -> None:
    model, optimizer = make_pair()
    run = tmp_path / "checkpoints" / "tiny-20260919-060101"
    save_checkpoint(
        run / "best.pt",
        step=5,
        model=model,
        optimizer=optimizer,
        cfg={},
        model_cfg=TOY.model_dump(),
    )
    payload = load_checkpoint(tmp_path / "checkpoints" / "tiny" / "best.pt")
    assert payload["step"] == 5