El Elo y la suite

"""Elo against Stockfish: play the rungs, then fit the rating that explains the results.

Rungs are Stockfish at ``UCI_Elo`` 1320/1500/1800/2000 (1320 is the engine's floor) plus four
``Skill Level`` rungs below it, so a weak model still has opponents it can score against: with
only losses there is nothing for the fit to latch onto. Both colours are played at every rung.

The fit is the standard logistic (Elo) model with one free parameter,

    P(score) = 1 / (1 + 10 ** ((opponent - elo) / 400)),

maximised by Newton's method: with ``c = ln(10) / 400`` the gradient is ``c * sum(s - p)`` and
the Hessian ``-c**2 * sum(p * (1 - p))``, so each step is ``sum(s - p) / (c * sum(p * (1 - p)))``.
Draws enter as a score of 0.5, which is the usual quasi-likelihood treatment. The 95 % interval
is a percentile bootstrap over the game results, fitted in a single vectorised pass so a
thousand resamples cost milliseconds and no SciPy is needed.

Two honesty notes are built into the numbers rather than left to the reader:

* a game cut short by the context limit is **adjudicated** (shallow engine analysis of the final
  position, or the material count when no engine is around) instead of being booked as a draw,
  because half a point per cut game biases the fit towards the middle of the rungs;
* when every game is a win (or every game is a loss) the likelihood has no maximum inside the
  range of opponents and the bootstrap collapses to a zero-width interval. That is reported as
  ``separated`` with a one-sided likelihood bound instead of a symmetric interval that would be
  a lie.
"""

from __future__ import annotations

import math
from collections.abc import Sequence

import chess
import numpy as np
from pydantic import BaseModel, ConfigDict, model_validator

from rukh.config import BaseConfig
from rukh.eval.cache import EvalCache
from rukh.infer import GameResult, SampleConfig, StockfishOpponent, adjudicate, play_game
from rukh.models import MoveDecoder
from rukh.tokenize.uci_vocab import UciTokenizer

LOG10_400 = math.log(10.0) / 400.0
SUITE = "elo"
MAX_STEP = 400.0
SLACK = 1000.0
"""How far outside the range of opponents the fit is allowed to wander (separation guard)."""
BOUND_SLACK = 4000.0
"""Search range of the one-sided bound when the results are separated."""
MOVE_TIME = 0.1
"""Seconds per move for the engine; below this ``UCI_Elo`` means very little (D-025)."""

class EloRung(BaseConfig):
    """One Stockfish setting and the rating it is assumed to play at."""

    name: str
    elo: int
    uci_elo: int | None = None
    skill: int | None = None

    @model_validator(mode="after")
    def _check(self) -> EloRung:
        if (self.uci_elo is None) == (self.skill is None):
            raise ValueError(f"rung {self.name}: set exactly one of uci_elo and skill")
        return self

# ``Skill Level`` ratings are nominal anchors for the rungs below the engine's 1320 floor: they
# are not measured strengths, so a model whose fit leans on them is reported with that caveat.
DEFAULT_RUNGS: list[EloRung] = [
    EloRung(name="skill-0", elo=800, skill=0),
    EloRung(name="skill-1", elo=950, skill=1),
    EloRung(name="skill-2", elo=1100, skill=2),
    EloRung(name="skill-3", elo=1250, skill=3),
    EloRung(name="uci-1320", elo=1320, uci_elo=1320),
    EloRung(name="uci-1500", elo=1500, uci_elo=1500),
    EloRung(name="uci-1800", elo=1800, uci_elo=1800),
    EloRung(name="uci-2000", elo=2000, uci_elo=2000),
]

class GameRecord(BaseModel):
    """One played game, from the model's point of view."""

    model_config = ConfigDict(extra="forbid")

    rung: str
    opponent_elo: int
    index: int
    model_white: bool
    result: str
    score: float
    plies: int
    illegal_proposals: int
    cut: bool = False
    """The game ran out of context instead of ending on the board."""
    adjudicated: str | None = None
    """How a cut game was decided (``engine depth 8``, ``material count``), or None."""

    def item_id(self) -> str:
        return f"{self.rung}:{self.index}"

class RungResult(BaseModel):
    """Aggregated results at one rung."""

    model_config = ConfigDict(extra="forbid")

    name: str
    opponent_elo: int
    games: int
    wins: int
    draws: int
    losses: int
    score: float
    """Average score in [0, 1]."""
    cut: int = 0
    """Games that hit the context limit instead of ending on the board."""
    adjudicated: int = 0
    """Cut games whose result was decided by adjudication."""

