Entrenar las cabezas: etiquetas, congelación y la curva

"""The supervised dataset of M3: value, blunder and result labels from the P1 evaluations.

Everything here comes out of ``data/evals/positions-eval.parquet``, the table P1 built by
crossing the sampled positions with ``Lichess/chess-position-evaluations`` (see
``rukh.data.evals``): one row per distinct four-field FEN, with the game it first came from
(``game_id``, ``ply``, ``last_move``, ``result``) and Stockfish's verdict on it (``best_move``,
``cp``, ``mate``, and every line in ``pvs``). No column is invented here; the mate convention
and the sign of a score are ``rukh.data.scoring``, shared with the DPO pairs.

Three labels:

``value``
    ``tanh(score / 400)`` from **White's** point of view, so the number is a position's value
    and not a player's, and a mate is ``±1`` rather than an outlier of ten thousand. Also
    published as five buckets (``< -200``, ``-200..-50``, ``-50..50``, ``50..200``, ``> 200``
    centipawns) so a regression and a classification can be compared on the same rows.

``blunder``
    whether ``last_move`` — the move that *led to* this position — threw away at least
    ``blunder_cp`` centipawns. It needs the previous position of the same game, and it is
    measured from the point of view of the side that moved: its best line there against what it
    actually got here, which is minus the value of this position for the opponent. A position
    whose predecessor is not in the table (the first ply, or a FEN that was deduplicated into
    another game) has no blunder label at all, ``null``, never ``0``.

``result``
    how the game the position came from ended: 0 White, 1 draw, 2 Black.

``game_moves`` is the one thing here that leaves ``positions-eval.parquet``: it reads the P1
games back so the ``moves`` scheme of the encoder can be fine-tuned on the line that led to a
position instead of on the position itself. It is deliberately **not** part of ``build_labels``,
because the ``squares`` path must not pay for a join it never uses.

The split is by ``game_id`` and never by position. Two positions of the same game are not
independent: the second one is the first one plus a move, and a model that memorised the game
would score well on both. With a per-position split that leak is invisible in the metrics and
fatal in the demo, so ``test_labels`` proves no ``game_id`` is ever on both sides.
"""

from __future__ import annotations

import zlib
from pathlib import Path

import polars as pl
from pydantic import Field

from rukh.config import BaseConfig
from rukh.data.manifest import FileHash, Manifest
from rukh.data.scoring import SCORE_WHITE_SQL
from rukh.data.uci import sha256_file
from rukh.paths import resolve

LABELS_FILE = "positions-labels.parquet"
GAMES_GLOB = "year=*/month=*/games.parquet"
"""The partition layout ``rukh.data.uci`` writes; one scan reads every month at once."""
RESULT_CLASSES: dict[str, int] = {"1-0": 0, "1/2-1/2": 1, "0-1": 2}
"""Result of the game the position came from: White wins, draw, Black wins."""
VALUE_EDGES: tuple[int, int, int, int] = (-200, -50, 50, 200)
"""Centipawn edges of the five value buckets, from White's point of view."""
SPLITS = ("train", "val")
SOURCE_COLUMNS = ["fen", "game_id", "ply", "last_move", "result", "phase", "cp", "mate"]

class LabelsConfig(BaseConfig):
    """Where the evaluations are, how a blunder is defined and how the split is drawn."""

    positions_eval: str = "data/evals/positions-eval.parquet"
    games_dir: str = "data/uci"
    """Where ``game_moves`` looks for the P1 games, ``year=*/month=*/games.parquet``.

    Only the ``moves`` scheme reads it: the ``squares`` path never touches ``data/uci``."""
    out_dir: str = "data/labels"
    value_scale: float = Field(default=400.0, gt=0.0)
    """``tanh(cp / value_scale)``: 400 centipawns (four pawns, a bit less than a rook) is 0.76."""
    blunder_cp: int = Field(default=100, ge=1)
    val_fraction: float = Field(default=0.1, gt=0.0, lt=1.0)
    seed: int = 42

def game_split(game_id: int, val_fraction: float, seed: int) -> str:
    """Which side of the split a game falls on: a pure function of its id and the seed.

    CRC-32 rather than ``hash()`` or a dataframe hash: it is the same number in every process,
    every polars version and every future run, which is what makes a split reproducible.
    """
    digest = zlib.crc32(f"{seed}:{game_id}".encode())
    return "val" if (digest % 1_000_000) / 1_000_000 < val_fraction else "train"


def row_order(game_id: int, ply: int, seed: int) -> int:
    """A stable per-row key: the order the label-count curve grows its nested subsets in."""
    return zlib.crc32(f"{seed}:{game_id}:{ply}".encode())

def _score_white() -> pl.Expr:
    """Centipawns from White's point of view, with a mate in ``n`` worth ``±(10000 - n)``.

    The convention is not restated here: ``rukh.data.scoring.SCORE_WHITE_SQL`` is the one the
    DPO pairs and the evaluation consolidation already run, and ``pl.sql_expr`` turns it into a
    polars expression, so a change to the mate convention cannot reach one consumer and miss
    this one.
    """
    return pl.sql_expr(SCORE_WHITE_SQL).cast(pl.Int64)


def _value_bucket(score: pl.Expr) -> pl.Expr:
    """Index of the five-bucket classification variant of ``value``."""
    bucket = pl.lit(0, dtype=pl.Int8)
    for index, edge in enumerate(VALUE_EDGES, start=1):
        bucket = pl.when(score >= edge).then(pl.lit(index, dtype=pl.Int8)).otherwise(bucket)
    return bucket

