Los tres vocabularios

"""Tokenization schemes for chess games: fixed UCI vocabulary, char-level SAN and BPE."""

"""Fixed UCI vocabulary: a deterministic enumeration shared with the TypeScript twin.

The vocabulary does not depend on data. It is built by enumerating squares in file-major order
(``a1, a2, ..., a8, b1, ..., h8``), every ``from``/``to`` pair reachable by a queen or a knight
(1 792 moves), every promotion (176) and a fixed block of special tokens in front. Any change to
this enumeration breaks parity with ``rukh-web/src/lib/chess-lm/tokenizer.ts`` and with
``artifacts/tokenizer/fixtures/games.json``.
"""

from __future__ import annotations

import json
from pathlib import Path

FILES = "abcdefgh"
PROMOTION_PIECES = "qrbn"
RESULT_TOKENS: dict[str, str] = {"1-0": "<1-0>", "0-1": "<0-1>", "1/2-1/2": "<1/2>"}
ELO_MIN = 600
ELO_MAX = 3200
ELO_STEP = 100

SPECIALS: list[str] = ["<pad>", "<bos>", "<eos>", "<mask>", "<unk>", "<1-0>", "<0-1>", "<1/2>"]
PAD_ID = 0
BOS_ID = 1
EOS_ID = 2
MASK_ID = 3
UNK_ID = 4

def squares() -> list[str]:
    """The 64 squares in file-major order (index ``file * 8 + rank``)."""
    return [f"{FILES[f]}{r + 1}" for f in range(8) for r in range(8)]


def _reachable(df: int, dr: int) -> bool:
    """Whether a queen (line or diagonal) or a knight can go from one square to another."""
    if df == 0 and dr == 0:
        return False
    if df == 0 or dr == 0 or abs(df) == abs(dr):
        return True
    return {abs(df), abs(dr)} == {1, 2}

def build_moves() -> list[str]:
    """1 792 ``fromto`` strings, ``from`` outer and ``to`` inner, both in file-major order."""
    moves: list[str] = []
    for ff in range(8):
        for fr in range(8):
            for tf in range(8):
                for tr in range(8):
                    if _reachable(tf - ff, tr - fr):
                        moves.append(f"{FILES[ff]}{fr + 1}{FILES[tf]}{tr + 1}")
    return moves


def build_promotions() -> list[str]:
    """176 ``fromto+piece`` strings: rank 7 to 8 (White) then rank 2 to 1 (Black)."""
    promotions: list[str] = []
    for from_rank, to_rank in ((7, 8), (2, 1)):
        for ff in range(8):
            for tf in range(8):
                if abs(tf - ff) <= 1:
                    for piece in PROMOTION_PIECES:
                        promotions.append(f"{FILES[ff]}{from_rank}{FILES[tf]}{to_rank}{piece}")
    return promotions

def elo_tokens(side: str) -> list[str]:
    """``<w0600>`` ... ``<w3200>`` (or ``b``): 27 bins of 100 Elo."""
    return [f"<{side}{elo:04d}>" for elo in range(ELO_MIN, ELO_MAX + 1, ELO_STEP)]


def build_vocab() -> list[str]:
    """The full token list in id order: specials, Elo bins, moves, promotions (2 030 tokens)."""
    return SPECIALS + elo_tokens("w") + elo_tokens("b") + build_moves() + build_promotions()


def elo_bin(elo: int) -> int:
    """Lower edge of the 100-Elo bin, clamped to ``[600, 3200]``."""
    return min(max(int(elo), ELO_MIN), ELO_MAX + ELO_STEP - 1) // ELO_STEP * ELO_STEP


def elo_token(elo: int, side: str) -> str:
    """``<w1800>``-style token for ``elo`` and side ``w`` or ``b``."""
    if side not in ("w", "b"):
        raise ValueError(f"side must be 'w' or 'b', got {side!r}")
    return f"<{side}{elo_bin(elo):04d}>"

