path: root/annotate_pgns/annotate.py

#!/usr/bin/env python3
"""Shard a PGN file across N single-threaded typhoon engine instances,
score every position in every game at a fixed search depth, and
reassemble an annotated PGN with { +score/depth PV } comments matching
the self-play annotation format already used in
../pgn/self_play_games.pgn.

Reusable across depths: --depth is a plain CLI argument, not baked into
any file name or code path, so re-running at sd12/sd16/etc. is just a
different --depth and --outdir.

Usage:
    python3 annotate.py --input ../pgn/twic_filtered.pgn \
        --depth 10 --shards 16 --hash 32m --outdir runs/sd10_twic

See the module-level functions below for the four stages: shard, build
engine scripts, run engines in parallel, parse + reassemble.
"""
import argparse
import os
import pty
import re
import select
import signal
import subprocess
import sys
import threading
import time
from pathlib import Path

import chess
import chess.pgn

from quiet_common import position_is_quiet

TELLOTHERS_RE = re.compile(
    r"^tellothers d(\d+), ([+-]?[\w.]+), [\d.]+s,\s*\d+knps, PV=(.*)$"
)
FLOAT_RE = re.compile(r"^[+-]?\d+\.\d+$")
MATE_RE = re.compile(r"^([+-])MATE(\d*)$")  # digit(s) optional -- the
# "FORCED MATE --> stop searching now" immediate-short-circuit path
# (search.c) prints a bare "+MATE"/"-MATE" with no distance number at
# all, not just the ordinary "+MATE<N>" case.

# --engine-protocol uci support: drives a real UCI engine (Stockfish)
# directly over stdin/stdout with the native UCI protocol -- no xboard,
# no PolyGlot adapter needed. Matches the standard "info depth N ...
# score cp/mate X ... pv ..." thinking-output line Stockfish prints once
# per completed depth.
UCI_ITER_RE = re.compile(
    r"^info depth (\d+) .*?\bscore (cp|mate) (-?\d+).*?\bpv (.*)$"
)
# A naive first attempt at this driver tried to interrupt each search
# early (xboard's "force", or UCI's "stop") the instant the target depth
# was seen, rather than waiting for "bestmove" -- motivated by "go depth
# N" appearing to hang forever with no bestmove ever printed. That hang
# turned out to be an artifact of this tool's own sandboxed subprocess
# environment, NOT a Stockfish or protocol issue (confirmed 2026-10-20:
# the identical "uci; isready; position fen ...; go depth 10" sequence
# got a normal, prompt "bestmove" reply when run directly in an
# interactive terminal outside that sandbox). So this driver just sends
# "go depth <depth>" and waits for "bestmove" like any normal UCI
# client -- no interrupt-and-abort workaround needed.
UCI_HANDSHAKE_TIMEOUT_SEC = 15
UCI_GO_TIMEOUT_SEC = 120  # per position, generous but bounded


def load_games(input_paths, max_games=None):
    """Load games, skipping ones with an unknown/unterminated result
    ('*') -- can't derive a WDL training label without a real outcome.
    Short-agreed-draw games are not separately filtered here: the
    >40-ply / >=2400-Elo filter already applied when twic_filtered.pgn
    was built (eval_tune/filter_pgn.py) rules those out by construction
    (confirmed empirically: no games below 40 plies in a 5000-game
    sample)."""
    games = []
    for path in input_paths:
        with open(path, encoding="utf-8", errors="replace") as f:
            while max_games is None or len(games) < max_games:
                game = chess.pgn.read_game(f)
                if game is None:
                    break
                if game.headers.get("Result", "*") == "*":
                    continue
                games.append(game)
        if max_games is not None and len(games) >= max_games:
            break
    return games


def shard_games(games, n_shards):
    shards = [[] for _ in range(n_shards)]
    for i, g in enumerate(games):
        shards[i % n_shards].append(g)
    return shards


