Estilo y texto PGN: los dos cortes del corpus

"""Tests for rukh.data.style: the game slices a LoRA adapter is trained to sound like."""

from __future__ import annotations

from pathlib import Path

import polars as pl
import pydantic
import pytest

from rukh.data.style import StyleConfig, StyleSpec, build_styles, predicate_of

pytestmark = pytest.mark.unit


def _stage(rukh_home: Path) -> None:
    """One month with three openings and two rating levels, all replayable."""
    rows = [
        ("e2e4 e7e5 g1f3", "C40", 2000),
        ("e2e4 c7c5 g1f3", "B20", 2100),
        ("e2e4 e7e6 d2d4", "C00", 1850),
        ("d2d4 d7d5 c2c4", "D06", 2000),
        ("d2d4 g8f6 c2c4", "E00", 1900),
        ("c2c4 e7e5", "A20", 2200),
    ]
    frame = pl.DataFrame(
        {
            "game_id": list(range(len(rows))),
            "uci": [uci for uci, _, _ in rows],
            "n_plies": [len(uci.split()) for uci, _, _ in rows],
            "white_elo": [elo for _, _, elo in rows],
            "black_elo": [elo for _, _, elo in rows],
            "result": ["1-0"] * len(rows),
            "time_control": ["300+0"] * len(rows),
            "eco": [eco for _, eco, _ in rows],
            "month": ["2025-01"] * len(rows),
        }
    )
    target = rukh_home / "data" / "uci" / "year=2025" / "month=01" / "games.parquet"
    target.parent.mkdir(parents=True)
    frame.write_parquet(target.as_posix())


def _read(rukh_home: Path, name: str) -> pl.DataFrame:
    return pl.read_parquet((rukh_home / "data" / "style" / f"{name}.parquet").as_posix())

def test_the_first_move_slice_only_holds_games_that_open_with_it(rukh_home: Path) -> None:
    """The axis the demo shows: swap the adapter, the model opens differently."""
    _stage(rukh_home)
    manifests = build_styles(
        StyleConfig(styles=[StyleSpec(name="e4", first_move="e2e4", n_games=10)])
    )
    frame = _read(rukh_home, "e4")
    assert frame.height == 3
    assert all(uci.startswith("e2e4 ") for uci in frame["uci"])
    assert manifests["e4"].counts == {"e4": 3}


def test_a_first_move_prefix_does_not_leak_into_another_move(rukh_home: Path) -> None:
    """`d2d4` must not also collect `d2d4` games that began `d2d3`; the match is on the token."""
    _stage(rukh_home)
    build_styles(StyleConfig(styles=[StyleSpec(name="d4", first_move="d2d4", n_games=10)]))
    frame = _read(rukh_home, "d4")
    assert sorted(frame["eco"].to_list()) == ["D06", "E00"]


def test_eco_prefixes_select_a_family_of_openings(rukh_home: Path) -> None:
    _stage(rukh_home)
    build_styles(StyleConfig(styles=[StyleSpec(name="sicilian", eco_prefix=["B2"], n_games=10)]))
    assert _read(rukh_home, "sicilian")["eco"].to_list() == ["B20"]


def test_a_rating_floor_keeps_a_style_adapter_from_being_a_weakness_adapter(
    rukh_home: Path,
) -> None:
    _stage(rukh_home)
    build_styles(
        StyleConfig(styles=[StyleSpec(name="e4", first_move="e2e4", min_elo=1900, n_games=10)])
    )
    assert _read(rukh_home, "e4").height == 2  # the 1850 game is out


def test_the_cap_limits_how_many_games_a_style_gets(rukh_home: Path) -> None:
    _stage(rukh_home)
    manifests = build_styles(
        StyleConfig(styles=[StyleSpec(name="e4", first_move="e2e4", n_games=2)])
    )
    assert manifests["e4"].counts == {"e4": 2}


def test_the_manifest_records_the_predicate_a_published_adapter_has_to_state(
    rukh_home: Path,
) -> None:
    _stage(rukh_home)
    manifests = build_styles(
        StyleConfig(styles=[StyleSpec(name="e4", first_move="e2e4", n_games=10)])
    )
    filters = manifests["e4"].filters
    assert "split_part(uci, ' ', 1) = 'e2e4'" in str(filters["predicate"])
    assert filters["style"]["first_move"] == "e2e4"  # type: ignore[index]

