El Elo y la suite

def play_rung(
    model: MoveDecoder,
    tok: UciTokenizer,
    rung: EloRung,
    games: int,
    cfg: SampleConfig,
    move_time: float = MOVE_TIME,
    max_plies: int | None = None,
    cache: EvalCache | None = None,
) -> list[GameRecord]:
    """Play ``games`` games against one rung, alternating colours, reusing cached games.

    A game that ends with ``*`` (the context ran out) is adjudicated with the rung's own engine
    before it is scored, so it enters the fit as a win, a loss or a draw on the merits of the
    final position.
    """
    records: list[GameRecord] = []
    for index in range(games):
        cached = cache.get(SUITE, f"{rung.name}:{index}") if cache is not None else None
        if cached is not None:
            records.append(GameRecord.model_validate(cached))
    done = {record.index for record in records}
    missing = [index for index in range(games) if index not in done]
    if not missing:
        return sorted(records, key=lambda r: r.index)

    opponent = StockfishOpponent(elo=rung.uci_elo or 1320, skill=rung.skill, move_time=move_time)
    try:
        for index in missing:
            model_white = index % 2 == 0
            outcome = play_game(
                model,
                tok,
                opponent,
                cfg.model_copy(update={"seed": None if cfg.seed is None else cfg.seed + index}),
                model_color=chess.WHITE if model_white else chess.BLACK,
                max_plies=max_plies,
            )
            record = record_of(outcome, rung, index, model_white, engine=opponent.engine)
            if cache is not None:
                cache.put(SUITE, record.item_id(), record.model_dump())
            records.append(record)
    finally:
        opponent.close()
    return sorted(records, key=lambda r: r.index)

def play_rungs(
    model: MoveDecoder,
    tok: UciTokenizer,
    rungs: Sequence[EloRung],
    games: int,
    cfg: SampleConfig,
    move_time: float = MOVE_TIME,
    max_plies: int | None = None,
    cache: EvalCache | None = None,
) -> list[GameRecord]:
    """Play every rung and return all the game records."""
    records: list[GameRecord] = []
    for rung in rungs:
        records.extend(
            play_rung(
                model,
                tok,
                rung,
                games,
                cfg,
                move_time=move_time,
                max_plies=max_plies,
                cache=cache,
            )
        )
    return records

from __future__ import annotations

import json
import logging
import re
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from pydantic import BaseModel, ConfigDict, Field

from rukh import paths
from rukh.config import BaseConfig
from rukh.eval.accuracy import AccuracyResult, accuracy
from rukh.eval.cache import EvalCache, config_sha, file_sha
from rukh.eval.elo import DEFAULT_RUNGS, EloResult, EloRung, estimate, play_rungs
from rukh.eval.legality import LegalityResult, legality, sample_positions
from rukh.eval.puzzles import PuzzleResult, load_puzzles, model_source, run_puzzles
from rukh.eval.report import ReportPaths, write_report
from rukh.infer import SampleConfig
from rukh.models import DecoderConfig
from rukh.tokenize.uci_vocab import UciTokenizer

log = logging.getLogger(__name__)

SUITES = ("full", "quick")
HUB_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9._-]+$")
HUB_WEIGHTS = ("model.safetensors", "pytorch_model.bin")
ELO_CAVEAT = (
    "the Elo interval covers sampling noise only: the four ``skill-*`` rungs are nominal "
    "``Skill Level`` anchors rather than measured ratings, and Stockfish plays at {move_time:g} s "
    "per move, far below any setting ``UCI_Elo`` is calibrated for"
)
LEGALITY_DEFINITIONS = (
    "legality_argmax is the share of validation positions where the single most likely token is "
    "a legal move (no temperature, no top-k, no mask): this is the >= 99 % bar of GOAL.md. "
    "legality_sampled draws the token the way the demo does (temperature {temperature:g}"
    "{top_k}) and is always the lower of the two."
)

