#!/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 re
import subprocess
import sys
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+), ([+-]?[\d.]+), [\d.]+s,\s*\d+knps, PV=(.*)$"
)


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 write_shard_script(shard_games_list, script_path, depth, prefilter_quiet=False):
    """Write the engine command script for one shard, plus the parallel
    index recording which (game_idx_in_shard, ply_idx) each 'go' line
    corresponds to -- the engine processes 'go' commands strictly in
    script order with no interleaving, so a flat index list zips back
    onto the flat sequence of tellothers output lines."""
    index = []
    with open(script_path, "w") as out:
        for gi, game in enumerate(shard_games_list):
            for pi, fen in enumerate(positions_for_game(game, prefilter_quiet)):
                out.write(f"setboard {fen}\n")
                out.write(f"sd {depth}\n")
                out.write("go\n")
                index.append((gi, pi))
    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 parse_shard_output(out_path):
    """Return a list of (depth, score, pv_string) in script order."""
    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, pv = m.groups()
                results.append((int(depth), float(score), pv.strip()))
    return results


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, pv = scores[pi]
                node.comment = f"{score:+.2f}/{depth} {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("--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.")
    print(f"Writing {args.shards} shard scripts at sd {args.depth} ...")
    indices = []
    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")

    print("Launching shard engines in parallel ...")
    t0 = time.time()
    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)
    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()