Legalidad, exactitud y puzles

"""Tests for rukh.eval.legality: sampling validation prefixes and the unmasked legality rate."""

from __future__ import annotations

import importlib
from pathlib import Path
from typing import Any

import chess
import pytest
import torch

from rukh.eval.legality import (
    Position,
    board_of,
    header,
    legality,
    position_at,
    sample_positions,
)
from rukh.infer import SampleConfig
from rukh.models import DecoderConfig, MoveDecoder
from rukh.tokenize.uci_vocab import UciTokenizer

pytestmark = pytest.mark.unit

# ``rukh.eval`` re-exports the function under the module's own name, so the module object for
# monkeypatching has to be asked for explicitly.
legality_module = importlib.import_module("rukh.eval.legality")

TOY = DecoderConfig(vocab_size=2030, n_layer=2, n_head=2, d_model=32, block=64)
GAME = "e2e4 e7e5 g1f3 b8c6 f1b5 a7a6 b5a4 g8f6"

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


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


@pytest.fixture
def games(tmp_path: Path) -> Path:
    import polars as pl

    path = tmp_path / "games.parquet"
    pl.DataFrame(
        {
            "game_id": [1, 2, 3],
            "uci": [GAME, GAME, "e2e4"],
            "white_elo": [1850, 2150, 1900],
            "black_elo": [1950, 2050, 1900],
            "n_plies": [8, 8, 1],
        }
    ).write_parquet(path)
    return path

def test_position_at_keeps_the_prefix_the_history_and_the_human_move(tok: UciTokenizer) -> None:
    moves = GAME.split()
    position = position_at(tok, 7, moves, 4, 1850, 1950)
    assert position.moves == moves[:4]
    assert position.target == moves[4]
    assert position.history[:3] == header(tok, 1850, 1950)
    assert len(position.history) == 3 + 4
    assert position.elo == 1850  # ply 4 is White's turn


def test_the_elo_of_a_position_is_the_side_to_move(tok: UciTokenizer) -> None:
    moves = GAME.split()
    assert position_at(tok, 1, moves, 3, 1850, 1950).elo == 1950


def test_the_crop_keeps_the_header_and_drops_the_oldest_moves(tok: UciTokenizer) -> None:
    moves = GAME.split()
    position = position_at(tok, 1, moves, 8, 1850, 1950, block=5)
    assert len(position.history) == 5
    # The Elo conditioning survives: a plain ``history[-block:]`` would have thrown it away.
    assert position.history[:3] == header(tok, 1850, 1950)
    assert position.history[3:] == [tok.vocab[uci] for uci in moves[6:8]]


def test_a_short_history_is_left_alone(tok: UciTokenizer) -> None:
    position = position_at(tok, 1, GAME.split(), 4, 1850, 1950, block=200)
    assert len(position.history) == 7
    assert position.history[:3] == header(tok, 1850, 1950)

def test_legality_is_measured_twice_with_different_definitions(
    tok: UciTokenizer, model: MoveDecoder
) -> None:
    positions = [position_at(tok, 1, GAME.split(), ply, 1850, 1950) for ply in range(1, 6)]
    cfg = SampleConfig(temperature=0.6, top_k=20, seed=0)
    argmax = legality(model, tok, positions, cfg, mode="argmax")
    sampled = legality(model, tok, positions, cfg, mode="sampled")
    assert argmax.mode == "argmax" and argmax.temperature is None and argmax.top_k is None
    assert sampled.mode == "sampled" and sampled.temperature == 0.6 and sampled.top_k == 20
    # argmax is deterministic: the same call twice gives the same rate.
    assert legality(model, tok, positions, cfg, mode="argmax").legal == argmax.legal


def test_an_unknown_legality_mode_is_an_error(tok: UciTokenizer, model: MoveDecoder) -> None:
    with pytest.raises(ValueError, match="unknown legality mode"):
        legality(model, tok, [], mode="greedy")

def test_board_of_replays_the_prefix(tok: UciTokenizer) -> None:
    position = position_at(tok, 1, GAME.split(), 4, 1850, 1950)
    board = board_of(position)
    assert board.turn == chess.WHITE
    assert board.fullmove_number == 3