class UciTokenizer:
    """Encode/decode games as ``[<bos>, <wXXXX>, <bXXXX>, moves..., <result>, <eos>]``."""

    specials: list[str] = SPECIALS
    pad_id = PAD_ID
    bos_id = BOS_ID
    eos_id = EOS_ID
    mask_id = MASK_ID
    unk_id = UNK_ID

    def __init__(self) -> None:
        self.ids: list[str] = build_vocab()
        self.vocab: dict[str, int] = {token: i for i, token in enumerate(self.ids)}
        if len(self.vocab) != len(self.ids):
            raise RuntimeError("duplicate tokens in the UCI vocabulary")

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

    def encode_game(
        self,
        uci: str,
        white_elo: int,
        black_elo: int,
        result: str,
        max_len: int = 200,
    ) -> list[int]:
        """Encode one game; unknown moves map to ``<unk>``; the sequence is cut at ``max_len``.

        When the whole game fits in ``max_len`` the sequence ends with ``<eos>``; otherwise it
        is truncated to the first ``max_len`` ids (no result, no ``<eos>``).
        """
        if result not in RESULT_TOKENS:
            raise ValueError(f"unknown result {result!r}; expected one of {list(RESULT_TOKENS)}")
        ids = [
            self.bos_id,
            self.vocab[elo_token(white_elo, "w")],
            self.vocab[elo_token(black_elo, "b")],
        ]
        ids.extend(self.vocab.get(move, self.unk_id) for move in uci.split())
        ids.append(self.vocab[RESULT_TOKENS[result]])
        ids.append(self.eos_id)
        return ids[:max_len]

    def decode(self, ids: list[int]) -> list[str]:
        """Map ids back to their token strings (raises ``IndexError`` on unknown ids)."""
        return [self.ids[i] for i in ids]

    def to_dict(self) -> dict[str, object]:
        return {
            "version": 1,
            "size": len(self.ids),
            "specials": list(self.specials),
            "elo_bins": {"min": ELO_MIN, "max": ELO_MAX, "step": ELO_STEP},
            "tokens": list(self.ids),
        }

    def export(self, path: Path) -> None:
        """Write ``vocab.json`` deterministically (``indent=0``, ASCII, trailing newline)."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        text = json.dumps(self.to_dict(), indent=0, ensure_ascii=True) + "\n"
        path.write_text(text, encoding="utf-8", newline="\n")

    @classmethod
    def from_file(cls, path: Path) -> UciTokenizer:
        """Load a ``vocab.json`` and check it matches the built-in enumeration."""
        data = json.loads(Path(path).read_text(encoding="utf-8"))
        tokenizer = cls()
        if data.get("tokens") != tokenizer.ids:
            raise ValueError(f"{path} does not match the built-in UCI vocabulary")
        return tokenizer

"""Char-level SAN tokenizer: a fixed 32-symbol alphabet plus ``<pad>``, ``<bos>``, ``<eos>``.

The text is the numbered SAN movetext without comments followed by the result:
``1.e4 e5 2.Nf3 Nc6 ... 1-0``. Every character must belong to the alphabet.
"""

from __future__ import annotations

ALPHABET = " #+-.0123456789=BKNOQRabcdefghx/"
SPECIALS: list[str] = ["<pad>", "<bos>", "<eos>"]
PAD_ID = 0
BOS_ID = 1
EOS_ID = 2

class SanCharTokenizer:
    """Encode text as ``[<bos>, char..., <eos>]`` and decode ids back to text."""

    specials: list[str] = SPECIALS
    pad_id = PAD_ID
    bos_id = BOS_ID
    eos_id = EOS_ID

    def __init__(self) -> None:
        self.ids: list[str] = SPECIALS + list(ALPHABET)
        self.vocab: dict[str, int] = {token: i for i, token in enumerate(self.ids)}
        if len(self.vocab) != len(self.ids):
            raise RuntimeError("duplicate symbols in the SAN alphabet")

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

    def encode(self, text: str) -> list[int]:
        """Encode ``text``; raises ``ValueError`` on a character outside the alphabet."""
        ids = [self.bos_id]
        for ch in text:
            try:
                ids.append(self.vocab[ch])
            except KeyError:
                raise ValueError(f"character {ch!r} is not in the SAN alphabet") from None
        ids.append(self.eos_id)
        return ids

    def decode(self, ids: list[int]) -> str:
        """Join the characters of ``ids``, skipping special tokens."""
        return "".join(self.ids[i] for i in ids if i >= len(SPECIALS))

def san_text(san_movetext: str, result: str) -> str:
    """The exact text this scheme encodes: numbered SAN plus a space and the result."""
    return f"{san_movetext} {result}"


def uci_to_san(uci: str) -> str:
    """Numbered SAN movetext (``1.e4 e5 2.Nf3 ...``) for a space-separated UCI game.

    Raises ``ValueError`` on an illegal move, like the rest of the pipeline.
    """
    import chess

    board = chess.Board()
    out: list[str] = []
    for token in uci.split():
        move = chess.Move.from_uci(token)
        if move not in board.legal_moves:
            raise ValueError(f"illegal move {token} in {board.fen()}")
        san = board.san(move)
        out.append(f"{board.fullmove_number}.{san}" if board.turn == chess.WHITE else san)
        board.push(move)
    return " ".join(out)

"""BPE over UCI text with Hugging Face ``tokenizers``.