def test_several_styles_are_cut_in_one_pass(rukh_home: Path) -> None:
    _stage(rukh_home)
    manifests = build_styles(
        StyleConfig(
            styles=[
                StyleSpec(name="e4", first_move="e2e4", n_games=10),
                StyleSpec(name="d4", first_move="d2d4", n_games=10),
            ]
        )
    )
    assert set(manifests) == {"e4", "d4"}
    assert (rukh_home / "data" / "style" / "d4.manifest.json").is_file()


def test_the_slice_is_deterministic_for_a_seed(rukh_home: Path) -> None:
    _stage(rukh_home)
    spec = StyleSpec(name="e4", first_move="e2e4", n_games=2)
    build_styles(StyleConfig(styles=[spec], seed=1))
    first = _read(rukh_home, "e4")
    build_styles(StyleConfig(styles=[spec], seed=1))
    assert _read(rukh_home, "e4").equals(first)


def test_a_style_that_selects_nothing_is_refused() -> None:
    with pytest.raises(pydantic.ValidationError, match="selects nothing"):
        StyleSpec(name="empty")


def test_a_first_move_that_is_not_a_uci_move_is_refused() -> None:
    with pytest.raises(pydantic.ValidationError, match="not a UCI move"):
        StyleSpec(name="bad", first_move="e4")


def test_duplicate_style_names_are_refused() -> None:
    with pytest.raises(pydantic.ValidationError, match="duplicate style names"):
        StyleConfig(
            styles=[
                StyleSpec(name="e4", first_move="e2e4"),
                StyleSpec(name="e4", first_move="d2d4"),
            ]
        )


def test_a_quote_in_a_field_cannot_break_out_of_the_predicate() -> None:
    """Typed fields are the reason the config cannot carry SQL; the escaping is the backstop."""
    spec = StyleSpec(name="odd", eco_prefix=["B'2"])
    assert "'B''2%'" in predicate_of(spec)


def test_no_styles_is_refused(rukh_home: Path) -> None:
    with pytest.raises(ValueError, match="no styles configured"):
        build_styles(StyleConfig())

# The two style slices the LoRA adapters of M4 are trained on, and their token packs.
#
# The axis is White's first move, for one reason above the others: it is *visible*. A reader can
# load the demo, pick an adapter and watch the model open `1. e4`, then pick the other and watch
# it open `1. d4`. An adapter that shifted something real but invisible would be a better piece of
# machine learning and a worse piece of teaching.
#
# `min_elo: 1800` on both sides so a style adapter is not secretly a weakness adapter: if the
# `1. d4` games came from weaker players, "the adapter changed the opening" and "the adapter made
# it worse" would be the same measurement.
duckdb:
  memory_limit: 32GB
  threads: 0
style:
  sources: [data/uci]
  out_dir: data/style
  seed: 42
  styles:
    - name: e4
      first_move: e2e4
      min_elo: 1800
      n_games: 200000
    - name: d4
      first_move: d2d4
      min_elo: 1800
      n_games: 200000

# Token pack of the `e4` style slice, for `configs/train/lora-e4.yaml`.
#
# The validation split is the project's own -- the first 100 000 games of 2025-02 -- and its
# remainder does **not** join training: an adapter trained on one opening has to be measured
# against the same held-out games as everything else, and topping it up with a general corpus
# would dilute the very slice that defines the style.
tokenize:
  uci_dir: data/uci
  out_dir: data/tokens-style-e4
  train_months: []
  extra_train_parquets: ['data/style/e4.parquet']
  val_month: '2025-02'
  val_games: 100000
  val_remainder_trains: false
  stats_games: data/uci/year=2025/month=01/games.parquet
  max_len: 200

# Token pack of the `d4` style slice, for `configs/train/lora-d4.yaml`.

  out_dir: data/tokens-style-d4
  train_months: []
  extra_train_parquets: ['data/style/d4.parquet']