def test_sample_positions_returns_one_prefix_per_game(tok: UciTokenizer, games: Path) -> None:
    positions = sample_positions(games, 10, tok, seed=0)
    assert len(positions) == 2  # the one-move game is too short to sample from
    for position in positions:
        assert 1 <= position.ply < 8
        assert position.target == GAME.split()[position.ply]
        board_of(position)


def test_sample_positions_is_deterministic(tok: UciTokenizer, games: Path) -> None:
    first = sample_positions(games, 10, tok, seed=3)
    second = sample_positions(games, 10, tok, seed=3)
    assert [p.model_dump() for p in first] == [p.model_dump() for p in second]


def test_sample_positions_stops_at_n(tok: UciTokenizer, games: Path) -> None:
    assert len(sample_positions(games, 1, tok, seed=0)) == 1

def test_legality_counts_the_legal_proposals(
    tok: UciTokenizer, model: MoveDecoder, monkeypatch: pytest.MonkeyPatch
) -> None:
    answers = iter([True, False, True, True])

    def fake_pick(*_args: Any, **_kwargs: Any) -> tuple[None, dict[str, Any]]:
        return None, {"legal": next(answers), "raw_token": "e2e4", "top5": [], "masked": False}

    monkeypatch.setattr(legality_module, "pick_move", fake_pick)
    positions = [position_at(tok, 1, GAME.split(), ply, 1850, 1950) for ply in range(1, 5)]
    result = legality(model, tok, positions)
    assert (result.positions, result.legal) == (4, 3)
    assert result.rate == pytest.approx(0.75)


def test_legality_asks_the_model_without_the_mask(
    tok: UciTokenizer, model: MoveDecoder, monkeypatch: pytest.MonkeyPatch
) -> None:
    seen: list[bool] = []

    def fake_pick(_model, _tok, _board, _history, cfg, _generator=None):  # type: ignore[no-untyped-def]
        seen.append(cfg.mask_illegal)
        return None, {"legal": True, "raw_token": "e2e4", "top5": [], "masked": False}

    monkeypatch.setattr(legality_module, "pick_move", fake_pick)
    legality(model, tok, [position_at(tok, 1, GAME.split(), 2, 1850, 1950)])
    assert seen == [False]

def test_legality_of_a_real_toy_model_is_a_rate(tok: UciTokenizer, model: MoveDecoder) -> None:
    positions = [position_at(tok, 1, GAME.split(), ply, 1850, 1950) for ply in range(1, 6)]
    result = legality(model, tok, positions)
    assert result.positions == 5
    assert 0.0 <= result.rate <= 1.0
    assert result.legal == int(result.rate * result.positions)


def test_legality_without_positions_is_zero(tok: UciTokenizer, model: MoveDecoder) -> None:
    empty: list[Position] = []
    assert legality(model, tok, empty).rate == 0.0

"""Tests for rukh.eval.accuracy: Elo bands and top-1/top-3 agreement with the human move."""

from __future__ import annotations

from collections.abc import Iterator

import pytest
import torch

from rukh.eval.accuracy import accuracy, elo_band
from rukh.eval.legality import position_at
from rukh.models import DecoderConfig, MoveDecoder
from rukh.tokenize.uci_vocab import UciTokenizer

pytestmark = pytest.mark.unit

TOY = DecoderConfig(vocab_size=2030, n_layer=2, n_head=2, d_model=32, block=64)
GAME = "e2e4 e7e5 g1f3 b8c6 f1b5 a7a6 b5a4 g8f6"


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


class RankedModel:
    """A stand-in for the decoder whose logits rank a fixed list of tokens first."""

    def __init__(self, tok: UciTokenizer, ranking: list[str]) -> None:
        self.ids = [tok.vocab[token] for token in ranking]
        self.vocab_size = len(tok)
        self.cfg = TOY  # the prompt is cropped to the model's block, so the stub needs one

    def parameters(self) -> Iterator[torch.Tensor]:
        return iter([torch.zeros(1)])

    def next_logits(self, idx: torch.Tensor) -> torch.Tensor:
        logits = torch.full((idx.shape[0], self.vocab_size), -10.0)
        for rank, token_id in enumerate(self.ids):
            logits[:, token_id] = 10.0 - rank
        return logits