def positions_for_game(game, prefilter_quiet=False):
    """Yield the FEN *before* each move is played, skipping any ply at
    or after the game is already over (no legal moves to search).

    If prefilter_quiet, also skip any ply whose position fails the
    quiet-position test (quiet_common.position_is_quiet) -- checked
    against the move actually played in this game, since that's the
    only move available before a search has run. This is the same
    check/exclude criteria as quiet_filter.py's post-hoc filter, but
    SEE-checked against the *played* move rather than the engine's own
    PV first move -- see quiet_common.py's module docstring for why
    those are different training distributions, not interchangeable.

    ply numbering (pi, used by write_shard_script/annotate_shard/
    build_games.py to re-align results) still counts only the plies
    actually yielded, so this must be called with the same
    prefilter_quiet value everywhere a run's ply indices are
    recomputed (build_games.py in particular)."""
    board = game.board()
    for move in game.mainline_moves():
        if board.is_game_over():
            break
        if not prefilter_quiet or position_is_quiet(board, move):
            yield board.fen()
        board.push(move)


def build_shard_positions(shard_games_list, prefilter_quiet=False):
    """Return (fens, index) for one shard: the flat list of FENs in
    script order, and the parallel index recording which
    (game_idx_in_shard, ply_idx) each FEN corresponds to -- shared by
    both write_shard_script (typhoon protocol, writes an engine command
    script to disk) and run_shard_uci (uci protocol, drives the engine
    interactively with no script file at all)."""
    fens = []
    index = []
    for gi, game in enumerate(shard_games_list):
        for pi, fen in enumerate(positions_for_game(game, prefilter_quiet)):
            fens.append(fen)
            index.append((gi, pi))
    return fens, index


def write_shard_script(shard_games_list, script_path, depth, prefilter_quiet=False):
    """Write the engine command script for one shard. The engine
    processes 'go' commands strictly in script order with no
    interleaving, so the index returned by build_shard_positions zips
    back onto the flat sequence of tellothers output lines."""
    fens, index = build_shard_positions(shard_games_list, prefilter_quiet)
    with open(script_path, "w") as out:
        for fen in fens:
            out.write(f"setboard {fen}\n")
            out.write(f"sd {depth}\n")
            out.write("go\n")
    return index


def run_shard(engine_path, script_path, out_path, hash_size, cpus):
    cmd = [
        str(engine_path),
        "--cpus", str(cpus),
        "--hash", hash_size,
        "--logfile", "-",
        "--batch",
        "--command",
        f"force; book name /nonexistent.book.bin; script {script_path}",
    ]
    with open(out_path, "w") as out:
        return subprocess.Popen(cmd, stdout=out, stderr=subprocess.STDOUT)


def _uci_format_tellothers(depth, score_kind, score_val, pv):
    """Render one uci-protocol result as a synthetic tellothers line, so
    the existing parse_shard_output/_format_comment/annotate_shard
    pipeline (written for typhoon's own tellothers output) can be reused
    unchanged for uci-mode results too. `pv` must already be SAN
    (see _uci_pv_to_san) -- Stockfish's own PV is UCI long algebraic
    ("e2e4 d7d5"), not the SAN typhoon's tellothers line uses ("e4
    d5"), and downstream readers (extract_positions.py,
    extract_selfplay_positions.py) key off SAN."""
    if score_kind == "mate":
        score_tok = "+MATE" if score_val > 0 else "-MATE"
    else:
        score_tok = f"{score_val / 100:+.2f}"
    return f"tellothers d{depth}, {score_tok}, 0.0s, 0knps, PV={pv}\n"


def _uci_pv_to_san(fen, uci_pv):
    """Convert a UCI-long-algebraic PV string ("e2e4 d7d5 ...") rooted
    at `fen` into typhoon's own tellothers SAN convention ("e4 d5
    ..."). Stops at the first move that doesn't parse/isn't legal
    (shouldn't happen for a PV the engine itself just reported, but an
    illegal/ambiguous tail is far less useful than no PV, not more)."""
    board = chess.Board(fen)
    san_moves = []
    for uci_move in uci_pv.split():
        try:
            move = chess.Move.from_uci(uci_move)
            if move not in board.legal_moves:
                break
            san_moves.append(board.san(move))
            board.push(move)
        except (ValueError, AssertionError):
            break
    return " ".join(san_moves)


