Muestrear jugadas

"""Playing a whole game: the model against an opponent, counting illegal proposals.

An opponent is anything with ``choose(board) -> chess.Move``: the seeded ``RandomOpponent``
used by the unit tests, ``StockfishOpponent`` for the Elo harness, or the browser in the demo.
"""

from __future__ import annotations

import random
from pathlib import Path
from types import TracebackType
from typing import Protocol

import chess
import chess.engine
from pydantic import BaseModel, ConfigDict

from rukh.infer.sampler import (
    HEADER_TOKENS,
    SampleConfig,
    legal_token_ids,
    model_generator,
    pick_move,
)
from rukh.models import MoveDecoder
from rukh.tokenize.uci_vocab import RESULT_TOKENS, UciTokenizer, elo_token

MOVE_TIME_SECONDS = 0.1
"""Default seconds per move for the Stockfish opponent (``docs/decisiones-de-ejecucion`` D-025)."""
REPETITION_CLOCK = 8
"""A threefold repetition needs at least this many reversible plies, so the check is gated on it."""
FIFTY_MOVE_PLIES = 100

class Opponent(Protocol):
    """Whatever plays the other side."""

    def choose(self, board: chess.Board) -> chess.Move: ...


class RandomOpponent:
    """A seeded random legal mover: the cheap baseline every test can afford."""

    def __init__(self, seed: int = 0) -> None:
        self.rng = random.Random(seed)

    def choose(self, board: chess.Board) -> chess.Move:
        return self.rng.choice(list(board.legal_moves))

class StockfishOpponent:
    """Stockfish limited to ``elo``; close it (or use it as a context manager) when done.

    ``UCI_Elo`` bottoms out at 1320, so a ``skill`` between 0 and 20 may be given instead: it
    selects ``Skill Level``, the only way to get rungs below that floor for the Elo harness.
    """

    def __init__(
        self,
        elo: int = 1400,
        move_time: float = MOVE_TIME_SECONDS,
        path: Path | None = None,
        skill: int | None = None,
    ) -> None:
        from rukh.engine import EngineNotFound, find_stockfish

        binary = path or find_stockfish()
        if binary is None:
            raise EngineNotFound("Stockfish not found; set RUKH_STOCKFISH or run get_stockfish.py")
        self.move_time = move_time
        self.engine = chess.engine.SimpleEngine.popen_uci(str(binary))
        if skill is None:
            self.engine.configure({"UCI_LimitStrength": True, "UCI_Elo": elo})
        else:
            self.engine.configure({"UCI_LimitStrength": False, "Skill Level": skill})

    def choose(self, board: chess.Board) -> chess.Move:
        played = self.engine.play(board, chess.engine.Limit(time=self.move_time))
        if played.move is None:
            raise chess.engine.EngineError("Stockfish returned no move")
        return played.move

    def close(self) -> None:
        self.engine.quit()

    def __enter__(self) -> StockfishOpponent:
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> None:
        self.close()

class GameResult(BaseModel):
    """One finished (or cut) game, from the point of view of the board."""

    model_config = ConfigDict(extra="forbid")

    result: str  # 1-0, 0-1, 1/2-1/2 or * when the game was cut short
    plies: int
    illegal_proposals: int
    moves: list[str]
    termination: str
    fen: str = chess.STARTING_FEN
    """The final position, so a game cut by the context limit can still be adjudicated."""

    @property
    def cut(self) -> bool:
        """True when the game ran out of context instead of ending on the board."""
        return self.result == "*"

def is_over(board: chess.Board) -> bool:
    """Whether the game is finished, including the claimable draws, without the O(n^2) check.

    ``board.is_game_over(claim_draw=True)`` tries every legal move looking for a claimable
    repetition, which is quadratic over a long game and is called once per ply. The claimable
    draws are covered here by the fifty-move clock and by ``is_repetition`` (a scan of the move
    stack) gated on a halfmove clock that makes a repetition possible at all.
    """
    if board.is_game_over():
        return True
    if board.halfmove_clock >= FIFTY_MOVE_PLIES:
        return True
    return board.halfmove_clock >= REPETITION_CLOCK and board.is_repetition(3)

