Cabezas y casillas: tres preguntas y una segunda entrada

"""The three supervised heads of M3, all of them reading one vector: ``PositionEncoder.pool``.

The heads are deliberately tiny — one linear layer each — because that is what makes the
experiment mean something. If a linear probe on the pooled representation can say how good a
position is, whether the last move threw the game away and who is going to win, then the
representation already contains those facts and masked move modeling put them there. A deep
head would be able to learn them by itself and would tell us nothing about the encoder.

* ``ValueHead``: one scalar through ``tanh``, matching the ``tanh(cp / 400)`` label, so the
  output is bounded and a mate cannot dominate the loss.
* ``BlunderHead``: one logit, because the label is missing for a good part of the rows (a
  position whose predecessor is not in the table) and a masked binary cross entropy is the
  honest way to skip them.
* ``ResultHead``: three classes, White / draw / Black.

``MultiHead`` trains them together over a shared encoder, with a weight per head: the three
tasks have different scales (an MSE in ``[0, 4]``, two cross entropies) and different amounts of
data, so "just add them up" would silently make one of them the only one that matters.
"""

from __future__ import annotations

from typing import Literal

import torch
from torch import Tensor, nn
from torch.nn import functional as F

from rukh.config import BaseConfig
from rukh.models.encoder import PositionEncoder

HEADS = ("value", "blunder", "result")
RESULT_CLASSES = 3

class HeadWeights(BaseConfig):
    """Weight of each head in the joint loss; ``0`` switches a head off without removing it."""

    value: float = 1.0
    blunder: float = 1.0
    result: float = 0.5

class ValueHead(nn.Module):
    """How good is this position for White: a scalar in ``(-1, 1)``."""

    def __init__(self, d_model: int) -> None:
        super().__init__()
        self.proj = nn.Linear(d_model, 1)

    def forward(self, pooled: Tensor) -> Tensor:
        return torch.tanh(self.proj(pooled)).squeeze(-1)


class BlunderHead(nn.Module):
    """Did the move that led here throw the game away: one logit (not a probability)."""

    def __init__(self, d_model: int) -> None:
        super().__init__()
        self.proj = nn.Linear(d_model, 1)

    def forward(self, pooled: Tensor) -> Tensor:
        return self.proj(pooled).squeeze(-1)


class ResultHead(nn.Module):
    """How the game ended: logits over White / draw / Black."""

    def __init__(self, d_model: int) -> None:
        super().__init__()
        self.proj = nn.Linear(d_model, RESULT_CLASSES)

    def forward(self, pooled: Tensor) -> Tensor:
        return self.proj(pooled)

class MultiHead(nn.Module):
    """An encoder plus the three heads, trained on one pooled representation."""

    def __init__(
        self,
        encoder: PositionEncoder,
        weights: HeadWeights | None = None,
        pooling: Literal["cls", "mean"] = "mean",
    ) -> None:
        super().__init__()
        self.encoder = encoder
        self.weights = weights or HeadWeights()
        self.pooling = pooling
        d_model = encoder.cfg.d_model
        self.value = ValueHead(d_model)
        self.blunder = BlunderHead(d_model)
        self.result = ResultHead(d_model)

    def pooled(self, idx: Tensor, attention_mask: Tensor | None = None) -> Tensor:
        """The position's representation, ``(B, d_model)``."""
        return self.encoder.pool(idx, self.pooling, attention_mask)

    def forward(self, idx: Tensor, attention_mask: Tensor | None = None) -> dict[str, Tensor]:
        """``{"value": (B,), "blunder": (B,) logits, "result": (B, 3) logits}``."""
        return self.from_pooled(self.pooled(idx, attention_mask))

    def from_pooled(self, pooled: Tensor) -> dict[str, Tensor]:
        """The three heads' outputs for an already pooled batch."""
        return {
            "value": self.value(pooled),
            "blunder": self.blunder(pooled),
            "result": self.result(pooled),
        }

    def loss(
        self, outputs: dict[str, Tensor], targets: dict[str, Tensor]
    ) -> tuple[Tensor, dict[str, Tensor]]:
        """``(total, per head)``; a head with no labelled row in the batch contributes zero.

        ``targets`` carries ``value`` (float), ``blunder`` (0/1 float) with its ``blunder_mask``
        (the rows that have a label at all) and ``result`` (class index).
        """
        zero = torch.zeros((), device=outputs["value"].device, dtype=outputs["value"].dtype)
        parts = {
            "value": F.mse_loss(outputs["value"], targets["value"].to(outputs["value"].dtype)),
            "blunder": zero,
            "result": F.cross_entropy(outputs["result"], targets["result"].long()),
        }
        mask = targets["blunder_mask"].bool()
        if bool(mask.any()):
            parts["blunder"] = F.binary_cross_entropy_with_logits(
                outputs["blunder"][mask], targets["blunder"][mask].to(outputs["blunder"].dtype)
            )
        total = sum(
            (getattr(self.weights, name) * parts[name] for name in HEADS),
            start=torch.zeros_like(zero),
        )
        return total, parts

    def head_parameters(self) -> list[nn.Parameter]:
        """Every parameter of the three heads, the ones a ``probe`` run is allowed to move."""
        return [
            param for head in (self.value, self.blunder, self.result) for param in head.parameters()
        ]