class EloResult(BaseModel):
    """The fitted rating with its interval (or one-sided bound), plus the rung breakdown."""

    model_config = ConfigDict(extra="forbid")

    elo: float
    ci_low: float | None = None
    ci_high: float | None = None
    games: int
    score: float
    rungs: list[RungResult]
    cut: int = 0
    adjudicated: int = 0
    separated: bool = False
    """Every game was a win (or every one a loss): the fit is not identified."""
    elo_lower: float | None = None
    """One-sided 95 % lower bound; set when the model won every game."""
    elo_upper: float | None = None
    """One-sided 95 % upper bound; set when the model lost every game."""

def score_of(result: str, model_white: bool) -> float:
    """Score of a finished game from the model's point of view.

    A cut game (``*``) that could not be adjudicated is the only case left at 0.5 by default;
    ``play_rung`` adjudicates cut games before calling this, so that branch is a fallback.
    """
    if result == "1-0":
        return 1.0 if model_white else 0.0
    if result == "0-1":
        return 0.0 if model_white else 1.0
    return 0.5

def _fit_batch(opponents: np.ndarray, scores: np.ndarray, iterations: int = 60) -> np.ndarray:
    """Newton's method on ``(B, n)`` batches of results; returns one rating per row."""
    opponents = np.asarray(opponents, dtype=np.float64)
    scores = np.asarray(scores, dtype=np.float64)
    low = opponents.min(axis=1) - SLACK
    high = opponents.max(axis=1) + SLACK
    elo = opponents.mean(axis=1)
    for _ in range(iterations):
        p = 1.0 / (1.0 + np.exp(-LOG10_400 * (elo[:, None] - opponents)))
        gradient = (scores - p).sum(axis=1)
        curvature = (p * (1.0 - p)).sum(axis=1)
        safe = np.maximum(curvature, 1e-12)
        step = np.where(curvature > 1e-12, gradient / (LOG10_400 * safe), 0.0)
        step = np.clip(step, -MAX_STEP, MAX_STEP)
        elo = np.clip(elo + step, low, high)
        if np.max(np.abs(step)) < 1e-6:
            break
    return elo

def fit_elo(opponents: Sequence[float], scores: Sequence[float]) -> float:
    """Maximum-likelihood rating for one set of games."""
    if not len(opponents):
        raise ValueError("cannot fit an Elo without games")
    return float(_fit_batch(np.asarray([opponents]), np.asarray([scores]))[0])

def bootstrap_ci(
    opponents: Sequence[float],
    scores: Sequence[float],
    samples: int = 1_000,
    seed: int = 0,
    level: float = 0.95,
) -> tuple[float, float]:
    """Percentile bootstrap interval over the games themselves (one fit per resample)."""
    opponents = np.asarray(opponents, dtype=np.float64)
    scores = np.asarray(scores, dtype=np.float64)
    n = len(opponents)
    if n == 0:
        raise ValueError("cannot bootstrap without games")
    rng = np.random.default_rng(seed)
    index = rng.integers(0, n, size=(samples, n))
    fitted = _fit_batch(opponents[index], scores[index])
    tail = (1.0 - level) / 2.0 * 100.0
    return float(np.percentile(fitted, tail)), float(np.percentile(fitted, 100.0 - tail))

def separation(scores: Sequence[float]) -> str | None:
    """``"wins"``, ``"losses"`` or None: whether every game went the same way.

    With no counter-example the likelihood keeps rising as the rating goes to infinity, the fit
    stops at the clamp and every bootstrap resample returns that same clamp, which is where the
    zero-width "95 % CI" came from.
    """
    values = np.asarray(scores, dtype=np.float64)
    if values.size == 0:
        return None
    if np.all(values >= 1.0):
        return "wins"
    if np.all(values <= 0.0):
        return "losses"
    return None

def one_sided_bound(
    opponents: Sequence[float], scores: Sequence[float], level: float = 0.95, lower: bool = True
) -> float:
    """Likelihood bound for separated results: the rating at which the run stops being plausible.

    For an all-win run the bound is the lowest rating under which winning every game still has
    probability ``1 - level``; for an all-loss run it is the highest rating under which losing
    every game does. Solved by bisection on a monotone log-likelihood, so no SciPy is needed.
    """
    rungs = np.asarray(opponents, dtype=np.float64)
    if rungs.size == 0:
        raise ValueError("cannot bound an Elo without games")
    target = math.log(1.0 - level)

    def loglik(elo: float) -> float:
        p = 1.0 / (1.0 + np.exp(-LOG10_400 * (elo - rungs)))
        p = np.clip(p if lower else 1.0 - p, 1e-300, 1.0)
        return float(np.sum(np.log(p)))

    low = float(rungs.min()) - BOUND_SLACK
    high = float(rungs.max()) + BOUND_SLACK
    # ``loglik`` rises with the rating when bounding from below and falls when bounding from
    # above; bisect for the crossing of ``target`` either way.
    for _ in range(200):
        middle = 0.5 * (low + high)
        if (loglik(middle) < target) == lower:
            low = middle
        else:
            high = middle
    return 0.5 * (low + high)