def _history(tok: UciTokenizer, white_elo: int, black_elo: int) -> list[int]:
    return [
        tok.bos_id,
        tok.vocab[elo_token(white_elo, "w")],
        tok.vocab[elo_token(black_elo, "b")],
    ]

PIECE_VALUES = {
    chess.PAWN: 1.0,
    chess.KNIGHT: 3.0,
    chess.BISHOP: 3.0,
    chess.ROOK: 5.0,
    chess.QUEEN: 9.0,
}
ADJUDICATION_DEPTH = 8
ADJUDICATION_CP = 200
"""Centipawn margin above which an adjudicated position counts as a win."""
ADJUDICATION_PAWNS = ADJUDICATION_CP / 100.0
MATE_SCORE = 10_000

def material_balance(board: chess.Board) -> float:
    """Material of the position in pawns, from White's point of view."""
    total = 0.0
    for piece_type, value in PIECE_VALUES.items():
        total += value * len(board.pieces(piece_type, chess.WHITE))
        total -= value * len(board.pieces(piece_type, chess.BLACK))
    return total

def adjudicate(
    board: chess.Board | str,
    engine: chess.engine.SimpleEngine | None = None,
    depth: int = ADJUDICATION_DEPTH,
    margin_cp: int = ADJUDICATION_CP,
) -> tuple[str, str]:
    """Decide a game that was cut short; returns ``(result, how)``.

    With an engine the final position is analysed to a shallow depth and a score of at least
    ``margin_cp`` for one side is a win for that side; without one (or when the analysis gives no
    centipawn score) the same margin is applied to the material count, which is crude but is
    still an answer, and is far better than booking every cut game as half a point.
    """
    position = chess.Board(board) if isinstance(board, str) else board
    if engine is not None:
        try:
            info = engine.analyse(position, chess.engine.Limit(depth=depth))
            score = info["score"].white().score(mate_score=MATE_SCORE)
        except (chess.engine.EngineError, chess.engine.EngineTerminatedError, KeyError):
            score = None
        if score is not None:
            return _verdict(float(score), float(margin_cp)), f"engine depth {depth}"
    pawns = material_balance(position)
    return _verdict(pawns * 100.0, float(margin_cp)), "material count"


def _verdict(centipawns: float, margin: float) -> str:
    """``1-0``/``0-1``/``1/2-1/2`` from a White-relative score in centipawns."""
    if centipawns >= margin:
        return "1-0"
    if centipawns <= -margin:
        return "0-1"
    return "1/2-1/2"

def result_token(tok: UciTokenizer, result: str) -> int | None:
    """Vocabulary id of a finished game's result token, or None for an unfinished game."""
    name = RESULT_TOKENS.get(result)
    return tok.vocab[name] if name else None

def legality_rate(
    model: MoveDecoder,
    tok: UciTokenizer,
    boards: list[chess.Board],
    histories: list[list[int]],
    cfg: SampleConfig,
) -> float:
    """Share of the positions actually queried where the unmasked network proposes a legal move.

    Positions with no legal move (checkmate, stalemate) are never asked, so they are not counted
    in the denominator either: the rate is over the boards the model was really given.
    """
    if not boards:
        return 0.0
    unmasked = cfg.model_copy(update={"mask_illegal": False})
    generator = model_generator(model, cfg)
    legal = 0
    counted = 0
    for board, history in zip(boards, histories, strict=True):
        if not legal_token_ids(board, tok):
            continue
        counted += 1
        _, report = pick_move(model, tok, board, history, unmasked, generator)
        legal += int(report["legal"])
    return legal / counted if counted else 0.0

"""Inference: masked sampling of one move and self-play against an opponent."""

from rukh.infer.game import (
    ADJUDICATION_CP,
    ADJUDICATION_DEPTH,
    GameResult,
    Opponent,
    RandomOpponent,
    StockfishOpponent,
    adjudicate,
    is_over,
    legality_rate,
    material_balance,
    play_game,
    result_token,
)
from rukh.infer.sampler import (
    HEADER_TOKENS,
    SampleConfig,
    legal_token_ids,
    model_generator,
    pick_move,
    prompt_ids,
)