The text of one game is its UCI moves concatenated without spaces (``e2e4e7e5g1f3...``), so
that BPE can merge whole openings into one token. ``WhitespaceSplit`` keeps games apart when
several are fed as one text. The model has no normalizer, no continuing-subword prefix and no
end-of-word suffix: the TypeScript twin re-applies ``model.merges`` by rank over each
whitespace-split word and must reproduce ``Tokenizer.encode(text).ids`` exactly.
"""

from __future__ import annotations

from collections.abc import Iterable, Iterator
from pathlib import Path
from typing import TYPE_CHECKING

from tokenizers import Tokenizer, models, pre_tokenizers, trainers

from rukh.tokenize.uci_vocab import UciTokenizer

if TYPE_CHECKING:
    import chess

PROMOTION_LETTERS = frozenset("qrbn")


def bpe_text(uci: str) -> str:
    """The text BPE sees for one game: moves joined without spaces."""
    return uci.replace(" ", "")

def split_moves(text: str) -> list[str]:
    """Split a spaceless UCI string back into moves.

    A move is four characters plus a promotion letter, and ``b`` is both a file and a
    promotion piece, so the split replays the game: the letter belongs to the move only when
    the piece that moves is a pawn reaching the last rank (or, when the position is corrupt
    and no piece sits on the origin square, when the ranks are 7 to 8 or 2 to 1).
    """
    import chess

    board = chess.Board()
    moves: list[str] = []
    i = 0
    n = len(text)
    while i < n:
        move = text[i : i + 4]
        i += 4
        if i < n and text[i] in PROMOTION_LETTERS and _is_promotion(board, move):
            move += text[i]
            i += 1
        moves.append(move)
        try:
            parsed = chess.Move.from_uci(move)
        except ValueError:
            break
        if parsed in board.legal_moves:
            board.push(parsed)
        else:
            break
    if i < n:
        moves.extend(_split_blind(text[i:]))
    return moves

def _is_promotion(board: chess.Board, move: str) -> bool:
    import chess

    if len(move) != 4:
        return False
    try:
        from_sq = chess.parse_square(move[:2])
        to_sq = chess.parse_square(move[2:])
    except ValueError:
        return False
    piece = board.piece_at(from_sq)
    if piece is None:
        return move[1] + move[3] in ("78", "21") and abs(ord(move[0]) - ord(move[2])) <= 1
    return piece.piece_type == chess.PAWN and chess.square_rank(to_sq) in (0, 7)


def _split_blind(text: str) -> list[str]:
    """Rank-based split for text after an illegal move (no board to consult)."""
    moves: list[str] = []
    i = 0
    n = len(text)
    while i < n:
        move = text[i : i + 4]
        i += 4
        if (
            i < n
            and text[i] in PROMOTION_LETTERS
            and len(move) == 4
            and move[1] + move[3] in ("78", "21")
            and abs(ord(move[0]) - ord(move[2])) <= 1
        ):
            move += text[i]
            i += 1
        moves.append(move)
    return moves

def train_bpe(
    texts: Iterable[str],
    vocab_size: int = 4096,
    specials: list[str] | None = None,
) -> Tokenizer:
    """Train a BPE model on ``texts`` (one game per text, already spaceless)."""
    tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
    tokenizer.pre_tokenizer = pre_tokenizers.WhitespaceSplit()
    trainer = trainers.BpeTrainer(
        vocab_size=vocab_size,
        special_tokens=list(UciTokenizer.specials if specials is None else specials),
        show_progress=False,
    )
    tokenizer.train_from_iterator(texts, trainer)
    return tokenizer


def save(tokenizer: Tokenizer, path: Path) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    tokenizer.save(str(path), pretty=True)


def load_bpe(path: Path) -> Tokenizer:
    return Tokenizer.from_file(str(path))

def encode(tokenizer: Tokenizer, text: str) -> list[int]:
    return tokenizer.encode(text).ids


def decode(tokenizer: Tokenizer, ids: list[int]) -> str:
    """Concatenate the tokens (skipping specials) and re-split into space-separated moves."""
    specials = set(UciTokenizer.specials)
    tokens = [tokenizer.id_to_token(i) for i in ids]
    text = "".join(t for t in tokens if t is not None and t not in specials)
    return " ".join(split_moves(text))


def iter_parquet_texts(games: Path, n_games: int) -> Iterator[str]:
    """Yield ``bpe_text`` for the first ``n_games`` rows of a UCI games parquet."""
    import polars as pl

    frame = pl.scan_parquet(Path(games).as_posix()).select("uci").head(n_games).collect()
    for uci in frame["uci"]:
        yield bpe_text(uci)


def train_bpe_from_parquet(games: Path, out: Path, vocab_size: int, n_games: int) -> Tokenizer:
    tokenizer = train_bpe(iter_parquet_texts(games, n_games), vocab_size=vocab_size)
    save(tokenizer, out)
    return tokenizer


def longest_tokens(tokenizer: Tokenizer, n: int = 10) -> list[dict[str, object]]:
    """The ``n`` longest learned tokens, with how many whole moves each spans."""
    specials = set(UciTokenizer.specials)
    vocab = tokenizer.get_vocab()
    tokens = sorted((t for t in vocab if t not in specials), key=lambda t: (-len(t), vocab[t]))[:n]
    return [{"token": t, "id": vocab[t], "length": len(t), "moves": len(t) // 4} for t in tokens]