class EvalConfig(BaseConfig):
    """What a suite measures and where it reads and writes."""

    stage: str | None = None
    """Name of the row in the results table; defaults to the checkpoint's run directory."""
    games: str = "data/uci/year=2025/month=02/games.parquet"
    puzzles: str = "data/puzzles/puzzles.parquet"
    puzzle_split: str = "test"
    legality_positions: int = 10_000
    accuracy_positions: int = 10_000
    position_pool: int = 50_000
    puzzles_per_band: int = 2_000
    elo_games: int = 100
    elo_rungs: list[EloRung] = Field(default_factory=lambda: list(DEFAULT_RUNGS))
    elo_move_time: float = 0.1
    elo_max_plies: int | None = None
    bootstrap: int = 1_000
    temperature: float = 0.6
    top_k: int | None = 20
    block: int = 200
    seed: int = 42
    out_dir: str = "artifacts/eval"
    web_results: str | None = "artifacts/web/results.json"
    cache_db: str = "artifacts/eval/cache.sqlite"
    device: str | None = None
    """Where the model runs; ``None`` means ``rukh.train.pick_device()`` (CUDA when present)."""
    track: bool = True
    """Log the suite to the local MLflow store (the model card reads those metrics back)."""

    def cache_fields(self) -> dict[str, Any]:
        """The settings that change what a cached game or puzzle means."""
        return {
            "temperature": self.temperature,
            "top_k": self.top_k,
            "elo_move_time": self.elo_move_time,
            "elo_max_plies": self.elo_max_plies,
            "elo_games": self.elo_games,
            "rungs": [rung.model_dump(mode="json") for rung in self.elo_rungs],
            "seed": self.seed,
            "block": self.block,
        }

    def sampling(self) -> SampleConfig:
        """The sampler the Elo games use: masked, seeded, as the demo plays."""
        return SampleConfig(
            temperature=self.temperature,
            top_k=self.top_k,
            mask_illegal=True,
            seed=self.seed,
        )

class SuiteResult(BaseModel):
    """Everything one evaluation produced; serialised verbatim as ``results.json``."""

    model_config = ConfigDict(extra="forbid")

    stage: str
    suite: str
    checkpoint: str
    model_sha: str
    params: int
    date: str
    device: str = "cpu"
    run_id: str | None = None
    legality_argmax: LegalityResult | None = None
    """The headline legality rate: the most likely token, unmasked."""
    legality_sampled: LegalityResult | None = None
    """Legality of a token drawn the way the demo draws it (temperature and top-k)."""
    accuracy: AccuracyResult | None = None
    puzzles: PuzzleResult | None = None
    elo: EloResult | None = None
    delta_cp: float | None = None
    """Mean centipawn loss; measured by a later milestone, ``null`` until then."""
    diversity: float | None = None
    """Opening entropy over self-play games; measured by a later milestone."""
    notes: list[str] = []
    config: dict[str, Any] = {}

def config_path(suite: str) -> Path:
    """Where the YAML of a named suite lives inside the source tree."""
    if suite not in SUITES:
        raise ValueError(f"unknown suite {suite!r}; expected one of {', '.join(SUITES)}")
    return paths.package_root() / "configs" / "eval" / f"{suite}.yaml"


def load_suite(suite: str, config: Path | None = None) -> EvalConfig:
    """Load the suite config: an explicit path wins over the named one."""
    from rukh.config import load_yaml

    return load_yaml(config if config is not None else config_path(suite), EvalConfig)

def _positions(cfg: EvalConfig, tok: UciTokenizer, notes: list[str]) -> list[Any]:
    games = paths.resolve(cfg.games)
    if not games.is_file():
        notes.append(f"validation games not found at {games}: legality and accuracy skipped")
        return []
    return sample_positions(
        games,
        max(cfg.legality_positions, cfg.accuracy_positions),
        tok,
        seed=cfg.seed,
        pool=cfg.position_pool,
        block=cfg.block,
    )