__all__ = [
    "ADJUDICATION_CP",
    "ADJUDICATION_DEPTH",
    "HEADER_TOKENS",
    "GameResult",
    "Opponent",
    "RandomOpponent",
    "SampleConfig",
    "StockfishOpponent",
    "adjudicate",
    "is_over",
    "legal_token_ids",
    "legality_rate",
    "material_balance",
    "model_generator",
    "pick_move",
    "play_game",
    "prompt_ids",
    "result_token",
]

@app.command("play")
def play_cmd(
    ckpt: Annotated[
        Path,
        typer.Option(
            "--ckpt", exists=True, dir_okay=False, readable=True, help="Checkpoint to play with."
        ),
    ],
    games: Annotated[int, typer.Option("--games", help="Number of games to play.")] = 1,
    opponent: Annotated[
        str, typer.Option("--opponent", help="Opponent: random or stockfish.")
    ] = "random",
    elo: Annotated[int, typer.Option("--elo", help="UCI_Elo of the Stockfish opponent.")] = 1400,
    temperature: Annotated[
        float, typer.Option("--temperature", help="Sampling temperature (0 = argmax).")
    ] = 0.6,
    top_k: Annotated[int, typer.Option("--top-k", help="Top-k truncation; 0 disables it.")] = 20,
    no_mask: Annotated[
        bool, typer.Option("--no-mask", help="Sample without the legality mask (measures it).")
    ] = False,
    seed: Annotated[int, typer.Option("--seed", help="Seed for sampling and the opponent.")] = 0,
    as_json: Annotated[bool, typer.Option("--json", help="Print the games as JSON only.")] = False,
) -> None:
    """Play games between a trained decoder and an opponent, with or without the legality mask."""
    import json

    from rukh.infer import RandomOpponent, SampleConfig, StockfishOpponent, play_game
    from rukh.tokenize.uci_vocab import UciTokenizer
    from rukh.train import load_model

    if opponent not in ("random", "stockfish"):
        typer.echo("error: --opponent must be random or stockfish", err=True)
        raise typer.Exit(code=2)
    model, _ = load_model(ckpt)
    tok = UciTokenizer()
    cfg = SampleConfig(
        temperature=temperature,
        top_k=top_k or None,
        mask_illegal=not no_mask,
        seed=seed,
    )
    rival = RandomOpponent(seed) if opponent == "random" else StockfishOpponent(elo=elo)
    try:
        results = [
            play_game(model, tok, rival, cfg, model_color=index % 2 == 0) for index in range(games)
        ]
    finally:
        if isinstance(rival, StockfishOpponent):
            rival.close()
    if as_json:
        typer.echo(json.dumps([result.model_dump() for result in results], indent=2))
        return
    typer.echo(f"opponent: {opponent}")
    typer.echo(f"mask:     {'on' if cfg.mask_illegal else 'off'}")
    for index, result in enumerate(results):
        typer.echo(
            f"  game {index + 1}: {result.result:<7} {result.plies:>3} plies  "
            f"{result.illegal_proposals} illegal  ({result.termination})"
        )
    typer.echo(f"illegal:  {sum(r.illegal_proposals for r in results)} proposals")

"""Tests for rukh.infer: legality masking, temperature, promotions and whole games."""

from __future__ import annotations

import chess
import pytest
import torch

from rukh.infer import (
    GameResult,
    RandomOpponent,
    SampleConfig,
    StockfishOpponent,
    adjudicate,
    is_over,
    legal_token_ids,
    legality_rate,
    material_balance,
    model_generator,
    pick_move,
    play_game,
    prompt_ids,
    result_token,
)
from rukh.models import DecoderConfig, MoveDecoder
from rukh.tokenize.uci_vocab import UciTokenizer

# The Stockfish game carries only `engine`, so `-m unit` never needs the binary.
unit = pytest.mark.unit

TOY = DecoderConfig(vocab_size=2030, n_layer=2, n_head=2, d_model=32, block=64)
PROMOTION_FEN = "8/P6k/8/8/8/8/6K1/8 w - - 0 1"


@pytest.fixture(scope="module")
def tok() -> UciTokenizer:
    return UciTokenizer()


@pytest.fixture(scope="module")
def model() -> MoveDecoder:
    torch.manual_seed(0)
    return MoveDecoder(TOY).eval()