def run_shard_uci(engine_path, fens, depth, out_path):
    """Drive one persistent native-UCI engine subprocess (e.g. a
    Stockfish binary) through every FEN in `fens`, in order, writing one
    synthetic tellothers line per FEN to out_path.

    Spawned under a pty (via pty.fork), not a plain subprocess.Popen
    pipe -- verified live (2026-10-20) that this particular Stockfish
    build hangs after printing just its banner + "id name" line, with
    no more output ever arriving, when its stdout is a plain pipe (the
    "id author"/options/"uciok" lines, and everything after, simply
    never get produced -- not merely buffered: a pty.fork-based read of
    the exact same binary gets the whole handshake instantly). The
    earlier guess that this was xboard/PolyGlot-specific, and later that
    it was this tool's own sandboxed subprocess environment, were both
    wrong -- it reproduces under dangerouslyDisableSandbox too. Root
    cause not fully isolated (likely an isatty()-gated code path in
    Stockfish's own startup or NUMA/processor-count detection), but the
    pty workaround is confirmed sufficient."""
    pid, fd = pty.fork()
    if pid == 0:
        try:
            os.execvp(str(engine_path), [str(engine_path)])
        finally:
            os._exit(127)  # os.execvp only returns on failure

    buf = b""

    def send(cmd):
        os.write(fd, (cmd + "\n").encode())

    def readline(timeout):
        nonlocal buf
        while b"\n" not in buf:
            r, _, _ = select.select([fd], [], [], timeout)
            if not r:
                return None
            try:
                chunk = os.read(fd, 65536)
            except OSError:
                return None
            if not chunk:
                return None
            buf += chunk
        line, buf = buf.split(b"\n", 1)
        # The pty's line-discipline translates outgoing "\n" to "\r\n"
        # (termios ONLCR), so every line read back here has a trailing
        # "\r" that a plain subprocess.Popen pipe would never produce --
        # strip it, or sentinel comparisons ("uciok", "readyok") and the
        # captured PV text silently never match/contain a stray "\r".
        return line.rstrip(b"\r").decode(errors="replace") + "\n"

    def wait_for(sentinel, timeout_sec):
        deadline = time.time() + timeout_sec
        while time.time() < deadline:
            line = readline(2)
            if line is None:
                continue
            if line.rstrip("\n") == sentinel:
                return True
        return False

    send("uci")
    wait_for("uciok", UCI_HANDSHAKE_TIMEOUT_SEC)
    # Match ../src/stockfish_xboard's own default: Stockfish's tbprobe.cpp
    # is Fathom-derived, so it reads the exact same .rtbw/.rtbz files
    # typhoon does -- no separate Stockfish-format tablebase needed.
    # SYZYGY_PATH="" explicitly disables TB probing rather than falling
    # back to the default, same as that wrapper.
    syzygy_path = os.environ.get("SYZYGY_PATH", "/zscratch/egtb")
    if syzygy_path:
        send(f"setoption name SyzygyPath value {syzygy_path}")
    send("isready")
    wait_for("readyok", UCI_HANDSHAKE_TIMEOUT_SEC)

    with open(out_path, "w") as out:
        for fen in fens:
            send(f"position fen {fen}")
            send(f"go depth {depth}")
            best = None
            deadline = time.time() + UCI_GO_TIMEOUT_SEC
            while time.time() < deadline:
                line = readline(2)
                if line is None:
                    continue
                line = line.rstrip("\n")
                if line.startswith("bestmove"):
                    break
                m = UCI_ITER_RE.match(line)
                if m:
                    d = int(m.group(1))
                    if d >= depth:
                        best = (d, m.group(2), int(m.group(3)), m.group(4))
                        # Keep reading up to "bestmove" even once the
                        # target depth is seen: with "go depth N",
                        # Stockfish always reports depth N as its final
                        # iteration before bestmove, so there's nothing
                        # more useful to capture, but exiting early here
                        # (like the earlier interrupt-based design did)
                        # would again require aborting the search instead
                        # of letting it finish on its own.
            # Always write exactly one line per "go", even if the
            # deadline was hit with no captured depth -- annotate_shard
            # zips index and parsed results positionally, so skipping a
            # write here would desync every later position in this
            # shard, not just leave this one unannotated (the same class
            # of bug the module docstring describes for a dropped
            # tellothers line). "NA" matches neither FLOAT_RE nor
            # MATE_RE, so _format_comment's fallback passes it through
            # as-is instead of mis-rendering it as a real score.
            if best is not None:
                d, kind, val, uci_pv = best
                out.write(_uci_format_tellothers(
                    d, kind, val, _uci_pv_to_san(fen, uci_pv)))
            else:
                out.write(f"tellothers d{depth}, NA, 0.0s, 0knps, PV=\n")

    send("quit")
    deadline = time.time() + 5
    while time.time() < deadline:
        wpid, _ = os.waitpid(pid, os.WNOHANG)
        if wpid == pid:
            break
        time.sleep(0.1)
    else:
        try:
            os.kill(pid, signal.SIGKILL)
            os.waitpid(pid, 0)
        except ProcessLookupError:
            pass
    os.close(fd)