def _elo(
    model: Any, tok: UciTokenizer, cfg: EvalConfig, cache: EvalCache | None, notes: list[str]
) -> EloResult | None:
    from rukh.engine import EngineNotFound

    try:
        records = play_rungs(
            model,
            tok,
            cfg.elo_rungs,
            cfg.elo_games,
            cfg.sampling(),
            move_time=cfg.elo_move_time,
            max_plies=cfg.elo_max_plies,
            cache=cache,
        )
    except EngineNotFound as exc:
        notes.append(f"Elo skipped: {exc}")
        return None
    if not records:
        notes.append("Elo skipped: no games were played")
        return None
    return estimate(records, samples=cfg.bootstrap, seed=cfg.seed)

def _metrics(result: SuiteResult) -> dict[str, float]:
    """Flat metrics for MLflow (only what was actually measured)."""
    metrics: dict[str, float] = {}
    if result.legality_argmax is not None:
        metrics["legality_argmax"] = result.legality_argmax.rate
    if result.legality_sampled is not None:
        metrics["legality_sampled"] = result.legality_sampled.rate
    if result.accuracy is not None:
        metrics["top1"] = result.accuracy.top1
        metrics["top3"] = result.accuracy.top3
        for band in result.accuracy.bands:
            metrics[f"top1/{band.band}"] = band.top1
            metrics[f"top3/{band.band}"] = band.top3
    if result.puzzles is not None:
        metrics["puzzles"] = result.puzzles.rate
        for band in result.puzzles.bands:
            metrics[f"puzzles/{band.band}"] = band.rate
    if result.elo is not None:
        metrics["elo"] = result.elo.elo
        if result.elo.ci_low is not None and result.elo.ci_high is not None:
            metrics["elo_ci_low"] = result.elo.ci_low
            metrics["elo_ci_high"] = result.elo.ci_high
        if result.elo.elo_lower is not None:
            metrics["elo_lower_bound"] = result.elo.elo_lower
        if result.elo.elo_upper is not None:
            metrics["elo_upper_bound"] = result.elo.elo_upper
    return metrics

def is_hub_id(spec: str) -> bool:
    """Whether ``spec`` looks like a Hub repository id (``owner/name``) rather than a path."""
    text = str(spec)
    if Path(text).expanduser().is_file() or text.endswith(".pt"):
        return False
    return HUB_ID.fullmatch(text) is not None


def hub_checkpoint(repo_id: str, out_dir: Path | None = None) -> Path:
    """Download a published model and write it back as a checkpoint the harness can read.

    A Hub repository holds ``config.json`` and ``model.safetensors`` (or the torch pickle when
    safetensors was not installed at publication time), not a training checkpoint, so the two are
    folded into the ``{model_state, model_cfg, step}`` payload every other entry point expects.
    Nothing here runs while the command line is being parsed: a Hub id is validated by its shape.
    """
    import torch
    from huggingface_hub import hf_hub_download

    directory = out_dir if out_dir is not None else paths.resolve("artifacts/hub")
    directory.mkdir(parents=True, exist_ok=True)
    target = directory / (repo_id.replace("/", "__") + ".pt")
    config = json.loads(Path(hf_hub_download(repo_id, "config.json")).read_text(encoding="utf-8"))
    state: dict[str, Any] = {}
    for name in HUB_WEIGHTS:
        try:
            weights = Path(hf_hub_download(repo_id, name))
        except Exception as exc:  # noqa: BLE001 - a missing file only means "try the next one"
            log.debug("%s has no %s (%s)", repo_id, name, exc)
            continue
        if name.endswith(".safetensors"):
            from safetensors.torch import load_file

            state = dict(load_file(str(weights)))
        else:
            state = dict(torch.load(weights, map_location="cpu", weights_only=True))
        break
    if not state:
        raise FileNotFoundError(f"{repo_id} holds none of {', '.join(HUB_WEIGHTS)}")
    model_cfg = {key: value for key, value in config.items() if key in DecoderConfig.model_fields}
    torch.save(
        {
            "step": int(config.get("step") or 0),
            "model_state": state,
            "opt_state": None,
            "cfg": {},
            "model_cfg": model_cfg,
            "vocab_hash": config.get("vocab_hash"),
            "data_manifest_sha": config.get("data_manifest_sha"),
            "git_sha": config.get("git_sha"),
            "best_val": None,
        },
        target,
    )
    return target