def history(tok: UciTokenizer) -> list[int]:
    return [tok.bos_id, tok.vocab["<w1800>"], tok.vocab["<b1800>"]]

@unit
def test_the_initial_position_has_twenty_legal_tokens(tok: UciTokenizer) -> None:
    board = chess.Board()
    ids = legal_token_ids(board, tok)
    assert len(ids) == 20
    assert ids == sorted(set(ids))
    for token_id in ids:
        assert board.is_legal(chess.Move.from_uci(tok.ids[token_id]))
    assert {tok.ids[i] for i in ids} == {move.uci() for move in board.legal_moves}


@unit
def test_castling_is_the_two_square_king_move(tok: UciTokenizer) -> None:
    board = chess.Board("r3k2r/8/8/8/8/8/8/R3K2R w KQkq - 0 1")
    tokens = {tok.ids[i] for i in legal_token_ids(board, tok)}
    assert {"e1g1", "e1c1"} <= tokens

@unit
def test_masked_sampling_never_proposes_an_illegal_move(
    model: MoveDecoder, tok: UciTokenizer
) -> None:
    board = chess.Board()
    cfg = SampleConfig(temperature=1.5, top_k=None, mask_illegal=True, seed=0)
    generator = cfg.generator()
    seen = set()
    for _ in range(1000):
        move, report = pick_move(model, tok, board, history(tok), cfg, generator)
        assert move is not None
        assert report["legal"] is True and report["masked"] is True
        assert board.is_legal(move)
        seen.add(move.uci())
    assert len(seen) > 1  # the mask does not collapse the distribution


@unit
def test_without_the_mask_an_illegal_token_is_reported_as_such(
    model: MoveDecoder, tok: UciTokenizer
) -> None:
    board = chess.Board()
    cfg = SampleConfig(temperature=1.0, top_k=None, mask_illegal=False, seed=1)
    generator = cfg.generator()
    illegal = 0
    for _ in range(50):
        move, report = pick_move(model, tok, board, history(tok), cfg, generator)
        assert report["masked"] is False
        assert report["raw_token"] in tok.vocab
        if move is None:
            illegal += 1
            assert report["legal"] is False
    assert illegal > 0  # an untrained model is almost never legal on its own

@unit
def test_zero_temperature_is_argmax(model: MoveDecoder, tok: UciTokenizer) -> None:
    board = chess.Board()
    cfg = SampleConfig(temperature=0.0, top_k=None, mask_illegal=True)
    ids = legal_token_ids(board, tok)
    idx = torch.tensor([history(tok)], dtype=torch.long)
    logits = model.next_logits(idx)[0]
    best = tok.ids[ids[int(torch.argmax(logits[torch.tensor(ids)]))]]
    for _ in range(5):
        move, report = pick_move(model, tok, board, history(tok), cfg)
        assert move is not None and move.uci() == best
        assert report["raw_token"] == best


@unit
def test_top_k_restricts_the_reported_distribution(model: MoveDecoder, tok: UciTokenizer) -> None:
    board = chess.Board()
    cfg = SampleConfig(temperature=1.0, top_k=3, mask_illegal=True, seed=2)
    drawn = set()
    generator = cfg.generator()
    for _ in range(200):
        move, report = pick_move(model, tok, board, history(tok), cfg, generator)
        assert move is not None
        drawn.add(move.uci())
        assert len(report["top5"]) <= 3
        assert sum(probability for _, probability in report["top5"]) == pytest.approx(1.0, abs=1e-5)
    assert len(drawn) <= 3

@unit
def test_a_forced_promotion_yields_a_promotion_token(model: MoveDecoder, tok: UciTokenizer) -> None:
    board = chess.Board(PROMOTION_FEN)
    tokens = {tok.ids[i] for i in legal_token_ids(board, tok)}
    assert {"a7a8q", "a7a8r", "a7a8b", "a7a8n"} <= tokens
    cfg = SampleConfig(temperature=1.0, top_k=None, mask_illegal=True, seed=3)
    generator = cfg.generator()
    for _ in range(20):
        move, report = pick_move(model, tok, board, history(tok), cfg, generator)
        assert move is not None
        if move.from_square == chess.A7:
            assert move.promotion is not None
            assert len(report["raw_token"]) == 5