"""Tests for the three heads and their staged fine-tuning: gradients, freezing and the curve."""

from __future__ import annotations

from pathlib import Path

import polars as pl
import pytest
import torch

from rukh.models import EncoderConfig, MultiHead, PositionEncoder
from rukh.models.heads import HEADS, BlunderHead, HeadWeights, ResultHead, ValueHead
from rukh.models.squares import SQUARE_TOKENS, fen_to_tokens
from rukh.train import HeadsConfig, freeze_encoder, label_curve, train_heads
from rukh.train.checkpoint import load_checkpoint, save_checkpoint
from rukh.train.heads import LabelledPositions, build_frames, set_training_mode

pytestmark = pytest.mark.unit

TOY = EncoderConfig(input="squares", n_layer=3, n_head=2, d_model=32, dropout=0.0)
# Four boards whose evaluation the encoder can actually read off the pieces: the full set, the
# same without Black's queen, without White's queen, and without either. A toy task has to be
# learnable or "the loss went down" measures nothing.
BOARDS = (
    "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR",
    "rnb1kbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR",
    "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNB1KBNR",
    "rnb1kbnr/pppppppp/8/8/8/8/PPPPPPPP/RNB1KBNR",
)
BOARD_CP = (0, 300, -300, 0)
BOARD_RESULT = ("1/2-1/2", "1-0", "0-1", "1/2-1/2")
WHITE_TO_MOVE = f"{BOARDS[0]} w KQkq -"