@pytest.mark.parametrize(
    ("elo", "expected"),
    [
        (1200, "<1800"),
        (1799, "<1800"),
        (1800, "1800-2000"),
        (1999, "1800-2000"),
        (2000, "2000-2200"),
        (2599, "2400-2600"),
        (2600, "2600+"),
        (3000, "2600+"),
    ],
)
def test_elo_band_cuts_every_two_hundred_points(elo: int, expected: str) -> None:
    assert elo_band(elo) == expected


def test_top1_and_top3_see_the_ranking(tok: UciTokenizer) -> None:
    moves = GAME.split()
    positions = [position_at(tok, 1, moves, 1, 1850, 1950)]  # the human played e7e5
    exact = accuracy(RankedModel(tok, ["e7e5", "a7a6", "b8c6"]), tok, positions)  # type: ignore[arg-type]
    assert (exact.top1, exact.top3) == (1.0, 1.0)
    third = accuracy(RankedModel(tok, ["a7a6", "b8c6", "e7e5"]), tok, positions)  # type: ignore[arg-type]
    assert (third.top1, third.top3) == (0.0, 1.0)
    missed = accuracy(RankedModel(tok, ["a7a6", "b8c6", "g8f6"]), tok, positions)  # type: ignore[arg-type]
    assert (missed.top1, missed.top3) == (0.0, 0.0)

def test_positions_are_filed_under_the_band_of_the_side_to_move(tok: UciTokenizer) -> None:
    moves = GAME.split()
    positions = [
        position_at(tok, 1, moves, 1, 1850, 2150),  # Black to move: 2150 -> 2000-2200
        position_at(tok, 1, moves, 2, 1850, 2150),  # White to move: 1850 -> 1800-2000
    ]
    result = accuracy(RankedModel(tok, ["e7e5"]), tok, positions)  # type: ignore[arg-type]
    bands = {band.band: band for band in result.bands}
    assert set(bands) == {"1800-2000", "2000-2200"}
    assert bands["2000-2200"].top1 == 1.0
    assert bands["1800-2000"].top1 == 0.0
    assert result.top1 == pytest.approx(0.5)


def test_positions_without_a_human_move_are_skipped(tok: UciTokenizer) -> None:
    moves = GAME.split()
    end = position_at(tok, 1, moves, len(moves), 1850, 1950)
    assert end.target is None
    assert accuracy(RankedModel(tok, ["e7e5"]), tok, [end]).positions == 0  # type: ignore[arg-type]


def test_accuracy_of_nothing_is_zero(tok: UciTokenizer) -> None:
    result = accuracy(RankedModel(tok, ["e7e5"]), tok, [])  # type: ignore[arg-type]
    assert (result.positions, result.top1, result.top3) == (0, 0.0, 0.0)


def test_accuracy_runs_on_the_real_decoder(tok: UciTokenizer) -> None:
    torch.manual_seed(0)
    model = MoveDecoder(TOY).eval()
    moves = GAME.split()
    positions = [position_at(tok, 1, moves, ply, 1850, 1950) for ply in range(1, 6)]
    result = accuracy(model, tok, positions)
    assert result.positions == 5
    assert result.top1 <= result.top3

def run_puzzles(
    source: MoveSource,
    tok: UciTokenizer,
    items: Sequence[PuzzleItem],
    cache: EvalCache | None = None,
    header_elo: int = HEADER_ELO,
) -> PuzzleResult:
    """Attempt every puzzle, reusing the cached attempts of a previous run of the same weights."""
    attempts: list[PuzzleAttempt] = []
    for item in items:
        cached = cache.get(SUITE, item.puzzle_id) if cache is not None else None
        if cached is not None:
            attempts.append(PuzzleAttempt.model_validate(cached))
            continue
        attempt = solve_puzzle(source, tok, item, header_elo=header_elo)
        if cache is not None:
            cache.put(SUITE, item.puzzle_id, attempt.model_dump())
        attempts.append(attempt)
    return summarize(attempts)