@unit
def test_a_finished_position_has_no_legal_tokens(model: MoveDecoder, tok: UciTokenizer) -> None:
    board = chess.Board("7k/5Q2/6K1/8/8/8/8/8 b - - 0 1")  # checkmate
    assert legal_token_ids(board, tok) == []
    move, report = pick_move(model, tok, board, history(tok), SampleConfig(seed=0))
    assert move is None and report["legal"] is False and report["top5"] == []

@unit
def test_a_full_game_against_a_random_opponent(model: MoveDecoder, tok: UciTokenizer) -> None:
    cfg = SampleConfig(temperature=1.0, top_k=20, mask_illegal=True, seed=7)
    result = play_game(model, tok, RandomOpponent(seed=7), cfg, max_plies=40)
    assert isinstance(result, GameResult)
    assert result.result in ("1-0", "0-1", "1/2-1/2", "*")
    assert result.plies == len(result.moves)
    assert result.illegal_proposals == 0  # the mask makes them impossible
    assert result.termination

    board = chess.Board()
    for uci in result.moves:
        move = chess.Move.from_uci(uci)
        assert board.is_legal(move)
        board.push(move)


@unit
def test_an_unmasked_game_counts_illegal_proposals_and_still_ends(
    model: MoveDecoder, tok: UciTokenizer
) -> None:
    cfg = SampleConfig(temperature=1.0, top_k=None, mask_illegal=False, seed=5)
    result = play_game(
        model, tok, RandomOpponent(seed=5), cfg, model_color=chess.BLACK, max_plies=20
    )
    assert result.illegal_proposals > 0
    assert result.plies == 20 or result.result != "*"


@unit
def test_the_game_stops_at_the_block_size_by_default(model: MoveDecoder, tok: UciTokenizer) -> None:
    cfg = SampleConfig(temperature=1.0, top_k=20, seed=11)
    result = play_game(model, tok, RandomOpponent(seed=11), cfg)
    assert result.plies <= TOY.block - 4


@unit
def test_the_same_seed_replays_the_same_game(model: MoveDecoder, tok: UciTokenizer) -> None:
    cfg = SampleConfig(temperature=1.0, top_k=20, seed=13)
    first = play_game(model, tok, RandomOpponent(seed=13), cfg, max_plies=30)
    second = play_game(model, tok, RandomOpponent(seed=13), cfg, max_plies=30)
    assert first.moves == second.moves

@unit
def test_legality_rate_is_a_share(model: MoveDecoder, tok: UciTokenizer) -> None:
    boards = [chess.Board() for _ in range(5)]
    histories = [history(tok) for _ in boards]
    rate = legality_rate(model, tok, boards, histories, SampleConfig(seed=0))
    assert 0.0 <= rate <= 1.0
    assert legality_rate(model, tok, [], [], SampleConfig()) == 0.0


@unit
def test_result_tokens_map_to_the_vocabulary(tok: UciTokenizer) -> None:
    assert result_token(tok, "1-0") == tok.vocab["<1-0>"]
    assert result_token(tok, "1/2-1/2") == tok.vocab["<1/2>"]
    assert result_token(tok, "*") is None


@unit
def test_the_generator_follows_the_model_device(model: MoveDecoder) -> None:
    cfg = SampleConfig(seed=0)
    generator = model_generator(model, cfg)
    assert generator is not None
    assert generator.device.type == next(model.parameters()).device.type
    assert model_generator(model, SampleConfig(seed=None)) is None

@pytest.mark.gpu
def test_sampling_runs_on_cuda(tok: UciTokenizer) -> None:
    if not torch.cuda.is_available():
        pytest.skip("CUDA is not available")
    torch.manual_seed(0)
    on_gpu = MoveDecoder(TOY).eval().to("cuda")
    cfg = SampleConfig(temperature=1.0, top_k=20, mask_illegal=True, seed=0)
    generator = model_generator(on_gpu, cfg)
    assert generator is not None and generator.device.type == "cuda"
    board = chess.Board()
    for _ in range(20):
        move, report = pick_move(on_gpu, tok, board, history(tok), cfg, generator)
        assert move is not None and board.is_legal(move)
        assert report["legal"] is True
    result = play_game(on_gpu, tok, RandomOpponent(seed=0), cfg, max_plies=20)
    assert result.plies > 0