def _rows(path: Path, limit: int, offset: int = 0) -> Iterator[dict[str, object]]:
    import polars as pl

    frame = (
        pl.scan_parquet(path.as_posix())
        .select(["uci", "white_elo", "black_elo", "result"])
        .slice(offset, limit)
        .collect()
    )
    yield from frame.rows(named=True)


def _write_split(
    source: Path, target: Path, limit: int, max_plies: int, offset: int = 0
) -> tuple[int, int]:
    """Write one JSONL split; return ``(written, dropped)``."""
    written = dropped = 0
    target.parent.mkdir(parents=True, exist_ok=True)
    with target.open("w", encoding="utf-8", newline="\n") as fh:
        for row in _rows(source, limit, offset):
            text = game_text(
                str(row["uci"]),
                int(row["white_elo"] or 0),
                int(row["black_elo"] or 0),
                str(row["result"]),
                max_plies,
            )
            if text is None:
                dropped += 1
                continue
            fh.write(json.dumps({"text": text}, ensure_ascii=False) + "\n")
            written += 1
    return written, dropped

def build(cfg: PgnTextConfig) -> Manifest:
    """Write ``train.jsonl`` and ``val.jsonl`` plus the manifest that says what went in."""
    out_dir = resolve(cfg.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    train_source = resolve(cfg.train_games)
    val_source = resolve(cfg.val_games)
    for path in (train_source, val_source):
        if not path.is_file():
            raise FileNotFoundError(f"games not found at {path}")
    train_written, train_dropped = _write_split(
        train_source, out_dir / TRAIN_FILE, cfg.n_train, cfg.max_plies
    )
    val_written, val_dropped = _write_split(
        val_source, out_dir / VAL_FILE, cfg.n_val, cfg.max_plies, cfg.val_offset
    )
    log.info("wrote %s train and %s val samples", train_written, val_written)
    files = [
        FileHash(
            path=name, sha256=sha256_file(out_dir / name), bytes=(out_dir / name).stat().st_size
        )
        for name in (TRAIN_FILE, VAL_FILE)
    ]
    manifest = Manifest(
        dataset="rukh-pgn-text",
        months=[],
        filters={
            "train_games": cfg.train_games,
            "val_games": cfg.val_games,
            "max_plies": cfg.max_plies,
            "val_offset": cfg.val_offset,
            "dropped": {"train": train_dropped, "val": val_dropped},
        },
        counts={"train": train_written, "val": val_written},
        files=files,
    )
    (out_dir / "manifest.json").write_text(manifest.model_dump_json(indent=2) + "\n", "utf-8")
    return manifest

"""Tests for rukh.data.pgn_text: the same games, spelled the way a general LLM reads them."""

from __future__ import annotations

import json
import shutil
from pathlib import Path

import chess
import polars as pl
import pytest

from rukh.data.pgn_text import PgnTextConfig, build, game_text, movetext, prompt_for

pytestmark = pytest.mark.unit

RUY_LOPEZ = ["e2e4", "e7e5", "g1f3", "b8c6", "f1b5"]


def _board(moves: list[str]) -> chess.Board:
    board = chess.Board()
    for uci in moves:
        board.push(chess.Move.from_uci(uci))
    return board


def test_movetext_is_numbered_san() -> None:
    assert movetext(_board(RUY_LOPEZ)) == "1. e4 e5 2. Nf3 Nc6 3. Bb5"


def test_san_is_disambiguated_by_the_position_it_is_written_in() -> None:
    """The whole reason the movetext is replayed instead of read off the move stack.

    With both knights able to reach the square, the same move is `Nbd2`, not `Nd2`; only the
    position at that moment knows which letter is needed.
    """
    board = _board(["g1f3", "d7d5", "d2d4", "g8f6", "b1d2"])
    assert movetext(board).endswith("Nbd2")

def test_a_game_carries_its_ratings_and_its_result() -> None:
    text = game_text(" ".join(RUY_LOPEZ), 1800, 1900, "1-0", 120)
    assert text is not None
    assert text.startswith('[WhiteElo "1800"] [BlackElo "1900"]\n')
    assert text.endswith(" 1-0")
    assert "1. e4 e5" in text


def test_a_game_longer_than_the_cap_is_cut_not_dropped() -> None:
    text = game_text(" ".join(RUY_LOPEZ), 1800, 1800, "1-0", 4)
    assert text is not None
    assert text.rstrip(" 1-0").endswith("Nc6")


def test_an_unplayable_game_is_dropped_rather_than_repaired() -> None:
    assert game_text("e2e4 e2e4", 1800, 1800, "1-0", 120) is None
    assert game_text("", 1800, 1800, "1-0", 120) is None


def test_the_prompt_never_ends_on_a_space() -> None:
    """A BPE writes a move as " Nc6" with the space attached.

    A prompt that already carries the space forces the model onto the spelling it almost never
    saw in training, which would show up as a worse model rather than as a worse prompt.
    """
    for plies in range(len(RUY_LOPEZ) + 1):
        prompt = prompt_for(_board(RUY_LOPEZ[:plies]))
        assert prompt == prompt.rstrip()


def test_the_prompt_ends_with_the_move_number_when_white_is_to_move() -> None:
    assert prompt_for(_board(RUY_LOPEZ[:4])).endswith("Nc6 3.")


def test_the_prompt_ends_after_the_white_move_when_black_is_to_move() -> None:
    assert prompt_for(_board(RUY_LOPEZ)).endswith("3. Bb5")


def test_a_position_from_a_fen_carries_the_fen_and_the_ellipsis() -> None:
    """A puzzle has no history; asking the model to invent one is asking a different question."""
    fen = "r1bqkbnr/pppp1ppp/2n5/4p3/2B1P3/5N2/PPPP1PPP/RNBQK2R b KQkq - 5 4"
    prompt = prompt_for(chess.Board(fen), 2000, 2000)
    assert f'[FEN "{fen}"]' in prompt
    assert prompt.endswith("4...")


def test_the_start_position_carries_no_fen_tag() -> None:
    assert "FEN" not in prompt_for(chess.Board())


def test_the_ratings_asked_for_are_the_ratings_written() -> None:
    assert '[WhiteElo "1200"] [BlackElo "2400"]' in prompt_for(chess.Board(), 1200, 2400)

def test_build_writes_both_splits_and_counts_what_it_dropped(
    rukh_home: Path, repo_root: Path
) -> None:
    src = repo_root / "tests" / "fixtures" / "games.parquet"
    for month in ("01", "02"):
        target = rukh_home / "data" / "uci" / "year=2025" / f"month={month}" / "games.parquet"
        target.parent.mkdir(parents=True)
        shutil.copy(src, target)
    rows = pl.read_parquet(src.as_posix()).height

    manifest = build(PgnTextConfig(n_train=rows, n_val=3))
    out = rukh_home / "data" / "pgn-text"
    train = [json.loads(line) for line in (out / "train.jsonl").read_text("utf-8").splitlines()]
    val = [json.loads(line) for line in (out / "val.jsonl").read_text("utf-8").splitlines()]
    assert manifest.counts == {"train": len(train), "val": len(val)}
    assert len(val) <= 3
    assert all(set(sample) == {"text"} for sample in train)
    assert all(sample["text"].startswith("[WhiteElo ") for sample in train)
    assert manifest.filters["dropped"] == {"train": 0, "val": 0}


def test_the_validation_split_can_skip_the_rows_the_harness_evaluates_on(
    rukh_home: Path, repo_root: Path
) -> None:
    src = repo_root / "tests" / "fixtures" / "games.parquet"
    for month in ("01", "02"):
        target = rukh_home / "data" / "uci" / "year=2025" / f"month={month}" / "games.parquet"
        target.parent.mkdir(parents=True)
        shutil.copy(src, target)
    plain = build(PgnTextConfig(n_train=2, n_val=2)).counts["val"]
    out = rukh_home / "data" / "pgn-text"
    first = (out / "val.jsonl").read_text("utf-8")
    build(PgnTextConfig(n_train=2, n_val=2, val_offset=2))
    assert plain == 2
    assert (out / "val.jsonl").read_text("utf-8") != first


def test_missing_games_say_where_they_were_looked_for(rukh_home: Path) -> None:
    with pytest.raises(FileNotFoundError, match="games not found"):
        build(PgnTextConfig())

DUCKDB_STEPS = ("positions", "evals", "puzzles", "elo_bins", "style")


class PipelineConfig(BaseConfig):
    """All step configs; each step reads only its own section.

    ``duckdb:`` is shared: the top-level settings are copied into every step that opens a
    DuckDB connection unless that step declares its own.
    """

    duckdb: DuckDbConfig = Field(default_factory=DuckDbConfig)
    uci: UciConfig = Field(default_factory=UciConfig)
    tokenize: TokenizeConfig = Field(default_factory=TokenizeConfig)
    positions: PositionsConfig = Field(default_factory=PositionsConfig)
    evals: EvalsConfig = Field(default_factory=EvalsConfig)
    puzzles: PuzzlesConfig = Field(default_factory=PuzzlesConfig)
    pairs: PairsConfig = Field(default_factory=PairsConfig)
    elite: EliteConfig = Field(default_factory=EliteConfig)
    elo_bins: EloBinsConfig = Field(default_factory=EloBinsConfig)
    style: StyleConfig = Field(default_factory=StyleConfig)
    pgn_text: PgnTextConfig = Field(default_factory=PgnTextConfig)
    publish: PublishConfig = Field(default_factory=PublishConfig)

def build_pairs(evals: Path, min_delta_cp: int) -> pl.DataFrame:
    """One pair per FEN, reading the evaluations parquet batch by batch."""
    reader = pq.ParquetFile(evals)
    rows: list[dict[str, object]] = []
    columns = ["fen", "phase", "pvs", "game_id", "ply"]
    for batch in reader.iter_batches(batch_size=BATCH_ROWS, columns=columns):
        records = zip(
            batch.column("fen").to_pylist(),
            batch.column("phase").to_pylist(),
            batch.column("pvs").to_pylist(),
            batch.column("game_id").to_pylist(),
            batch.column("ply").to_pylist(),
            strict=True,
        )
        for fen, ph, pvs, game_id, ply in records:
            pair = make_pair(fen, pvs, min_delta_cp)
            if pair is not None:
                pair["phase"] = ph
                # Provenance, and the only way the decoder can use a pair at all: it reads a move
                # sequence, not a board, so a pair is useless to it without the prefix that led
                # there. ``game_id`` and ``ply`` are what join a pair back to its game (D-070).
                pair["game_id"] = game_id
                pair["ply"] = ply
                rows.append(pair)
    schema = {
        "fen": pl.String,
        "chosen": pl.String,
        "rejected": pl.String,
        "cp_chosen": pl.Int32,
        "cp_rejected": pl.Int32,
        "phase": pl.String,
        "game_id": pl.String,
        "ply": pl.Int32,
    }
    return pl.DataFrame(rows, schema=schema)

@data_app.command("pgn-text")
def data_pgn_text(config: PipelineOption = None, as_json: JsonOption = False) -> None:
    """Render the same games as PGN text, for the general-model comparison of M4."""
    from rukh.data.pgn_text import PgnTextConfig, build
    from rukh.data.pipeline import load_pipeline

    logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(message)s")
    cfg = load_pipeline(config).pgn_text if config is not None else PgnTextConfig()
    try:
        manifest = build(cfg)
    except FileNotFoundError as exc:
        typer.echo(f"error: {exc}", err=True)
        raise typer.Exit(code=1) from exc
    _echo_manifest(manifest, as_json, cfg.out_dir)


@data_app.command("style")
def data_style(config: PipelineOption = None, as_json: JsonOption = False) -> None:
    """Cut one parquet per style: the games a LoRA adapter is trained to sound like."""
    from rukh.data.pipeline import load_pipeline
    from rukh.data.style import build_styles

    logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(message)s")
    cfg = load_pipeline(config).style
    try:
        manifests = build_styles(cfg)
    except (FileNotFoundError, ValueError) as exc:
        typer.echo(f"error: {exc}", err=True)
        raise typer.Exit(code=1) from exc
    if as_json:
        typer.echo("[" + ", ".join(m.model_dump_json(indent=2) for m in manifests.values()) + "]")
        return
    typer.echo(f"out_dir:  {cfg.out_dir}")
    typer.echo("styles:")
    for name, manifest in manifests.items():
        typer.echo(f"  {name}: {manifest.counts[name]:,} games")