def build_labels(cfg: LabelsConfig, frame: pl.DataFrame | None = None) -> pl.DataFrame:
    """The labelled table; ``frame`` overrides reading ``cfg.positions_eval`` (used by tests)."""
    source = frame if frame is not None else pl.read_parquet(resolve(cfg.positions_eval))
    rows = (
        source.select(SOURCE_COLUMNS)
        .with_columns(
            score_white=_score_white(),
            mover=pl.when(pl.col("fen").str.split(" ").list.get(1) == "w").then(1).otherwise(-1),
        )
        .drop_nulls("score_white")
        .with_columns(score_mover=pl.col("mover") * pl.col("score_white"))
    )
    # The move into this position is judged against the best line of the position before it, and
    # from the mover's side: what it could have had there, minus what it actually got here
    # (which is minus the value of this position for the opponent, now to move).
    previous = rows.select(
        pl.col("game_id"),
        (pl.col("ply") + 1).alias("ply"),
        pl.col("score_mover").alias("before_best"),
    )
    joined = rows.join(previous, on=["game_id", "ply"], how="left").with_columns(
        loss_cp=(pl.col("before_best") + pl.col("score_mover")).cast(pl.Int64)
    )

    games = joined.select("game_id").unique().sort("game_id")
    split = pl.DataFrame(
        {
            "game_id": games["game_id"],
            "split": [game_split(int(g), cfg.val_fraction, cfg.seed) for g in games["game_id"]],
        }
    )
    labelled = (
        joined.join(split, on="game_id", how="left")
        .with_columns(
            value=(pl.col("score_white") / cfg.value_scale).tanh().cast(pl.Float32),
            value_bucket=_value_bucket(pl.col("score_white")),
            result_class=pl.col("result")
            .replace_strict(RESULT_CLASSES, default=None)
            .cast(pl.Int8),
            blunder=pl.when(pl.col("loss_cp").is_null())
            .then(None)
            .otherwise((pl.col("loss_cp") >= cfg.blunder_cp).cast(pl.Int8)),
        )
        .drop_nulls("result_class")
    )
    order = [
        row_order(int(game), int(ply), cfg.seed)
        for game, ply in zip(labelled["game_id"], labelled["ply"], strict=True)
    ]
    return labelled.with_columns(order=pl.Series("order", order, dtype=pl.UInt32)).select(
        "game_id",
        "ply",
        "fen",
        "last_move",
        "phase",
        "split",
        "value",
        "value_bucket",
        "blunder",
        "result_class",
        pl.col("score_white").alias("cp"),
        "loss_cp",
        "order",
    )

GAME_COLUMNS = ("game_id", "uci", "white_elo", "black_elo")


def game_moves(frame: pl.DataFrame, games_dir: str = "data/uci") -> pl.DataFrame:
    """The move prefix source of ``frame``: one row per **game**, never per position.

    ``build_labels`` deliberately knows nothing about this: the ``squares`` scheme reads a FEN
    and nothing else, and joining two months of games into it would make the cheap path pay for
    the expensive one. The ``moves`` scheme of ``rukh.train.heads`` calls this instead, and it
    scans ``games.parquet`` lazily and semi-joins on the labels' distinct ``game_id``s, so only
    the games that actually carry a label are ever read.

    The caveat that travels with the result: the supervised table is deduplicated by ``fen4``,
    so the row's ``game_id`` is the game that *first* reached that position. The prefix built
    from it is **a** line reaching the position, not necessarily the one the labelled game
    played — the position is the same, its history may not be.
    """
    wanted = frame.select(pl.col("game_id").unique()).lazy()
    pattern = (resolve(games_dir) / GAMES_GLOB).as_posix()
    return (
        pl.scan_parquet(pattern)
        .select(*GAME_COLUMNS)
        .join(wanted, on="game_id", how="semi")
        .unique(subset=["game_id"], keep="first")
        .collect()
    )

def label_subsets(frame: pl.DataFrame, fractions: list[float]) -> dict[float, pl.DataFrame]:
    """Nested subsets for the label-count curve: 10 % is inside 25 %, inside 50 %, inside 100 %.

    Nested on purpose. With independent samples, a dip at 50 % could just be an unlucky draw and
    the curve would measure sampling noise instead of the value of more labels; growing one set
    means every point is the previous one plus new rows.
    """
    if any(not 0.0 < fraction <= 1.0 for fraction in fractions):
        raise ValueError(f"every fraction must be in (0, 1], got {fractions}")
    ordered = frame.sort(["order", "game_id", "ply"])
    total = ordered.height
    return {
        fraction: ordered.head(max(1, round(total * fraction)))
        for fraction in sorted(set(fractions))
    }


def counts(frame: pl.DataFrame) -> dict[str, int]:
    """Row counts worth recording: total, per split and how many carry a blunder label."""
    per_split = {name: int(frame.filter(pl.col("split") == name).height) for name in SPLITS}
    return {
        "positions": int(frame.height),
        **per_split,
        "blunder_labelled": int(frame.drop_nulls("blunder").height),
        "blunders": int(frame.filter(pl.col("blunder") == 1).height),
    }