def summarize(attempts: Sequence[PuzzleAttempt]) -> PuzzleResult:
    """Fold attempts into overall and per-band rates."""
    counts: dict[str, int] = {}
    solved: dict[str, int] = {}
    for attempt in attempts:
        counts[attempt.band] = counts.get(attempt.band, 0) + 1
        solved[attempt.band] = solved.get(attempt.band, 0) + int(attempt.solved)
    total = sum(counts.values())
    return PuzzleResult(
        attempted=total,
        solved=sum(solved.values()),
        rate=sum(solved.values()) / total if total else 0.0,
        bands=[
            BandPuzzles(
                band=band,
                attempted=counts[band],
                solved=solved.get(band, 0),
                rate=solved.get(band, 0) / counts[band],
            )
            for band in sorted(counts)
        ],
    )

def load_puzzles(path: Path, per_band: int, seed: int = 0, split: str = "test") -> list[PuzzleItem]:
    """Read up to ``per_band`` puzzles of the given split from the P1 puzzle parquet."""
    import polars as pl

    frame = (
        pl.scan_parquet(Path(path).as_posix())
        .select(PUZZLE_COLUMNS)
        .filter(pl.col("split") == split)
        .collect()
    )
    items: list[PuzzleItem] = []
    for band in sorted(frame["band"].unique().to_list()):
        rows = (
            frame.filter(pl.col("band") == band)
            .sample(n=min(per_band, frame.filter(pl.col("band") == band).height), seed=seed)
            .rows(named=True)
        )
        items.extend(
            PuzzleItem(
                puzzle_id=str(row["puzzle_id"]),
                fen=str(row["fen"]),
                moves=str(row["moves"]).split(),
                rating=int(row["rating"]),
                band=str(row["band"]),
            )
            for row in rows
        )
    return items

"""SQLite cache of finished evaluation items, keyed by ``(model_sha, suite, item_id)``.

Games against Stockfish and puzzle lines are the expensive part of the harness: a full suite is
hundreds of games and thousands of puzzles. Every finished item is stored as JSON under the SHA
of the weights that produced it, so re-running a suite after a crash (or after adding one rung)
only pays for what is missing. ``--no-cache`` builds a disabled cache: it reads nothing and
writes nothing, which is what a benchmark of the harness itself wants.

The key is the weights **and** the settings that change what an item means: temperature, top-k,
the engine's move time, the ply limit, the rung definitions, the number of games and the seed.
Without them, lowering the temperature or raising the move time would silently reuse games
played under the old settings and report them as the new ones.
"""

from __future__ import annotations

import hashlib
import json
import sqlite3
from collections.abc import Mapping
from pathlib import Path
from types import TracebackType
from typing import Any

SCHEMA = """
CREATE TABLE IF NOT EXISTS items (
    model_sha TEXT NOT NULL,
    suite     TEXT NOT NULL,
    item_id   TEXT NOT NULL,
    payload   TEXT NOT NULL,
    PRIMARY KEY (model_sha, suite, item_id)
)
"""
CHUNK = 1 << 20
SHA_PREFIX = 16
"""Characters of each SHA kept in the key: enough to identify, short enough to read."""

def config_sha(fields: Mapping[str, Any]) -> str:
    """SHA-256 of the evaluation-relevant settings, as canonical JSON."""
    payload = json.dumps(dict(fields), sort_keys=True, default=str, ensure_ascii=False)
    return hashlib.sha256(payload.encode("utf-8")).hexdigest()


def file_sha(path: Path) -> str:
    """SHA-256 of a file, read in chunks (a checkpoint does not fit comfortably in memory)."""
    digest = hashlib.sha256()
    with Path(path).open("rb") as handle:
        while chunk := handle.read(CHUNK):
            digest.update(chunk)
    return digest.hexdigest()