def parse_shard_output(out_path):
    """Return a list of (depth, score_token, pv_string) in script
    order. score_token is kept as the engine's raw tellothers token
    (an ordinary "+1.23"-style float string, or "+MATE<N>"/"-MATE<N>")
    rather than eagerly parsed to float here -- a MATE token has no
    pawns magnitude to parse in the first place, and silently dropping
    any line that doesn't match a plain float is exactly the bug that
    used to desync this function's output from `index` (every position
    after a dropped mate/EGTB line in a shard got the wrong score) --
    see docs/NNUE.md's "Bug found and fixed in annotate.py's output
    parser" status entry, previously patched only in the separate
    extract_positions.py reader, not here."""
    results = []
    with open(out_path, encoding="utf-8", errors="replace") as f:
        for line in f:
            m = TELLOTHERS_RE.match(line.rstrip("\n"))
            if m:
                depth, score_tok, pv = m.groups()
                results.append((int(depth), score_tok, pv.strip()))
    return results


def _format_comment(depth, score_tok, pv):
    """Render one (depth, score_token, pv) result as a PGN comment in
    the same "{+-}SCORE/DEPTH PV" convention match_play.py's
    _format_eval_comment uses for self-play games (and
    extract_selfplay_positions.py's COMMENT_RE expects on read) --
    translates the engine's own raw "+MATE<N>"/"-MATE<N>" tellothers
    token into that convention's "+M"/"-M" (no distance digit, matching
    match_play.py exactly -- mate distance isn't independently verified
    here either). The PV's own sentinel token (<EGTB>/<REC>/<=>/<#>),
    when present, is passed through verbatim since that is exactly what
    both extract_positions.py and extract_selfplay_positions.py key
    their exactness detection off of."""
    m = MATE_RE.match(score_tok)
    if FLOAT_RE.match(score_tok):
        score_str = score_tok
    elif m:
        score_str = "+M" if m.group(1) == "+" else "-M"
    else:
        score_str = score_tok  # unrecognized token: pass through as-is
    comment = f"{score_str}/{depth}"
    if pv:
        comment += f" {pv}"
    return comment