def run(cfg: LabelsConfig) -> Manifest:
    """Write ``out_dir/positions-labels.parquet`` plus the manifest, and return the manifest."""
    frame = build_labels(cfg)
    out_dir = resolve(cfg.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    target = out_dir / LABELS_FILE
    frame.write_parquet(target, compression="zstd")
    source = resolve(cfg.positions_eval)
    manifest = Manifest(
        dataset="rukh-positions-eval",
        months=[],
        filters={
            "positions_eval": Path(cfg.positions_eval).as_posix(),
            # The labels are a pure function of this file and of the settings below, so its
            # digest is what says whether two label tables are the same table.
            "positions_eval_sha256": sha256_file(source) if source.is_file() else None,
            "value_scale": cfg.value_scale,
            "value_edges": list(VALUE_EDGES),
            "blunder_cp": cfg.blunder_cp,
            "split": {"by": "game_id", "val_fraction": cfg.val_fraction, "seed": cfg.seed},
        },
        counts=counts(frame),
        files=[FileHash(path=LABELS_FILE, sha256=sha256_file(target), bytes=target.stat().st_size)],
    )
    (out_dir / "manifest.json").write_text(manifest.model_dump_json(indent=2) + "\n", "utf-8")
    return manifest

"""Staged fine-tuning of the three heads, and the curve that asks how many labels are needed.

Three modes, in the order the lesson walks through them:

``probe``
    the encoder is frozen and only the three linear heads move. It answers the only question
    that matters about the pretraining: *is the information already in the representation?* The
    encoder's weights come out of a probe run bit for bit identical, and a test proves it.
``last-n``
    the last ``last_n`` blocks and the final norm are unfrozen. The early layers keep the general
    features, the late ones specialise.
``full``
    everything moves. Best numbers, most labels needed, easiest to overfit.

On top of that, ``label_curve`` trains the same recipe on 10 %, 25 %, 50 % and 100 % of the
training labels, on **nested** subsets (see ``rukh.data.labels.label_subsets``), so the curve
measures the value of more labels and not the luck of a draw. That curve is the module's lesson:
labelling is the expensive part of this project, and the curve says where it stops paying.

Both input schemes can be fine-tuned here, which is what makes the spec's comparison possible:

``squares``
    ``data/evals/positions-eval.parquet`` holds a FEN per position, and 69 fixed tokens come
    straight out of it.
``moves``
    the position is the *line that reached it*: ``rukh.data.labels.game_moves`` joins the P1
    games back in and the item is ``[<bos>, <wXXXX>, <bXXXX>] + uci.split()[:ply]``, cropped
    header-first by ``rukh.infer.sampler.prompt_ids`` and padded per batch. This is the scheme
    the masked-move pretraining runs on, so it is the only one an MMM checkpoint can be loaded
    into — before, the pretrained encoder had nowhere to go.

    The caveat of that path: the supervised table is deduplicated by ``fen4``, so the prefix is
    **a** line reaching the position, not necessarily the labelled game's. The position, the
    value and the blunder verdict are the same either way; the history may not be.
"""

from __future__ import annotations

import logging
import math
import time
from collections.abc import Callable, Sequence
from pathlib import Path
from typing import Literal

import polars as pl
import torch
from pydantic import Field
from torch import Tensor, nn
from torch.utils.data import DataLoader, Dataset

from rukh.config import BaseConfig
from rukh.data.labels import LabelsConfig, build_labels, game_moves, label_subsets
from rukh.infer.sampler import HEADER_TOKENS, prompt_ids
from rukh.models import EncoderConfig, PositionEncoder
from rukh.models.encoder import PAD_ID
from rukh.models.heads import HEADS, HeadWeights, MultiHead
from rukh.models.squares import fen_to_tokens
from rukh.tokenize.uci_vocab import UciTokenizer, elo_token
from rukh.train.checkpoint import BEST_NAME, CURVE_KEY, attach_payload, load_encoder, step_name
from rukh.train.common import (
    RunConfig,
    autocast_for,
    forever,
    load_resume,
    log_metrics,
    maybe_compile,
    param_groups,
    pick_device,
    run_dir,
    set_lr,
    write_checkpoint,
)

log = logging.getLogger(__name__)

Mode = Literal["probe", "last-n", "full"]
Scheme = Literal["squares", "moves"]
Item = dict[str, Tensor]
CURVE = (0.1, 0.25, 0.5, 1.0)

class LabelledPositions(Dataset[Item]):
    """The labelled rows as tensors; one item is a position and its three labels.

    The tokens are built on demand rather than up front: a handful of small integers per row is
    cheap to build and expensive to keep for millions of rows. On ``squares`` that is the 69
    tokens of the FEN; on ``moves`` it is the header plus the first ``ply`` moves of the game,
    which is why ``moves`` needs the ``(game_id, uci)`` frame of
    ``rukh.data.labels.game_moves`` — stored per **game**, never copied per row.
    """

    def __init__(
        self,
        frame: pl.DataFrame,
        scheme: Scheme = "squares",
        moves: pl.DataFrame | None = None,
        block: int = 200,
    ):
        self.scheme = scheme
        self.block = block
        self.value = frame["value"].to_numpy().astype("float32")
        self.result = frame["result_class"].to_numpy().astype("int64")
        self.blunder_mask = frame["blunder"].is_not_null().to_numpy()
        self.blunder = frame["blunder"].fill_null(0).to_numpy().astype("float32")
        self.fens: list[str] = frame["fen"].to_list()
        self.game_ids: list[int] = [int(value) for value in frame["game_id"].to_list()]
        self.plies: list[int] = [int(value) for value in frame["ply"].to_list()]
        self.prefixes: dict[int, list[int]] = {}
        if scheme == "moves":
            if moves is None:
                raise ValueError(
                    "the 'moves' scheme needs the games of the labelled rows: pass the frame "
                    "of rukh.data.labels.game_moves(frame, cfg.labels.games_dir)"
                )
            self.prefixes = _move_prefixes(moves)
            unknown = {game for game in self.game_ids if game not in self.prefixes}
            if unknown:
                raise ValueError(
                    f"{len(unknown)} labelled games are missing from the games frame "
                    f"(first: {sorted(unknown)[:3]}); check labels.games_dir"
                )

    def __len__(self) -> int:
        return len(self.fens)

    def tokens(self, index: int) -> list[int]:
        """The input ids of one row, in the scheme this dataset was built for."""
        if self.scheme == "squares":
            return fen_to_tokens(self.fens[index])
        prefix = self.prefixes[self.game_ids[index]]
        # ``ply`` counts the move that *led to* this position, so the prefix includes it.
        ids = prefix[: HEADER_TOKENS + self.plies[index]]
        # A prefix longer than the context keeps <bos> and the two Elo tokens and drops the
        # oldest moves: the same header-first crop the sampler uses, so the fine-tuned model
        # sees the shape the pretraining and the demo do.
        return prompt_ids(ids, self.block)

    def __getitem__(self, index: int) -> Item:
        return {
            "idx": torch.tensor(self.tokens(index), dtype=torch.long),
            "value": torch.tensor(float(self.value[index])),
            "blunder": torch.tensor(float(self.blunder[index])),
            "blunder_mask": torch.tensor(bool(self.blunder_mask[index])),
            "result": torch.tensor(int(self.result[index])),
        }

def _move_prefixes(moves: pl.DataFrame) -> dict[int, list[int]]:
    """``game_id -> [<bos>, elo(white), elo(black), *every move]``, one entry per game."""
    tok = UciTokenizer()
    out: dict[int, list[int]] = {}
    for game_id, uci, white_elo, black_elo in moves.select(
        "game_id", "uci", "white_elo", "black_elo"
    ).rows():
        ids = [
            tok.bos_id,
            tok.vocab[elo_token(int(white_elo), "w")],
            tok.vocab[elo_token(int(black_elo), "b")],
        ]
        ids.extend(tok.vocab.get(move, tok.unk_id) for move in str(uci).split())
        out[int(game_id)] = ids
    return out

def collate(items: list[Item]) -> Item:
    """Pad a batch to its longest sequence and say which tokens are real.

    ``squares`` items are all 69 tokens long and the mask is all ``True``; ``moves`` items are
    as long as the game was, so the short ones are padded with ``<pad>`` and the mask keeps the
    encoder's attention off them. The padding is never silent: ``attention_mask`` travels with
    the batch and ``PositionEncoder.pool`` averages the real tokens only.
    """
    width = max(int(item["idx"].shape[0]) for item in items)
    idx = torch.full((len(items), width), PAD_ID, dtype=torch.long)
    mask = torch.zeros((len(items), width), dtype=torch.bool)
    for row, item in enumerate(items):
        tokens = item["idx"]
        idx[row, : tokens.shape[0]] = tokens
        mask[row, : tokens.shape[0]] = True
    batch: Item = {"idx": idx, "attention_mask": mask}
    for key in ("value", "blunder", "blunder_mask", "result"):
        # The labels are optional so that inference (``rukh eval encoder``) can pad a batch of
        # bare token sequences with this very function instead of a second copy of it.
        if key in items[0]:
            batch[key] = torch.stack([item[key] for item in items])
    return batch

def make_label_loader(
    dataset: LabelledPositions,
    batch_size: int,
    seed: int = 0,
    workers: int = 0,
    shuffle: bool = True,
) -> DataLoader[Item]:
    """A ``DataLoader`` whose shuffling is fully determined by ``seed``."""
    generator = torch.Generator()
    generator.manual_seed(seed)
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        drop_last=False,
        num_workers=workers,
        generator=generator,
        persistent_workers=workers > 0,
        collate_fn=collate,
    )