def summarize_rungs(records: Sequence[GameRecord]) -> list[RungResult]:
    """Wins, draws, losses and average score per rung, ordered by opponent rating."""
    by_rung: dict[str, list[GameRecord]] = {}
    for record in records:
        by_rung.setdefault(record.rung, []).append(record)
    results = [
        RungResult(
            name=name,
            opponent_elo=games[0].opponent_elo,
            games=len(games),
            wins=sum(g.score == 1.0 for g in games),
            draws=sum(g.score == 0.5 for g in games),
            losses=sum(g.score == 0.0 for g in games),
            score=sum(g.score for g in games) / len(games),
            cut=sum(g.cut for g in games),
            adjudicated=sum(g.adjudicated is not None for g in games),
        )
        for name, games in by_rung.items()
    ]
    return sorted(results, key=lambda r: (r.opponent_elo, r.name))

def estimate(
    records: Sequence[GameRecord], samples: int = 1_000, seed: int = 0, level: float = 0.95
) -> EloResult:
    """Fit the rating and its interval (or its one-sided bound) from played games."""
    if not records:
        raise ValueError("cannot estimate an Elo without games")
    opponents = [float(record.opponent_elo) for record in records]
    scores = [record.score for record in records]
    elo = fit_elo(opponents, scores)
    rungs = summarize_rungs(records)
    result = EloResult(
        elo=elo,
        games=len(records),
        score=float(np.mean(scores)),
        rungs=rungs,
        cut=sum(rung.cut for rung in rungs),
        adjudicated=sum(rung.adjudicated for rung in rungs),
    )
    apart = separation(scores)
    if apart is None:
        result.ci_low, result.ci_high = bootstrap_ci(
            opponents, scores, samples=samples, seed=seed, level=level
        )
        return result
    result.separated = True
    bound = one_sided_bound(opponents, scores, level=level, lower=apart == "wins")
    if apart == "wins":
        result.elo_lower = bound
    else:
        result.elo_upper = bound
    return result

def record_of(
    outcome: GameResult,
    rung: EloRung,
    index: int,
    model_white: bool,
    engine: chess.engine.SimpleEngine | None = None,
) -> GameRecord:
    """One played game as a record, adjudicating it first when it was cut short."""
    result = outcome.result
    how: str | None = None
    if outcome.cut:
        result, how = adjudicate(outcome.fen, engine=engine)
    return GameRecord(
        rung=rung.name,
        opponent_elo=rung.elo,
        index=index,
        model_white=model_white,
        result=result,
        score=score_of(result, model_white),
        plies=outcome.plies,
        illegal_proposals=outcome.illegal_proposals,
        cut=outcome.cut,
        adjudicated=how,
    )

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