def resolve_model(spec: str | Path) -> Path:
    """A checkpoint path from either a local file or a Hub id (downloaded on demand)."""
    text = str(spec)
    if is_hub_id(text):
        return hub_checkpoint(text)
    path = Path(text)
    if not path.is_file():
        raise FileNotFoundError(f"{path} is neither a checkpoint nor a Hub id (owner/name)")
    return path


def _legality_note(cfg: EvalConfig) -> str:
    top_k = "" if cfg.top_k is None else f", top-k {cfg.top_k}"
    return LEGALITY_DEFINITIONS.format(temperature=cfg.temperature, top_k=top_k)

def _elo_notes(elo: EloResult | None, cfg: EvalConfig, notes: list[str]) -> None:
    """Everything the Elo number has to be read with, in words."""
    if elo is None:
        return
    notes.append(ELO_CAVEAT.format(move_time=cfg.elo_move_time))
    if elo.cut:
        notes.append(
            f"{elo.cut} of {elo.games} games hit the context limit and were adjudicated "
            f"({elo.adjudicated} of them) instead of being scored as draws"
        )
    if elo.separated:
        notes.append(
            "every game went the same way, so the rating is not identified: the report gives a "
            "one-sided bound instead of a 95 % interval"
        )

def track_result(result: SuiteResult, cfg: EvalConfig) -> str | None:
    """Log the suite to the local MLflow store and return the run id."""
    import mlflow

    from rukh.tracking import start_run

    with start_run(
        f"eval-{result.stage}",
        {"suite": result.suite, "checkpoint": result.checkpoint, **cfg.model_dump(mode="json")},
        tags={"kind": "eval", "stage": result.stage, "model_sha": result.model_sha},
    ) as run:
        metrics = _metrics(result)
        if metrics:
            mlflow.log_metrics(metrics)
        return str(run.info.run_id)

def run_suite(
    ckpt: Path | str,
    cfg: EvalConfig,
    suite: str = "full",
    use_cache: bool = True,
    device: str | None = None,
) -> tuple[SuiteResult, ReportPaths]:
    """Evaluate, track and report: the whole of ``rukh eval`` in one call."""
    result = evaluate(ckpt, cfg, suite=suite, use_cache=use_cache, device=device)
    if cfg.track:
        result.run_id = track_result(result, cfg)
    report = write_report(
        result,
        paths.resolve(cfg.out_dir),
        paths.resolve(cfg.web_results) if cfg.web_results else None,
    )
    return result, report

"""Reporting: ``report.md`` and ``results.json`` per model, one shared row for the web.

Every stage of the course (base, tiny, small, masters, Elo-conditioned, LoRA, DPO, GRPO, the
external baselines) ends up as one row of the single table of ``docs/spec/02`` §"Componente 5".
``artifacts/web/results.json`` is that table: the course site and the demo read it, so writing a
row is an upsert keyed by ``stage`` rather than an append, and a metric that has not been
measured yet (``delta_cp``, ``diversity``) is written as ``null`` instead of being invented.

Two numbers carry their definition with them rather than a footnote somewhere else: legality is
written twice (``argmax``, the headline, and ``sampled``, the demo's own setting), and an Elo
whose games were all wins or all losses is printed as a one-sided bound, in words, instead of a
zero-width "95 % CI".
"""

from __future__ import annotations

import json
from datetime import UTC, datetime
from pathlib import Path
from typing import TYPE_CHECKING, Any

from pydantic import BaseModel, ConfigDict

from rukh.eval.elo import EloResult

if TYPE_CHECKING:  # pragma: no cover - typing only
    from rukh.eval.suite import SuiteResult