class HeadsConfig(RunConfig):
    """One fine-tuning run of the three heads; unknown keys in the YAML are an error."""

    labels: LabelsConfig = Field(default_factory=LabelsConfig)
    encoder_ckpt: str | None = None
    """The masked-move checkpoint to start from; ``None`` trains the encoder from scratch,
    which is the honest baseline the pretraining has to beat."""
    model: EncoderConfig | None = None  # only used when there is no checkpoint to load
    input: Scheme = "squares"
    """The representation the heads are fine-tuned on; it must match ``encoder_ckpt``'s."""
    pooling: Literal["cls", "mean"] = "mean"
    mode: Mode = "probe"
    last_n: int = Field(default=2, ge=1)
    fraction: float = Field(default=1.0, gt=0.0, le=1.0)
    curve: list[float] = list(CURVE)
    weights: HeadWeights = Field(default_factory=HeadWeights)

    @property
    def default_name(self) -> str:
        return f"heads-{self.mode}"

    def encoder(self) -> EncoderConfig:
        """The encoder config used when no checkpoint is loaded."""
        base = self.model if self.model is not None else EncoderConfig()
        return base.model_copy(update={"input": self.input, "block": self.block})

class HeadsResult(BaseConfig):
    """What one fine-tuning run produced: where the weights are and how they scored."""

    checkpoint: str
    mode: Mode
    fraction: float
    train_labels: int
    val_labels: int
    metrics: dict[str, float]

def freeze_encoder(model: MultiHead, mode: Mode, last_n: int = 2) -> int:
    """Apply a fine-tuning mode and return how many encoder tensors stay trainable."""
    for param in model.encoder.parameters():
        param.requires_grad = mode == "full"
    if mode == "last-n":
        blocks = list(model.encoder.blocks)
        for block in blocks[max(0, len(blocks) - last_n) :]:
            for param in block.parameters():
                param.requires_grad = True
        for param in model.encoder.ln_f.parameters():
            param.requires_grad = True
    return sum(param.requires_grad for param in model.encoder.parameters())