class EvalCache:
    """Key-value store of evaluation items; ``enabled=False`` turns every call into a no-op."""

    def __init__(
        self,
        path: Path | None,
        model_sha: str,
        enabled: bool = True,
        config_sha: str | None = None,
    ) -> None:
        self.model_sha = model_sha
        self.config_sha = config_sha
        self.key = (
            model_sha if not config_sha else f"{model_sha[:SHA_PREFIX]}:{config_sha[:SHA_PREFIX]}"
        )
        self.enabled = enabled and path is not None
        self.path = Path(path) if path is not None else None
        self._conn: sqlite3.Connection | None = None
        if self.enabled and self.path is not None:
            self.path.parent.mkdir(parents=True, exist_ok=True)
            self._conn = sqlite3.connect(self.path)
            self._conn.execute(SCHEMA)
            self._conn.commit()

    def get(self, suite: str, item_id: str) -> dict[str, Any] | None:
        """The stored payload of one item, or None when it has not been computed yet."""
        if self._conn is None:
            return None
        row = self._conn.execute(
            "SELECT payload FROM items WHERE model_sha = ? AND suite = ? AND item_id = ?",
            (self.key, suite, item_id),
        ).fetchone()
        if row is None:
            return None
        payload = json.loads(row[0])
        return payload if isinstance(payload, dict) else None

    def put(self, suite: str, item_id: str, payload: Mapping[str, Any]) -> None:
        """Store (or replace) one finished item."""
        if self._conn is None:
            return
        self._conn.execute(
            "INSERT OR REPLACE INTO items (model_sha, suite, item_id, payload) VALUES (?, ?, ?, ?)",
            (self.key, suite, item_id, json.dumps(dict(payload), default=str)),
        )
        self._conn.commit()

    def count(self, suite: str | None = None) -> int:
        """How many items are stored for this model (optionally for one suite only)."""
        if self._conn is None:
            return 0
        if suite is None:
            query = "SELECT COUNT(*) FROM items WHERE model_sha = ?"
            args: tuple[str, ...] = (self.key,)
        else:
            query = "SELECT COUNT(*) FROM items WHERE model_sha = ? AND suite = ?"
            args = (self.key, suite)
        return int(self._conn.execute(query, args).fetchone()[0])

    def close(self) -> None:
        if self._conn is not None:
            self._conn.close()
            self._conn = None

    def __enter__(self) -> EvalCache:
        return self

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

"""Tests for rukh.eval.cache: keying by weights, isolation and skipping finished work."""

from __future__ import annotations

from pathlib import Path

import chess
import pytest

from rukh.eval.cache import EvalCache, config_sha, file_sha
from rukh.eval.puzzles import PuzzleItem, run_puzzles
from rukh.tokenize.uci_vocab import UciTokenizer

pytestmark = pytest.mark.unit

SCHOLAR = PuzzleItem(
    puzzle_id="p1",
    fen="r1bqkbnr/pppp1ppp/2n5/4p2Q/2B1P3/8/PPPP1PPP/RNB1K1NR b KQkq - 3 3",
    moves=["g8f6", "h5f7"],
    rating=1100,
    band="1000-1500",
)


@pytest.fixture
def tok() -> UciTokenizer:
    return UciTokenizer()


class CountingSource:
    """A scripted move source that records how many times it was asked."""

    def __init__(self, moves: list[str]) -> None:
        self.moves = moves
        self.calls = 0

    def __call__(self, board: chess.Board, history: list[int]) -> chess.Move | None:
        move = self.moves[min(self.calls, len(self.moves) - 1)]
        self.calls += 1
        return chess.Move.from_uci(move)

def test_put_and_get_round_trip(tmp_path: Path) -> None:
    with EvalCache(tmp_path / "cache.sqlite", "sha-a") as cache:
        cache.put("elo", "uci-1500:0", {"score": 1.0, "plies": 42})
        assert cache.get("elo", "uci-1500:0") == {"score": 1.0, "plies": 42}
        assert cache.get("elo", "uci-1500:1") is None
        assert cache.count("elo") == 1


def test_different_weights_do_not_share_results(tmp_path: Path) -> None:
    path = tmp_path / "cache.sqlite"
    with EvalCache(path, "sha-a") as first:
        first.put("puzzles", "p1", {"solved": True})
    with EvalCache(path, "sha-b") as second:
        assert second.get("puzzles", "p1") is None
    with EvalCache(path, "sha-a") as again:
        assert again.get("puzzles", "p1") == {"solved": True}