REPORT_NAME = "report.md"
RESULTS_NAME = "results.json"

class WebRow(BaseModel):
    """One row of the single results table, exactly as the web reads it."""

    model_config = ConfigDict(extra="forbid")

    stage: str
    params: int
    legality: float | None = None
    """Legality of the most likely token: the headline definition."""
    legality_sampled: float | None = None
    """Legality of a token drawn with the suite's temperature and top-k."""
    top1: float | None = None
    top3: float | None = None
    top1_by_band: dict[str, float] = {}
    top3_by_band: dict[str, float] = {}
    puzzles: dict[str, float] = {}
    elo: float | None = None
    elo_ci: list[float] | None = None
    """95 % bootstrap interval, or ``null`` when the results were separated."""
    elo_lower: float | None = None
    elo_upper: float | None = None
    elo_separated: bool = False
    delta_cp: float | None = None
    diversity: float | None = None
    date: str
    run_id: str | None = None

def row_of(result: SuiteResult) -> WebRow:
    """The table row a finished suite produces."""
    elo = result.elo
    interval = (
        [elo.ci_low, elo.ci_high]
        if elo is not None and elo.ci_low is not None and elo.ci_high is not None
        else None
    )
    accuracy = result.accuracy
    return WebRow(
        stage=result.stage,
        params=result.params,
        legality=result.legality_argmax.rate if result.legality_argmax else None,
        legality_sampled=result.legality_sampled.rate if result.legality_sampled else None,
        top1=accuracy.top1 if accuracy else None,
        top3=accuracy.top3 if accuracy else None,
        top1_by_band={band.band: band.top1 for band in accuracy.bands} if accuracy else {},
        top3_by_band={band.band: band.top3 for band in accuracy.bands} if accuracy else {},
        puzzles=result.puzzles.by_band() if result.puzzles else {},
        elo=elo.elo if elo else None,
        elo_ci=interval,
        elo_lower=elo.elo_lower if elo else None,
        elo_upper=elo.elo_upper if elo else None,
        elo_separated=bool(elo.separated) if elo else False,
        delta_cp=result.delta_cp,
        diversity=result.diversity,
        date=result.date,
        run_id=result.run_id,
    )

def upsert_row(path: Path, row: WebRow) -> list[dict[str, Any]]:
    """Insert or replace ``row`` in the shared results file and return every row."""
    path = Path(path)
    rows: list[dict[str, Any]] = []
    if path.is_file():
        try:
            payload = json.loads(path.read_text(encoding="utf-8"))
        except ValueError:
            payload = {}
        found = payload.get("rows") if isinstance(payload, dict) else payload
        if isinstance(found, list):
            rows = [item for item in found if isinstance(item, dict)]
    new = row.model_dump()
    rows = [item for item in rows if item.get("stage") != row.stage]
    rows.append(new)
    rows.sort(key=lambda item: str(item.get("stage")))
    path.parent.mkdir(parents=True, exist_ok=True)
    document = {"updated_at": datetime.now(UTC).isoformat(timespec="seconds"), "rows": rows}
    path.write_text(
        json.dumps(document, indent=2, ensure_ascii=False) + "\n", encoding="utf-8", newline="\n"
    )
    return rows

def _percent(value: float | None) -> str:
    return "n/a" if value is None else f"{value * 100:.1f} %"


def elo_line(elo: EloResult | None) -> str:
    """The Elo cell: a point estimate with an interval, or with a one-sided bound."""
    if elo is None:
        return "n/a"
    if not elo.separated and elo.ci_low is not None and elo.ci_high is not None:
        return f"{elo.elo:.0f} (95 % CI {elo.ci_low:.0f}-{elo.ci_high:.0f})"
    if elo.elo_lower is not None:
        return f"> {elo.elo_lower:.0f} (one-sided 95 % bound; every game won)"
    if elo.elo_upper is not None:
        return f"< {elo.elo_upper:.0f} (one-sided 95 % bound; every game lost)"
    return f"{elo.elo:.0f} (no interval)"