def set_training_mode(model: MultiHead, mode: Mode, last_n: int = 2) -> None:
    """Training mode for the whole model, with every **frozen** block kept in ``eval``.

    A frozen block in ``train`` mode would still apply dropout, so the same position would give
    the heads a different vector every epoch: noise the heads cannot learn away and the block
    cannot absorb, because it is not learning at all. Under ``last-n`` that applies to the
    frozen prefix too — the early blocks and the embeddings go to ``eval`` and only the last
    ``last_n`` blocks and the final norm keep their dropout, which is the regularisation of the
    part that is actually being trained.
    """
    model.train()
    if mode == "probe":
        model.encoder.eval()
    elif mode == "last-n":
        encoder = model.encoder
        encoder.eval()
        for block in list(encoder.blocks)[max(0, len(encoder.blocks) - last_n) :]:
            block.train()
        encoder.ln_f.train()

def build_frames(cfg: HeadsConfig) -> tuple[pl.DataFrame, pl.DataFrame]:
    """``(train, val)`` label frames, split by ``game_id`` and cut to ``cfg.fraction``."""
    frame = build_labels(cfg.labels)
    train = frame.filter(pl.col("split") == "train")
    if cfg.fraction < 1.0:
        train = label_subsets(train, [cfg.fraction])[cfg.fraction]
    return train, frame.filter(pl.col("split") == "val")


def build_datasets(
    cfg: HeadsConfig, train_frame: pl.DataFrame, val_frame: pl.DataFrame
) -> tuple[LabelledPositions, LabelledPositions]:
    """The two datasets of a run; on ``moves`` the games are read once and shared."""
    moves = None
    if cfg.input == "moves":
        moves = game_moves(pl.concat([train_frame, val_frame]), cfg.labels.games_dir)
        log.info("%d games joined back for the 'moves' scheme", moves.height)
    return (
        LabelledPositions(train_frame, cfg.input, moves, cfg.block),
        LabelledPositions(val_frame, cfg.input, moves, cfg.block),
    )

def load_encoder_for(cfg: HeadsConfig, where: torch.device) -> PositionEncoder:
    """The pretrained encoder of ``encoder_ckpt``, or a fresh one when there is none."""
    if cfg.encoder_ckpt is None:
        return PositionEncoder(cfg.encoder())
    from rukh import paths

    encoder, _payload = load_encoder(paths.resolve(cfg.encoder_ckpt), map_location=where)
    if encoder.cfg.input != cfg.input:
        raise ValueError(
            f"{cfg.encoder_ckpt} was trained on the {encoder.cfg.input!r} scheme, "
            f"the run asks for {cfg.input!r}"
        )
    return encoder

def to_device(batch: Item, where: torch.device) -> Item:
    return {key: value.to(where) for key, value in batch.items()}


def _compile_probe(
    model: MultiHead, dataset: LabelledPositions, where: torch.device, batch_size: int
) -> Callable[[nn.Module], None]:
    """One real step, so ``torch.compile`` runs here and not in the middle of the loop.

    The probe uses the dataset's own items, mask included: probing with ``attention_mask=None``
    while the loop always passes one would compile a graph the loop never runs and pay for a
    recompilation at step 1.
    """

    def probe(compiled: nn.Module) -> None:
        rows = [dataset[index] for index in range(min(batch_size, len(dataset)))]
        batch = to_device(collate(rows), where)
        outputs = compiled(batch["idx"], batch["attention_mask"])
        total, _ = model.loss(outputs, batch)
        total.backward()

    return probe

@torch.no_grad()
def evaluate_heads(
    model: MultiHead, loader: DataLoader[Item], batches: int, where: torch.device
) -> dict[str, float]:
    """Loss and one metric per head over at most ``batches`` validation batches."""
    was_training = model.training
    model.eval()
    totals = {f"val/{name}_loss": 0.0 for name in HEADS}
    seen = 0
    value_error = 0.0
    blunder_hits = blunder_seen = 0
    result_hits = 0
    for index, raw in enumerate(loader):
        if index >= batches:
            break
        batch = to_device(raw, where)
        outputs = model(batch["idx"], batch.get("attention_mask"))
        total, parts = model.loss(outputs, batch)
        rows = int(batch["idx"].shape[0])
        seen += rows
        totals["val/loss"] = totals.get("val/loss", 0.0) + float(total) * rows
        for name in HEADS:
            totals[f"val/{name}_loss"] += float(parts[name]) * rows
        value_error += float((outputs["value"] - batch["value"]).abs().sum())
        mask = batch["blunder_mask"].bool()
        if bool(mask.any()):
            predicted = (outputs["blunder"][mask] > 0).float()
            blunder_hits += int((predicted == batch["blunder"][mask]).sum())
            blunder_seen += int(mask.sum())
        result_hits += int((outputs["result"].argmax(dim=-1) == batch["result"]).sum())
    if was_training:
        model.train()
    if not seen:
        return {"val/loss": math.nan}
    metrics = {key: value / seen for key, value in totals.items()}
    metrics["val/value_mae"] = value_error / seen
    metrics["val/blunder_acc"] = blunder_hits / blunder_seen if blunder_seen else math.nan
    metrics["val/result_acc"] = result_hits / seen
    return metrics