@unit
def test_the_prompt_keeps_the_header_when_it_has_to_be_cropped(tok: UciTokenizer) -> None:
    ids = history(tok) + list(range(100, 200))
    cropped = prompt_ids(ids, 10)
    assert len(cropped) == 10
    assert cropped[:3] == history(tok)
    assert cropped[3:] == list(range(193, 200))
    assert prompt_ids(ids, 500) == ids


@unit
def test_a_long_game_still_carries_its_elo_tokens(model: MoveDecoder, tok: UciTokenizer) -> None:
    board = chess.Board()
    long_history = history(tok) + [tok.vocab["e2e4"]] * (TOY.block * 2)
    cfg = SampleConfig(temperature=0.0, top_k=None, mask_illegal=True)
    move, _ = pick_move(model, tok, board, long_history, cfg)
    assert move is not None  # the decoder would have refused a prompt longer than the block

@unit
def test_is_over_sees_the_claimable_draws() -> None:
    assert is_over(chess.Board("7k/5Q2/6K1/8/8/8/8/8 b - - 0 1"))  # checkmate
    assert not is_over(chess.Board())
    fifty = chess.Board("7k/8/6K1/8/8/8/8/R7 w - - 100 80")
    assert is_over(fifty)


@unit
def test_adjudication_reads_material_when_there_is_no_engine() -> None:
    assert material_balance(chess.Board()) == 0.0
    a_queen_up = chess.Board("7k/8/8/8/8/8/8/Q5K1 w - - 0 1")
    assert material_balance(a_queen_up) == 9.0
    assert adjudicate(a_queen_up) == ("1-0", "material count")
    assert adjudicate("7k/8/8/8/8/8/8/q5K1 w - - 0 1") == ("0-1", "material count")
    assert adjudicate(chess.Board())[0] == "1/2-1/2"
    one_pawn_up = chess.Board("4k3/p7/8/8/8/8/PP6/4K3 w - - 0 1")
    assert adjudicate(one_pawn_up) == ("1/2-1/2", "material count")  # under the two-pawn margin
    two_pawns_up = chess.Board("4k3/8/8/8/8/8/PP6/4K3 w - - 0 1")
    assert adjudicate(two_pawns_up) == ("1-0", "material count")  # the margin is inclusive

@unit
def test_a_cut_game_reports_its_final_position(model: MoveDecoder, tok: UciTokenizer) -> None:
    cfg = SampleConfig(temperature=1.0, top_k=20, seed=3)
    result = play_game(model, tok, RandomOpponent(seed=3), cfg, max_plies=6)
    assert result.result == "*" and result.cut is True
    board = chess.Board(result.fen)
    assert board.fullmove_number == 4
    assert adjudicate(result.fen)[0] in ("1-0", "0-1", "1/2-1/2")


@unit
def test_sample_config_rejects_impossible_values() -> None:
    with pytest.raises(ValueError, match="non-negative"):
        SampleConfig(temperature=-0.1)
    with pytest.raises(ValueError, match="top_k"):
        SampleConfig(top_k=0)
    with pytest.raises(ValueError):
        SampleConfig(temp=0.5)  # type: ignore[call-arg]
    assert SampleConfig(seed=None).generator() is None
    assert SampleConfig(seed=1).generator().device.type == "cpu"


@unit
def test_cli_play_is_registered() -> None:
    from typer.testing import CliRunner

    from rukh.cli import app

    result = CliRunner().invoke(app, ["play", "--help"])
    assert result.exit_code == 0, result.output
    for option in ("--ckpt", "--games", "--opponent", "--no-mask"):
        assert option in result.output

@pytest.mark.engine
def test_a_game_against_stockfish(model: MoveDecoder, tok: UciTokenizer) -> None:
    from rukh.engine import find_stockfish

    if find_stockfish() is None:
        pytest.skip("Stockfish is not available")
    cfg = SampleConfig(temperature=1.0, top_k=20, seed=17)
    with StockfishOpponent(elo=1400, move_time=0.02) as rival:
        result = play_game(model, tok, rival, cfg, model_color=chess.WHITE, max_plies=20)
    assert result.plies > 0 and result.illegal_proposals == 0