def toy_batch(rows: int = 4) -> dict[str, torch.Tensor]:
    idx = torch.tensor([fen_to_tokens(WHITE_TO_MOVE) for _ in range(rows)])
    return {
        "idx": idx,
        "value": torch.linspace(-0.9, 0.9, rows),
        "blunder": torch.tensor([1.0, 0.0] * (rows // 2)),
        "blunder_mask": torch.tensor([True, False] * (rows // 2)),
        "result": torch.tensor([0, 1, 2, 0][:rows]),
    }

def test_each_head_has_the_shape_its_label_needs() -> None:
    pooled = torch.randn(5, 32)
    value = ValueHead(32)(pooled)
    assert value.shape == (5,) and bool((value.abs() < 1.0).all())
    assert BlunderHead(32)(pooled).shape == (5,)  # a logit, not a probability
    assert ResultHead(32)(pooled).shape == (5, 3)


def test_multihead_answers_the_three_questions_at_once() -> None:
    model = MultiHead(toy_encoder())
    batch = toy_batch()
    outputs = model(batch["idx"])
    assert set(outputs) == set(HEADS)
    assert outputs["value"].shape == (4,)
    assert outputs["blunder"].shape == (4,)
    assert outputs["result"].shape == (4, 3)
    assert model.pooled(batch["idx"]).shape == (4, 32)
    assert model.encoder.cfg.seq == SQUARE_TOKENS

def test_multihead_propagates_the_three_gradients() -> None:
    model = MultiHead(toy_encoder(), HeadWeights(value=1.0, blunder=1.0, result=1.0))
    batch = toy_batch()
    total, parts = model.loss(model(batch["idx"]), batch)
    assert set(parts) == set(HEADS)
    assert all(torch.isfinite(part) for part in parts.values())
    total.backward()
    for name in HEADS:
        grad = getattr(model, name).proj.weight.grad
        assert grad is not None and float(grad.abs().sum()) > 0.0, name
    body = model.encoder.blocks[0].mlp.fc.weight.grad
    assert body is not None and float(body.abs().sum()) > 0.0  # the encoder gets them too

def test_a_head_with_weight_zero_stops_contributing() -> None:
    model = MultiHead(toy_encoder(), HeadWeights(value=1.0, blunder=0.0, result=0.0))
    batch = toy_batch()
    total, parts = model.loss(model(batch["idx"]), batch)
    assert torch.allclose(total, parts["value"], atol=1e-6)
    total.backward()
    assert float(model.blunder.proj.weight.grad.abs().sum()) == 0.0
    assert float(model.value.proj.weight.grad.abs().sum()) > 0.0


def test_a_batch_without_a_blunder_label_costs_nothing_instead_of_nan() -> None:
    model = MultiHead(toy_encoder())
    batch = toy_batch()
    batch["blunder_mask"] = torch.zeros_like(batch["blunder_mask"])
    total, parts = model.loss(model(batch["idx"]), batch)
    assert float(parts["blunder"]) == 0.0
    assert torch.isfinite(total)

"""The ``squares`` input scheme: a FEN as 69 fixed tokens, the didactic rival of ``moves``.

The ``moves`` scheme feeds the encoder the game so far (the decoder's own UCI vocabulary); this
one feeds it the position itself, so the two can be compared on the same heads. That comparison
is the point of the module: it is the "what is a good representation?" lesson of M3.

Layout (always exactly ``SQUARE_TOKENS`` = 69 positions, so no padding is ever needed)::

    0        <cls>                 pooling anchor, the counterpart of the decoder's <bos>
    1..64    the 64 squares        piece or <empty>, file-major (a1, a2, ..., a8, b1, ..., h8)
    65       side to move          turn:w / turn:b
    66       castling rights       castle:<KQkq subset>, 16 combinations in one token
    67       en passant            ep:none or the file, ep:a ... ep:h
    68       halfmove clock        clock:0 ... clock:3, the bucketed 50-move counter

The square order is ``rukh.tokenize.uci_vocab.squares()``, file-major, the same enumeration the
UCI vocabulary uses for move endpoints: one ordering for the whole project.

The plan's arithmetic (64 + turn + 4 castling + en passant + clock) adds up to 71, not to the 69
it fixes; 69 is the binding number, so the four castling rights travel in one token with 16
values (they are four bits of one fact, and this keeps every other component in a slot of its
own) and the position that frees up became ``<cls>``, which gives ``pool("cls")`` a real anchor
in this scheme instead of borrowing square a1.

Halfmove-clock buckets, chosen so that the boundaries mean something over the board:

===========  ===================================================================
``clock:0``  0-5 half-moves: a pawn moved or a piece was captured very recently
``clock:1``  6-24: a normal manoeuvring stretch
``clock:2``  25-49: the 50-move rule is in sight and shapes the plan
``clock:3``  50 or more: a draw can be claimed
===========  ===================================================================

A four-field FEN (``fen4``, the key of the P1 positions and evaluations) carries no counters, so
it reads back as ``clock:0``; that is the value every label built from P1 data will have.

``SQUARE_VOCAB`` is fixed by construction, never learned from data, and ``vocab_hash`` pins it:
a change to the enumeration invalidates every encoder trained with it, exactly like the UCI
vocabulary of P1.
"""

from __future__ import annotations

import hashlib
import json

from rukh.tokenize.uci_vocab import FILES, squares

PIECES = "PNBRQKpnbrqk"
CASTLING_ORDER = "KQkq"
N_SQUARES = 64
SQUARE_TOKENS = 69
"""Length of every ``squares`` sequence: 1 + 64 + 4."""

CLOCK_EDGES = (6, 25, 50)
"""Upper edges (exclusive) of the halfmove-clock buckets; see the module docstring."""


def castling_strings() -> list[str]:
    """The 16 castling combinations in a fixed order: ``-``, ``K``, ``Q``, ``KQ``, ``k``, ..."""
    out: list[str] = []
    for mask in range(16):
        rights = "".join(right for bit, right in enumerate(CASTLING_ORDER) if mask >> bit & 1)
        out.append(rights or "-")
    return out


def build_square_vocab() -> list[str]:
    """The token list in id order; see the module docstring for the layout it serves."""
    tokens = ["<pad>", "<mask>", "<cls>", "<empty>"]
    tokens.extend(PIECES)
    tokens.extend(["turn:w", "turn:b"])
    tokens.extend(f"castle:{rights}" for rights in castling_strings())
    tokens.append("ep:none")
    tokens.extend(f"ep:{file}" for file in FILES)
    tokens.extend(f"clock:{index}" for index in range(len(CLOCK_EDGES) + 1))
    return tokens


SQUARE_VOCAB: list[str] = build_square_vocab()
SQUARE_IDS: dict[str, int] = {token: index for index, token in enumerate(SQUARE_VOCAB)}
SQUARE_VOCAB_SIZE = len(SQUARE_VOCAB)

PAD_ID = SQUARE_IDS["<pad>"]
MASK_ID = SQUARE_IDS["<mask>"]
CLS_ID = SQUARE_IDS["<cls>"]
EMPTY_ID = SQUARE_IDS["<empty>"]
CONTROL_IDS: frozenset[int] = frozenset({PAD_ID, MASK_ID, CLS_ID})
"""Never masked by ``apply_masking``: they are the frame of the sequence, not a prediction."""

def vocab_hash() -> str:
    """SHA-256 of the vocabulary as a JSON list: the identity of the ``squares`` scheme."""
    payload = json.dumps(SQUARE_VOCAB, ensure_ascii=True, separators=(",", ":"))
    return hashlib.sha256(payload.encode("ascii")).hexdigest()


def clock_bucket(halfmove: int) -> int:
    """Index of the halfmove-clock bucket of ``halfmove`` (see ``CLOCK_EDGES``)."""
    if halfmove < 0:
        raise ValueError(f"the halfmove clock cannot be negative, got {halfmove}")
    return sum(halfmove >= edge for edge in CLOCK_EDGES)

def _placement_ids(placement: str) -> list[int]:
    """The 64 square ids, file-major, from the piece-placement field of a FEN."""
    ranks = placement.split("/")
    if len(ranks) != 8:
        raise ValueError(f"a FEN placement needs 8 ranks, got {len(ranks)}: {placement!r}")
    ids = [EMPTY_ID] * N_SQUARES
    for row, rank_text in enumerate(ranks):
        rank = 7 - row  # the placement is written from rank 8 down to rank 1
        file = 0
        for char in rank_text:
            if char.isdigit():
                file += int(char)
            elif char in PIECES:
                if file >= 8:
                    raise ValueError(f"rank {rank + 1} of {placement!r} is too long")
                ids[file * 8 + rank] = SQUARE_IDS[char]
                file += 1
            else:
                raise ValueError(f"unknown piece {char!r} in {placement!r}")
        if file != 8:
            raise ValueError(f"rank {rank + 1} of {placement!r} covers {file} files, not 8")
    return ids

def _castling_id(field: str) -> int:
    """Token id of a FEN castling field, normalized to the ``KQkq`` order."""
    if field in ("-", ""):
        return SQUARE_IDS["castle:-"]
    unknown = [char for char in field if char not in CASTLING_ORDER]
    if unknown:
        raise ValueError(f"unsupported castling field {field!r} (Chess960 is out of scope)")
    rights = "".join(right for right in CASTLING_ORDER if right in field)
    return SQUARE_IDS[f"castle:{rights}"]

def fen_to_tokens(fen: str) -> list[int]:
    """The ``SQUARE_TOKENS`` ids of a FEN; four fields (``fen4``) or the full six are accepted.

    Raises ``ValueError`` on a malformed FEN: the encoder must never be fed a position that was
    silently repaired into a different one.

    Token 68, the halfmove clock, is **constant** over every dataset this project builds. A
    ``fen4`` has no counters, so it reads back as ``clock:0``, and P1's positions table is all
    ``fen4``: the bucket carries exactly zero information there and the encoder learns a bias
    term for it. It stays in the layout because a live position from the demo does have the
    counters, and a representation that changes shape between training and serving would be
    worse than one slot of constant.
    """
    fields = fen.split()
    if len(fields) < 4:
        raise ValueError(f"a FEN needs at least 4 fields, got {len(fields)}: {fen!r}")
    placement, turn, castling, ep = fields[:4]
    if turn not in ("w", "b"):
        raise ValueError(f"the side to move must be 'w' or 'b', got {turn!r}")
    halfmove = int(fields[4]) if len(fields) > 4 else 0
    if ep in ("-", ""):
        ep_token = "ep:none"
    elif len(ep) == 2 and ep[0] in FILES and ep[1] in "36":
        ep_token = f"ep:{ep[0]}"
    else:
        raise ValueError(f"unknown en-passant square {ep!r}")
    return [
        CLS_ID,
        *_placement_ids(placement),
        SQUARE_IDS[f"turn:{turn}"],
        _castling_id(castling),
        SQUARE_IDS[ep_token],
        SQUARE_IDS[f"clock:{clock_bucket(halfmove)}"],
    ]

def tokens_to_strings(tokens: list[int]) -> list[str]:
    """The token strings of a sequence, for debugging and for the lesson's figures."""
    return [SQUARE_VOCAB[token] for token in tokens]


def tokens_to_fen(tokens: list[int]) -> str:
    """The four-field FEN a token sequence encodes: the inverse of ``fen_to_tokens``.

    The halfmove-clock bucket is dropped because it is lossy by construction, so the round trip
    is exact for a ``fen4`` and exact up to the counters for a full FEN.
    """
    if len(tokens) != SQUARE_TOKENS:
        raise ValueError(f"a squares sequence has {SQUARE_TOKENS} tokens, got {len(tokens)}")
    names = tokens_to_strings(tokens)
    if names[0] != "<cls>":
        raise ValueError(f"a squares sequence starts with <cls>, got {names[0]!r}")
    board = names[1 : 1 + N_SQUARES]
    ranks: list[str] = []
    for rank in range(7, -1, -1):
        row, empty = "", 0
        for file in range(8):
            piece = board[file * 8 + rank]
            if piece == "<empty>":
                empty += 1
                continue
            if piece not in PIECES:
                raise ValueError(f"{piece!r} is not a piece token")
            row += (str(empty) if empty else "") + piece
            empty = 0
        ranks.append(row + (str(empty) if empty else ""))
    turn = names[1 + N_SQUARES].removeprefix("turn:")
    castling = names[2 + N_SQUARES].removeprefix("castle:")
    ep_file = names[3 + N_SQUARES].removeprefix("ep:")
    ep = "-" if ep_file == "none" else f"{ep_file}{'6' if turn == 'w' else '3'}"
    return f"{'/'.join(ranks)} {turn} {castling} {ep}"


def square_name(index: int) -> str:
    """Name of the board square at position ``index`` of the 64-square block (0 = ``a1``)."""
    return squares()[index]

"""Tests for rukh.models.squares: the 69-token FEN scheme, its vocabulary and the round trip."""

from __future__ import annotations

import chess
import pytest

from rukh.models.squares import (
    CLOCK_EDGES,
    SQUARE_IDS,
    SQUARE_TOKENS,
    SQUARE_VOCAB,
    SQUARE_VOCAB_SIZE,
    castling_strings,
    clock_bucket,
    fen_to_tokens,
    square_name,
    tokens_to_fen,
    tokens_to_strings,
    vocab_hash,
)
from rukh.tokenize.uci_vocab import squares

pytestmark = pytest.mark.unit

# Pinned like the UCI vocabulary: an encoder is only meaningful with the vocabulary it saw.
SQUARE_VOCAB_SHA256 = "b0f530895c035c07423dfa246ad5e874c7406ffbaa1b2155601ebf09e89c98e9"
OPENING = ["e2e4", "c7c5", "g1f3", "d7d6", "d2d4", "c5d4", "f3d4", "g8f6", "b1c3", "a7a6"]


def played(moves: list[str]) -> chess.Board:
    board = chess.Board()
    for move in moves:
        board.push_uci(move)
    return board


def test_the_vocabulary_is_stable() -> None:
    assert vocab_hash() == SQUARE_VOCAB_SHA256
    assert len(SQUARE_VOCAB) == SQUARE_VOCAB_SIZE == 47
    assert len(set(SQUARE_VOCAB)) == SQUARE_VOCAB_SIZE
    assert SQUARE_VOCAB[:4] == ["<pad>", "<mask>", "<cls>", "<empty>"]
    assert SQUARE_VOCAB[4:16] == list("PNBRQKpnbrqk")
    assert len(castling_strings()) == 16
    assert castling_strings()[0] == "-" and castling_strings()[-1] == "KQkq"

def test_a_position_is_exactly_69_tokens() -> None:
    board = played(OPENING)
    assert len(fen_to_tokens(board.fen())) == SQUARE_TOKENS == 69
    fen4 = " ".join(board.fen().split()[:4])
    assert len(fen_to_tokens(fen4)) == SQUARE_TOKENS  # the P1 four-field FEN is accepted too


def test_the_layout_is_cls_board_turn_castling_en_passant_clock() -> None:
    names = tokens_to_strings(fen_to_tokens(chess.Board().fen()))
    assert names[0] == "<cls>"
    assert names[65:] == ["turn:w", "castle:KQkq", "ep:none", "clock:0"]
    assert names[1] == "R" and names[2] == "P" and names[3] == "<empty>"  # a1, a2, a3


def test_every_square_matches_python_chess() -> None:
    board = chess.Board()
    for move in OPENING:
        board.push_uci(move)
        names = tokens_to_strings(fen_to_tokens(board.fen()))[1:65]
        for index, name in enumerate(names):
            piece = board.piece_at(chess.square(index // 8, index % 8))
            assert name == (piece.symbol() if piece else "<empty>"), square_name(index)

def test_the_64_squares_round_trip_through_the_four_field_fen() -> None:
    board = chess.Board()
    for move in [*OPENING, "f1e2", "c8d7", "e1g1"]:
        board.push_uci(move)
        fen4 = " ".join(board.fen().split()[:4])
        assert tokens_to_fen(fen_to_tokens(fen4)) == fen4
        assert tokens_to_fen(fen_to_tokens(board.fen())) == fen4


def test_castling_rights_are_distinguished() -> None:
    board = played(OPENING)
    fen4 = " ".join(board.fen().split()[:4])
    fields = fen4.split()
    variants = {
        rights: fen_to_tokens(f"{fields[0]} {fields[1]} {rights} {fields[3]}")[66]
        for rights in ("KQkq", "KQk", "Kq", "-")
    }
    assert len(set(variants.values())) == 4
    assert variants["-"] == SQUARE_IDS["castle:-"]
    # The order inside the field does not matter, the rights do.
    assert fen_to_tokens(f"{fields[0]} {fields[1]} qK {fields[3]}")[66] == variants["Kq"]
    with pytest.raises(ValueError, match="unsupported castling field"):
        fen_to_tokens(f"{fields[0]} {fields[1]} AHah {fields[3]}")

def test_en_passant_is_distinguished_and_does_not_disturb_the_board() -> None:
    board = played(["e2e4", "a7a6", "e4e5", "d7d5"])  # d6 is a legal en-passant target
    with_ep = fen_to_tokens(board.fen())
    without = fen_to_tokens(board.fen().replace(" d6 ", " - "))
    assert with_ep[67] == SQUARE_IDS["ep:d"] and without[67] == SQUARE_IDS["ep:none"]
    assert with_ep[:67] == without[:67]  # only the en-passant slot changes
    assert tokens_to_fen(with_ep).split()[3] == "d6"
    assert fen_to_tokens("8/8/8/8/8/8/8/8 b - c3")[67] == SQUARE_IDS["ep:c"]
    assert tokens_to_fen(fen_to_tokens("8/8/8/8/8/8/8/8 b - c3")).split()[3] == "c3"
    with pytest.raises(ValueError, match="en-passant square"):
        fen_to_tokens("8/8/8/8/8/8/8/8 w - e4")


def test_the_halfmove_clock_is_bucketed() -> None:
    assert CLOCK_EDGES == (6, 25, 50)
    assert [clock_bucket(n) for n in (0, 5, 6, 24, 25, 49, 50, 120)] == [0, 0, 1, 1, 2, 2, 3, 3]
    start = chess.Board().fen().split()
    fen = " ".join(start[:4])
    assert fen_to_tokens(f"{fen} 0 1")[68] == SQUARE_IDS["clock:0"]
    assert fen_to_tokens(f"{fen} 30 40")[68] == SQUARE_IDS["clock:2"]
    assert fen_to_tokens(fen)[68] == SQUARE_IDS["clock:0"]  # a fen4 carries no counter
    with pytest.raises(ValueError, match="cannot be negative"):
        clock_bucket(-1)

def test_the_square_order_is_the_uci_vocabulary_order() -> None:
    assert square_name(0) == "a1" and square_name(63) == "h8"
    assert [square_name(i) for i in range(64)] == squares()


def test_a_malformed_fen_is_refused() -> None:
    with pytest.raises(ValueError, match="at least 4 fields"):
        fen_to_tokens("8/8/8/8/8/8/8/8 w")
    with pytest.raises(ValueError, match="8 ranks"):
        fen_to_tokens("8/8/8 w - -")
    with pytest.raises(ValueError, match="covers 7 files"):
        fen_to_tokens("7/8/8/8/8/8/8/8 w - -")
    with pytest.raises(ValueError, match="unknown piece"):
        fen_to_tokens("xxxxxxxx/8/8/8/8/8/8/8 w - -")
    with pytest.raises(ValueError, match="side to move"):
        fen_to_tokens("8/8/8/8/8/8/8/8 x - -")
    with pytest.raises(ValueError, match="69 tokens"):
        tokens_to_fen([1, 2, 3])