def train_heads(
    cfg: HeadsConfig, resume: Path | None = None, device: str | None = None
) -> HeadsResult:
    """Fine-tune the three heads over a (possibly frozen) encoder and return the run's result."""
    cfg.check()
    torch.manual_seed(cfg.seed)
    where = torch.device(device or pick_device())
    train_frame, val_frame = build_frames(cfg)
    if not train_frame.height:
        raise ValueError("the training split is empty; check labels.val_fraction and the source")
    train_set, val_set = build_datasets(cfg, train_frame, val_frame)
    train_loader = make_label_loader(train_set, cfg.batch_size, cfg.seed, cfg.workers)
    val_loader = make_label_loader(val_set, cfg.batch_size, cfg.seed, 0, shuffle=False)

    encoder = load_encoder_for(cfg, where)
    model = MultiHead(encoder, cfg.weights, cfg.pooling).to(where)
    trainable = freeze_encoder(model, cfg.mode, cfg.last_n)
    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:
        start_step, best_val, run_id = load_resume(resume, model, optimizer, where)
    autocast, _ = autocast_for(cfg, where)
    with autocast:  # compile under the same precision the loop will use
        probe = _compile_probe(model, train_set, where, cfg.batch_size)
        runnable = maybe_compile(model, cfg.compile, probe=probe)
    set_training_mode(model, cfg.mode, cfg.last_n)

    out_dir = run_dir(cfg, resume)
    out_dir.mkdir(parents=True, exist_ok=True)
    batches = forever(train_loader)
    final = out_dir / step_name(cfg.max_steps)
    metrics: dict[str, float] = {}

    from rukh.tracking import start_run

    params = {
        # ``tokens_dir`` is the one inherited knob this loop has no use for: the labels are a
        # parquet of positions, not a packed token stream.
        **cfg.model_dump(mode="json", exclude={"tokens_dir"}),
        "device": str(where),
        "train_labels": train_frame.height,
        "val_labels": val_frame.height,
        "trainable_encoder_tensors": trainable,
        "num_params": sum(p.numel() for p in model.parameters()),
    }
    tags = {"objective": "heads", "mode": cfg.mode, "fraction": str(cfg.fraction)}

    with start_run(out_dir.name, params, tags=tags, run_id=run_id) as run:
        log.info("run %s in %s on %s (%s)", run.info.run_id, out_dir, where, cfg.mode)
        this_run = str(run.info.run_id)

        def save(path: Path, step: int) -> Path:
            return write_checkpoint(
                path,
                step=step,
                model=model,
                optimizer=optimizer,
                cfg=cfg,
                model_cfg=encoder.cfg,
                vocab_hash=None,
                manifest_sha=None,
                best_val=best_val,
                run_id=this_run,
            )

        clock = time.perf_counter()
        for step in range(start_step, cfg.max_steps):
            lr = set_lr(optimizer, step, cfg)
            optimizer.zero_grad(set_to_none=True)
            step_losses = {name: 0.0 for name in HEADS}
            step_total = 0.0
            for _ in range(cfg.grad_accum):
                batch = to_device(next(batches), where)
                with autocast:
                    outputs = runnable(batch["idx"], batch.get("attention_mask"))
                    total, parts = model.loss(outputs, batch)
                (total / cfg.grad_accum).backward()
                step_total += float(total.detach()) / cfg.grad_accum
                for name in HEADS:
                    step_losses[name] += float(parts[name].detach()) / cfg.grad_accum
            grad_norm = float(
                nn.utils.clip_grad_norm_(
                    [p for p in model.parameters() if p.requires_grad], cfg.grad_clip
                )
            )
            optimizer.step()

            done = step + 1
            last = done == cfg.max_steps
            if done % cfg.log_every == 0 or last:
                elapsed = max(time.perf_counter() - clock, 1e-9)
                steps = min(cfg.log_every, done - start_step)
                log_metrics(
                    {
                        "train/loss": step_total,
                        **{f"train/{name}_loss": step_losses[name] for name in HEADS},
                        "lr": lr,
                        "grad_norm": grad_norm,
                        "positions_per_s": cfg.batch_size * cfg.grad_accum * steps / elapsed,
                    },
                    step=done,
                )
                clock = time.perf_counter()

            if done % cfg.eval_every == 0 or last:
                metrics = evaluate_heads(model, val_loader, cfg.eval_batches, where)
                set_training_mode(model, cfg.mode, cfg.last_n)
                log_metrics(metrics, step=done)
                val_loss = metrics.get("val/loss", math.nan)
                log.info("step %d  val/loss %.4f", done, val_loss)
                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 HeadsResult(
        checkpoint=str(final),
        mode=cfg.mode,
        fraction=cfg.fraction,
        train_labels=train_frame.height,
        val_labels=val_frame.height,
        metrics=metrics,
    )

def label_curve(
    cfg: HeadsConfig, fractions: list[float] | None = None, device: str | None = None
) -> list[HeadsResult]:
    """Train the same recipe on nested subsets of the labels; one result per fraction.

    The points are written back into the checkpoints of the **last** run (the one with the most
    labels, which is the checkpoint anybody would then evaluate or publish), so
    ``rukh eval encoder`` can render the "labels needed" table without being told where the
    other runs are. Without that write the curve would exist only in MLflow and in this return
    value, and the table — a ``GOAL.md`` deliverable — could never be filled.
    """
    wanted = sorted(set(fractions if fractions is not None else cfg.curve))
    results: list[HeadsResult] = []
    for fraction in wanted:
        run = cfg.model_copy(
            update={
                "fraction": fraction,
                "run_name": f"{cfg.run_name or cfg.default_name}-{int(fraction * 100)}",
            }
        )
        results.append(train_heads(run, device=device))
    if results:
        write_curve(Path(results[-1].checkpoint).parent, results)
    return results


def write_curve(run_dir: Path, results: Sequence[HeadsResult]) -> list[Path]:
    """Record the curve in every checkpoint of ``run_dir``; returns the files it touched."""
    points = [result.model_dump(mode="json") for result in results]
    written: list[Path] = []
    for path in sorted(run_dir.glob("*.pt")):
        attach_payload(path, **{CURVE_KEY: points})
        written.append(path)
    log.info("label curve of %d points written into %d checkpoints", len(points), len(written))
    return written