def annotate_shard(shard_games_list, index, results):
    """Attach comments to each game's mainline nodes in place. index and
    results are both in script order and the same length; index maps
    each result back to (game_idx_in_shard, ply_idx)."""
    by_game = {}
    for (gi, pi), res in zip(index, results):
        by_game.setdefault(gi, {})[pi] = res

    for gi, game in enumerate(shard_games_list):
        scores = by_game.get(gi, {})
        node = game
        pi = 0
        while node.variations:
            node = node.variations[0]
            if pi in scores:
                depth, score_tok, pv = scores[pi]
                node.comment = _format_comment(depth, score_tok, pv)
            pi += 1


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--input", nargs="+", required=True,
                     help="one or more input PGN files")
    ap.add_argument("--depth", type=int, required=True,
                     help="fixed search depth (sd N) for every position")
    ap.add_argument("--shards", type=int, default=16)
    ap.add_argument("--hash", default="32m",
                     help="hash table size per shard engine, e.g. 32m")
    ap.add_argument("--cpus", type=int, default=1,
                     help="--cpus passed to each shard's engine")
    ap.add_argument("--engine", default="../src/typhoon")
    ap.add_argument("--engine-protocol", choices=["typhoon", "uci"],
                     default="typhoon",
                     help="'typhoon' drives --engine with its own "
                          "'--command \"script <path>\"' batch mode and "
                          "parses its tellothers output (default). 'uci' "
                          "drives --engine (a UCI engine binary, e.g. a "
                          "Stockfish build) directly over stdin/stdout "
                          "with the native UCI protocol instead ('uci', "
                          "'position fen ...', 'go depth N', wait for "
                          "'bestmove'). --hash/--cpus are ignored in this "
                          "mode; set hash/threads via UCI setoption "
                          "yourself if needed.")
    ap.add_argument("--outdir", required=True)
    ap.add_argument("--max-games", type=int, default=None,
                     help="limit to the first N games total, for smoke testing")
    ap.add_argument("--prefilter-quiet", action="store_true",
                     help="skip in-check positions and SEE-positive-capture "
                          "played moves before ever sending them to the "
                          "engine, instead of searching everything and "
                          "filtering after the fact (see quiet_filter.py / "
                          "quiet_filter_games.py). Changes the quiet "
                          "criterion's move source from the engine's own PV "
                          "to the move actually played -- see quiet_common.py.")
    args = ap.parse_args()

    outdir = Path(args.outdir)
    outdir.mkdir(parents=True, exist_ok=True)
    engine_path = Path(args.engine).resolve()
    if not engine_path.exists():
        sys.exit(f"engine not found: {engine_path}")

    print(f"Loading games from {args.input} ...")
    games = load_games(args.input, args.max_games)
    print(f"  {len(games)} games loaded")

    shards = shard_games(games, args.shards)

    if args.prefilter_quiet:
        print("Pre-filtering: skipping in-check / SEE-positive-played-move "
              "positions before search.")

    indices = []
    shard_fens = []  # only populated/used in uci mode
    if args.engine_protocol == "typhoon":
        print(f"Writing {args.shards} shard scripts at sd {args.depth} ...")
        for i, shard in enumerate(shards):
            script_path = outdir / f"shard_{i:02d}.script"
            idx = write_shard_script(shard, script_path, args.depth,
                                      args.prefilter_quiet)
            indices.append(idx)
            print(f"  shard {i:02d}: {len(shard)} games, {len(idx)} positions")
    else:
        print(f"Building {args.shards} shard position lists at sd "
              f"{args.depth} ...")
        for i, shard in enumerate(shards):
            fens, idx = build_shard_positions(shard, args.prefilter_quiet)
            shard_fens.append(fens)
            indices.append(idx)
            print(f"  shard {i:02d}: {len(shard)} games, {len(idx)} positions")

    print("Launching shard engines in parallel ...")
    t0 = time.time()
    if args.engine_protocol == "typhoon":
        procs = []
        for i in range(args.shards):
            script_path = outdir / f"shard_{i:02d}.script"
            out_path = outdir / f"shard_{i:02d}.out"
            if indices[i]:
                procs.append((i, run_shard(engine_path, script_path, out_path,
                                            args.hash, args.cpus)))

        # Wait for the slowest shard.
        for i, p in procs:
            rc = p.wait()
            if rc != 0:
                print(f"  WARNING: shard {i:02d} exited with code {rc}", file=sys.stderr)
    else:
        # No per-engine OS process launched up front here the way
        # run_shard does -- run_shard_uci owns its own subprocess
        # internally and blocks for the whole shard, so parallelism
        # across shards comes from running each one in its own thread
        # instead (these are I/O-bound waits on each engine subprocess's
        # pipe, so the GIL isn't a bottleneck here).
        threads = []
        for i in range(args.shards):
            if not indices[i]:
                continue
            out_path = outdir / f"shard_{i:02d}.out"
            t = threading.Thread(
                target=run_shard_uci,
                args=(engine_path, shard_fens[i], args.depth, out_path),
            )
            t.start()
            threads.append(t)
        for t in threads:
            t.join()
    elapsed = time.time() - t0
    total_positions = sum(len(idx) for idx in indices)
    print(f"All shards done in {elapsed:.1f}s "
          f"({total_positions / elapsed:.2f} positions/sec aggregate)")

    print("Parsing output and reassembling annotated PGN ...")
    final_path = outdir / "annotated.pgn"
    total_annotated = 0
    with open(final_path, "w") as out:
        for i, shard in enumerate(shards):
            out_path = outdir / f"shard_{i:02d}.out"
            results = parse_shard_output(out_path) if indices[i] else []
            if len(results) != len(indices[i]):
                print(f"  WARNING: shard {i:02d} expected {len(indices[i])} "
                      f"results, got {len(results)} -- some positions may be "
                      f"unannotated", file=sys.stderr)
            annotate_shard(shard, indices[i], results)
            total_annotated += len(results)
            for game in shard:
                print(game, file=out, end="\n\n")

    print(f"Wrote {final_path} ({len(games)} games, "
          f"{total_annotated} annotated positions)")


if __name__ == "__main__":
    main()