def test_a_disabled_cache_never_reads_or_writes(tmp_path: Path) -> None:
    path = tmp_path / "cache.sqlite"
    with EvalCache(path, "sha-a") as warm:
        warm.put("puzzles", "p1", {"solved": True})
    with EvalCache(path, "sha-a", enabled=False) as cold:
        assert cold.get("puzzles", "p1") is None
        cold.put("puzzles", "p2", {"solved": False})
        assert cold.count() == 0
    with EvalCache(path, "sha-a") as check:
        assert check.get("puzzles", "p2") is None

def test_a_cache_without_a_path_is_a_no_op() -> None:
    with EvalCache(None, "sha-a") as cache:
        cache.put("puzzles", "p1", {"solved": True})
        assert cache.get("puzzles", "p1") is None


def test_the_cache_stops_a_puzzle_from_being_replayed(tmp_path: Path, tok: UciTokenizer) -> None:
    source = CountingSource(["h5f7"])
    with EvalCache(tmp_path / "cache.sqlite", "sha-a") as cache:
        first = run_puzzles(source, tok, [SCHOLAR], cache=cache)
        assert source.calls == 1
        second = run_puzzles(source, tok, [SCHOLAR], cache=cache)
    assert source.calls == 1
    assert second.model_dump() == first.model_dump()


def test_without_a_cache_the_puzzle_is_replayed(tok: UciTokenizer) -> None:
    source = CountingSource(["h5f7"])
    run_puzzles(source, tok, [SCHOLAR])
    run_puzzles(source, tok, [SCHOLAR])
    assert source.calls == 2

def test_different_settings_do_not_share_results(tmp_path: Path) -> None:
    path = tmp_path / "cache.sqlite"
    slow = config_sha({"temperature": 0.6, "elo_move_time": 0.1})
    fast = config_sha({"temperature": 0.6, "elo_move_time": 0.05})
    assert slow != fast
    with EvalCache(path, "sha-a", config_sha=slow) as first:
        first.put("elo", "uci-1500:0", {"score": 1.0})
    with EvalCache(path, "sha-a", config_sha=fast) as second:
        assert second.get("elo", "uci-1500:0") is None  # the move time changed the games
        assert second.count("elo") == 0
    with EvalCache(path, "sha-a", config_sha=slow) as again:
        assert again.get("elo", "uci-1500:0") == {"score": 1.0}


def test_the_config_hash_ignores_key_order_and_reads_every_field() -> None:
    fields = {"temperature": 0.6, "top_k": 20, "seed": 42}
    assert config_sha(fields) == config_sha(dict(reversed(list(fields.items()))))
    assert config_sha(fields) != config_sha({**fields, "top_k": 10})

def test_the_weights_sha_is_still_what_the_report_shows(tmp_path: Path) -> None:
    cache = EvalCache(tmp_path / "cache.sqlite", "sha-a", config_sha="deadbeef")
    assert cache.model_sha == "sha-a"
    assert cache.key.startswith("sha-a:")
    cache.close()


def test_file_sha_identifies_the_weights(tmp_path: Path) -> None:
    first = tmp_path / "a.pt"
    second = tmp_path / "b.pt"
    first.write_bytes(b"weights")
    second.write_bytes(b"weights")
    assert file_sha(first) == file_sha(second)
    second.write_bytes(b"other")
    assert file_sha(first) != file_sha(second)

It also turns off ``preserve_insertion_order``. With it on (the default) a ``COPY`` of a large
streaming scan buffers the whole result so the output rows keep the input order; the first real
fetch of a month grew to 20 GB of resident memory for a 2.5 GB file. Nothing downstream cares
about the order of the games, so the buffering buys nothing and costs the machine.
"""

    memory_limit: str = "32GB"
    threads: int = Field(default=0, ge=0)
    """``0`` = ``cpu_count()``."""
    preserve_insertion_order: bool = False
    """Keep the input order in the output. Off: it buffers the whole result of a ``COPY``."""

    con.execute("SET temp_directory = ?", [temp_directory().as_posix()])
    con.execute("SET memory_limit = ?", [settings.memory_limit])
    con.execute("SET threads = ?", [settings.threads or (os.cpu_count() or 4)])
    con.execute("SET preserve_insertion_order = ?", [settings.preserve_insertion_order])