BEST_NAME = "best.pt"
CURVE_KEY = "label_curve"
"""Where ``rukh.train.heads.label_curve`` writes its points and ``rukh eval encoder`` reads them.

The curve is the lesson of M3 ("how many labels does this actually take?") and a ``GOAL.md``
deliverable, so it travels **inside** the checkpoint: a number that only ever existed in a log
line is a number nobody can put in a table."""
TIED_HEAD = "lm_head.weight"
"""Tied to ``tokens.weight``; a published state dict leaves it out and it is re-tied on load."""
TIED_HEADS = frozenset({TIED_HEAD, "mlm_head.weight", "encoder.mlm_head.weight"})
"""Every tied head in the project: the decoder's and the encoder's masked-move head.

The last name is the same head seen from a ``MultiHead``, which holds the encoder as a child."""
TIED_SOURCES: dict[str, str] = {
    TIED_HEAD: "tokens.weight",
    "mlm_head.weight": "tokens.weight",
    "encoder.mlm_head.weight": "encoder.tokens.weight",
}
"""The embedding each tied head shares its storage with, when the tie is switched on."""

def attach_payload(path: Path, **extra: Any) -> Path:
    """Add plain values to an existing checkpoint, rewriting it atomically.

    Used for facts a run only knows once it is over — the label-count curve is the whole of it —
    so they do not need a second file next to the weights that can be lost or go stale.
    """
    path = Path(path)
    payload = load_checkpoint(path)
    payload.update(extra)
    tmp = path.with_suffix(path.suffix + ".tmp")
    torch.save(payload, tmp)
    tmp.replace(path)
    return path

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


def build_heads(payload: Mapping[str, Any]) -> MultiHead:
    """A ``MultiHead`` shaped by a payload, with no weights loaded yet."""
    run_cfg = payload.get("cfg") or {}
    encoder = PositionEncoder(EncoderConfig.model_validate(payload["model_cfg"]))
    pooling = run_cfg.get("pooling") or "mean"
    if pooling not in ("cls", "mean"):
        raise ValueError(f"unknown pooling {pooling!r} in the checkpoint's config")
    return MultiHead(encoder, HeadWeights.model_validate(run_cfg.get("weights") or {}), pooling)


def load_heads(
    path: Path, map_location: str | torch.device = "cpu"
) -> tuple[MultiHead, dict[str, Any]]:
    """Rebuild a fine-tuned ``MultiHead`` (encoder plus the three heads) from a checkpoint.

    ``rukh.train.heads`` saves the whole ``MultiHead`` state under the *encoder's* ``model_cfg``,
    because the heads have no shape of their own beyond ``d_model``; the pooling and the head
    weights come from the run's config, which travels in the same payload.
    """
    payload = load_checkpoint(path, map_location=map_location)
    model = build_heads(payload)
    load_state(model, payload["model_state"])
    return model.eval(), payload

def checkpoint_kind(payload: Mapping[str, Any]) -> str:
    """``"encoder"`` or ``"decoder"``: what kind of model a checkpoint describes.

    Read off the weights rather than off a flag nobody wrote: a ``MultiHead`` keeps the encoder
    as a child (``encoder.*``), a bare ``PositionEncoder`` is identified by the ``input`` field
    only its config has, and everything else is the decoder.
    """
    state = payload.get("model_state") or {}
    if any(str(name).startswith("encoder.") for name in state):
        return "encoder"
    return "encoder" if "input" in (payload.get("model_cfg") or {}) else "decoder"


def load_any(
    path: Path, map_location: str | torch.device = "cpu"
) -> tuple[nn.Module, dict[str, Any], str]:
    """Rebuild whatever model a checkpoint describes: ``(model, payload, kind)``.

    ``rukh export`` and ``rukh publish`` take a checkpoint and have to work out what is in it;
    this is the one place that decides, so the two commands can never disagree.
    """
    payload = load_checkpoint(path, map_location=map_location)
    kind = checkpoint_kind(payload)
    if kind == "decoder":
        model: nn.Module = MoveDecoder(DecoderConfig.model_validate(payload["model_cfg"]))
    elif any(str(name).startswith("encoder.") for name in payload["model_state"]):
        model = build_heads(payload)
    else:
        model = PositionEncoder(EncoderConfig.model_validate(payload["model_cfg"]))
    load_state(model, payload["model_state"])
    return model.eval(), payload, kind

from rukh.train.common import (
    RunConfig,
    forever,
    maybe_compile,
    param_groups,
    pick_device,
    run_dir,
    skip_batches,
)
from rukh.train.heads import (
    HeadsConfig,
    HeadsResult,
    evaluate_heads,
    freeze_encoder,
    label_curve,
    train_heads,
)
from rukh.train.loop import TrainConfig, evaluate, train
from rukh.train.mmm import MaskingConfig, MmmConfig, apply_masking, evaluate_mmm, train_mmm

from rukh.train.schedule import lr_at

__all__ = [
    "BEST_NAME",
    "TIED_HEAD",
    "TIED_HEADS",
    "TIED_SOURCES",
    "HeadsConfig",
    "HeadsResult",
    "MaskingConfig",
    "MmmConfig",
    "RunConfig",
    "TrainConfig",
    "apply_masking",
    "build_heads",
    "checkpoint_kind",
    "evaluate",
    "evaluate_heads",
    "evaluate_mmm",
    "forever",
    "freeze_encoder",
    "label_curve",
    "load_any",
    "load_checkpoint",
    "load_encoder",
    "load_heads",
    "load_model",
    "load_state",
    "lr_at",
    "maybe_compile",
    "param_groups",
    "pick_device",
    "read_manifest_sha",
    "read_vocab_hash",
    "restore",
    "run_dir",
    "save_checkpoint",
    "skip_batches",
    "step_name",
    "train",
    "train_heads",
    "train_mmm",
]