"""The evaluation suite: one command that produces one row of the results table.

``run_suite`` loads a checkpoint, measures legality, next-move accuracy, puzzles and Elo, and
hands the result to ``report``. Every part is optional at run time: a missing validation
parquet, a missing puzzle parquet or a missing Stockfish binary becomes a note in the report
instead of a crash, because a partial evaluation that says what it could not measure is more
useful than no evaluation at all.

The suite runs on the same device the model was trained on (``pick_device`` by default, so CUDA
when it is there): a harness pinned to the CPU turns ten thousand forward passes and a few
hundred games into hours. The device that was used is recorded in the report.
"""

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 evaluate(
    ckpt: Path | str,
    cfg: EvalConfig,
    suite: str = "full",
    use_cache: bool = True,
    device: str | None = None,
) -> SuiteResult:
    """Measure one checkpoint; missing inputs become notes, never exceptions."""
    from rukh.train import load_model, pick_device

    ckpt = resolve_model(ckpt)
    where = device or cfg.device or pick_device()
    model, _payload = load_model(ckpt, map_location=where)
    model = model.to(where).eval()
    log.info("evaluating %s on %s", ckpt, where)
    tok = UciTokenizer()
    notes: list[str] = [_legality_note(cfg)]
    cache = EvalCache(
        paths.resolve(cfg.cache_db) if use_cache else None,
        file_sha(ckpt),
        enabled=use_cache,
        config_sha=config_sha(cfg.cache_fields()),
    )
    try:
        positions = _positions(cfg, tok, notes)
        for_legality = positions[: cfg.legality_positions]
        result = SuiteResult(
            stage=cfg.stage or ckpt.parent.name,
            suite=suite,
            checkpoint=ckpt.as_posix(),
            model_sha=cache.model_sha,
            params=model.num_params(non_embedding=False),
            date=datetime.now(UTC).date().isoformat(),
            device=str(where),
            legality_argmax=(
                legality(model, tok, for_legality, cfg.sampling(), mode="argmax")
                if positions
                else None
            ),
            legality_sampled=(
                legality(model, tok, for_legality, cfg.sampling(), mode="sampled")
                if positions
                else None
            ),
            accuracy=(
                accuracy(model, tok, positions[: cfg.accuracy_positions]) if positions else None
            ),
            notes=notes,
            config=cfg.model_dump(mode="json"),
        )
        puzzle_path = paths.resolve(cfg.puzzles)
        if puzzle_path.is_file():
            items = load_puzzles(
                puzzle_path, cfg.puzzles_per_band, seed=cfg.seed, split=cfg.puzzle_split
            )
            result.puzzles = run_puzzles(model_source(model, tok), tok, items, cache=cache)
        else:
            notes.append(f"puzzles not found at {puzzle_path}: puzzle suite skipped")
        result.elo = _elo(model, tok, cfg, cache, notes)
        _elo_notes(result.elo, cfg, notes)
        result.notes = notes  # pydantic copied the list at construction time
    finally:
        cache.close()
    return result

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

def render_markdown(result: SuiteResult) -> str:
    """The human-readable report: headline numbers first, then every breakdown."""
    argmax = result.legality_argmax.rate if result.legality_argmax else None
    sampled = result.legality_sampled.rate if result.legality_sampled else None
    sampling = result.legality_sampled
    drawn = (
        ""
        if sampling is None or sampling.temperature is None
        else f" (T={sampling.temperature:g}"
        + ("" if sampling.top_k is None else f", top-k {sampling.top_k}")
        + ")"
    )
    lines: list[str] = [
        f"# Evaluation of `{result.stage}`",
        "",
        f"- Suite: `{result.suite}`",
        f"- Checkpoint: `{result.checkpoint}`",
        f"- Weights SHA-256: `{result.model_sha}`",
        f"- Parameters: {result.params:,}",
        f"- Device: `{result.device}`",
        f"- Date: {result.date}",
        f"- MLflow run: {result.run_id or 'not tracked'}",
        "",
        "## Headline",
        "",
        "| Metric | Value |",
        "|---|---|",
        f"| Legality without the mask, argmax | {_percent(argmax)} |",
        f"| Legality without the mask, sampled{drawn} | {_percent(sampled)} |",
        f"| Top-1 next move | {_percent(result.accuracy.top1 if result.accuracy else None)} |",
        f"| Top-3 next move | {_percent(result.accuracy.top3 if result.accuracy else None)} |",
        f"| Puzzles solved | {_percent(result.puzzles.rate if result.puzzles else None)} |",
        f"| Estimated Elo | {elo_line(result.elo)} |",
    ]
    delta_cp = "n/a" if result.delta_cp is None else f"{result.delta_cp:.1f}"
    diversity = "n/a" if result.diversity is None else f"{result.diversity:.3f}"
    lines.extend(
        [
            f"| Mean centipawn loss | {delta_cp} |",
            f"| Opening diversity | {diversity} |",
            "",
        ]
    )

    if result.legality_argmax is not None or result.legality_sampled is not None:
        lines.extend(
            [
                "## Legality",
                "",
                "Two rates, because they answer different questions. **argmax** is the share of "
                "validation positions whose single most likely token is a legal move, with no "
                "temperature, no top-k and no mask: it is a property of the weights and it is "
                "the definition behind the \u2265 99 % bar of `GOAL.md`. **sampled** draws the "
                "token exactly as the demo does" + drawn.strip() + ", so it is what a player "
                "would meet with the mask switched off, and it is always the lower of the two.",
                "",
                "| Definition | Positions | Legal | Rate |",
                "|---|---:|---:|---:|",
            ]
        )
        for measured in (result.legality_argmax, result.legality_sampled):
            if measured is not None:
                lines.append(
                    f"| {measured.mode} | {measured.positions} | {measured.legal} | "
                    f"{_percent(measured.rate)} |"
                )
        lines.append("")