# Fine-tuning of the three heads (value, blunder, result) on the `squares` scheme: the position
# itself, 69 fixed tokens. The default is `probe`: the encoder is frozen and only the three
# linear heads learn, which is the run that answers whether the representation already carries
# the information. `--mode last-n` and `--mode full` are the next two stages, and `--curve` runs
# the whole thing on 10/25/50/100 % of the labels and writes the points into the checkpoint.
#
# `encoder_ckpt` stays null here: masked move modeling only runs on `moves` (there is no packed
# stream of board tokens in P1), so a `squares` encoder starts from random weights. That is the
# baseline the pretrained `moves` run of `encoder-heads-moves.yaml` is compared against, and
# pointing this file at an MMM checkpoint is refused rather than silently reinterpreted.
encoder_ckpt: null
input: squares # the labels are FENs, one row per position; see encoder-heads-moves.yaml
pooling: mean
mode: probe
last_n: 2
fraction: 1.0
curve: [0.1, 0.25, 0.5, 1.0]
weights:
  # The three tasks have different scales and different amounts of data; `result` is the noisiest
  # (one label per game, repeated over every position of it), so it counts for half.
  value: 1.0
  blunder: 1.0
  result: 0.5
labels:
  positions_eval: data/evals/positions-eval.parquet
  out_dir: data/labels
  value_scale: 400.0
  blunder_cp: 100
  val_fraction: 0.1
  seed: 42
block: 200 # ignored by the squares scheme, which always uses its 69 positions
batch_size: 256
grad_accum: 1
lr: 1.0e-3 # a frozen encoder with three linear heads takes a much larger rate than pretraining
min_lr_ratio: 0.1
warmup: 200
max_steps: 4000
weight_decay: 0.01
betas: [0.9, 0.95]
grad_clip: 1.0
precision: bf16
compile: false # three linear heads over 69 fixed positions: compiling costs more than it saves
eval_every: 250
eval_batches: 50
ckpt_every: 1000
out_dir: checkpoints
seed: 42
run_name: encoder-heads
unique_run_name: true
workers: 4
log_every: 10

encoder_ckpt: checkpoints/encoder-mmm-20260919-093554//best.pt
input: moves

  games_dir: data/uci # only this scheme reads it: the prefixes come from the P1 games

block: 200 # the decoder's context, so a pretrained encoder fits without reshaping its positions
batch_size: 128 # half of `squares`: a prefix is up to 200 tokens, a board is always 69

compile: false # variable-length prefixes recompile on every new shape; not worth it here

def position(
    game_id: int,
    ply: int,
    cp: int | None = 0,
    mate: int | None = None,
    result: str = "1-0",
) -> dict[str, object]:
    """One row of ``positions-eval.parquet``; the side to move follows the ply, as in a game."""
    return {
        # A ply-N position is the one *after* the Nth half-move, so White is to move on even
        # plies: the move that led to an even-ply position was Black's.
        "fen": WHITE_TO_MOVE if ply % 2 == 0 else BLACK_TO_MOVE,
        "game_id": game_id,
        "ply": ply,
        "last_move": "e2e4",
        "result": result,
        "phase": "middlegame",
        "cp": cp,
        "mate": mate,
        "n_seen": 1,
        "best_move": "e2e4",
        "depth": 20,
    }

def toy_frame() -> pl.DataFrame:
    """Two games whose evaluations swing by the same 300 cp, once for each side."""
    return rows_to_frame(
        [
            # Black moved into ply 2 and handed White 300 centipawns: a blunder.
            position(10, 1, cp=0),
            position(10, 2, cp=300),
            # White moved into ply 3 and gained the same 300: the very same swing, not a blunder.
            position(11, 2, cp=0),
            position(11, 3, cp=300),
        ]
    )


def test_a_blunder_is_measured_from_the_side_that_moved() -> None:
    frame = build_labels(LabelsConfig(), toy_frame()).sort(["game_id", "ply"])
    by_key = {(int(r["game_id"]), int(r["ply"])): r for r in frame.to_dicts()}

    black_blundered = by_key[(10, 2)]
    assert black_blundered["loss_cp"] == 300
    assert black_blundered["blunder"] == 1

    white_gained = by_key[(11, 3)]
    assert white_gained["loss_cp"] == -300  # the same swing, the other side to move
    assert white_gained["blunder"] == 0

def test_a_position_without_its_predecessor_has_no_blunder_label() -> None:
    frame = build_labels(LabelsConfig(), toy_frame())
    first = frame.filter((pl.col("game_id") == 10) & (pl.col("ply") == 1)).to_dicts()[0]
    assert first["blunder"] is None and first["loss_cp"] is None
    assert frame.filter(pl.col("blunder").is_null()).height == 2  # one per game
    # A null is not a zero: a position we cannot judge must never train the head as "fine".
    assert frame.drop_nulls("blunder").height == 2


def test_the_blunder_threshold_is_configurable() -> None:
    rows = rows_to_frame([position(1, 2, cp=0), position(1, 3, cp=-120)])
    strict = build_labels(LabelsConfig(blunder_cp=100), rows)
    lenient = build_labels(LabelsConfig(blunder_cp=200), rows)
    assert strict.filter(pl.col("ply") == 3)["blunder"].item() == 1
    assert lenient.filter(pl.col("ply") == 3)["blunder"].item() == 0