#!/usr/bin/env python3
"""
Texel-style tuner for typhoon's eval.c constants, driving the engine
as an xboard-protocol subprocess and using its own `evaldna` command
as the read/write interface for the parameter vector.

Requires: python-chess (pip install chess)  -- only for PGN parsing.

This is a skeleton: the pieces that are engine-specific (parsing the
score out of xboard "thinking output" lines, the exact evaldna file
format) are marked TODO/ADAPT and need to be checked against a real
run of the engine.
"""
import atexit
import hashlib
import re
import subprocess
import random
import math
import sys

from dna_diff import DNA_NAMES
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import chess
import chess.pgn

# Python fully block-buffers stdout when it isn't a tty (e.g. redirected
# to a log file via `> run.log`), so print()s otherwise sit in a buffer
# until it fills or the process exits -- on a multi-hour tuning run,
# `tail -f run.log` can look totally stuck even though the process is
# making progress. Force line buffering so every print() lands
# immediately regardless of where stdout is pointed.
sys.stdout.reconfigure(line_buffering=True)


# ---------------------------------------------------------------------------
# 1. Engine process wrapper (xboard protocol)
# ---------------------------------------------------------------------------

class Engine:
    """A live typhoon subprocess.

    IMPORTANT: always use `with Engine(...) as engine:` (or
    EnginePool) rather than a bare constructor call. The engine holds
    a kernel SysV semaphore (semget(IPC_PRIVATE, ...) in unix.c) for
    its input-dispatch loop; that semaphore is only released by a
    clean `quit`, not by the OS when the process dies (kill/pkill
    leaves it behind). This bit us for real: kern.ipc.semmni on this
    box is only 50, a night of force-killed test engines leaked past
    that limit, and every engine spawned afterward silently fell back
    to a 100ms-per-command polling path (_WaitUntilTheresInputToRead
    in input.c) instead of the fast semaphore wait -- a ~1300x
    slowdown that looked like "the machine is contended" for hours
    before the real cause was found. quit()/__exit__ here are
    deliberately defensive (never raise, always fall through to a
    hard kill) so a crash mid-tuning-run can't repeat that.
    """

    def __init__(self, path, dna_scratch_dir, extra_args=None):
        # Default to no logfile: every spawn site here (EnginePool's
        # n_workers instances, self-play generation) is a short-lived
        # scratch subprocess, and multiple of them writing/backing-up
        # the same default logfile concurrently is exactly the lock-
        # contention bug match_play.py hit and fixed with this same
        # flag (see its "--logfile matters here" comment). A caller
        # that actually wants a logfile can still pass extra_args
        # explicitly to override this default.
        if extra_args is None:
            extra_args = ["--logfile", "-"]
        self.proc = subprocess.Popen(
            [path] + extra_args,
            stdin=subprocess.PIPE, stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL, text=True, bufsize=1,
        )
        assert self.proc.stdin is not None and self.proc.stdout is not None
        self.stdin = self.proc.stdin
        self.stdout = self.proc.stdout
        self.dna_scratch_dir = Path(dna_scratch_dir)
        self.dna_scratch_dir.mkdir(parents=True, exist_ok=True)
        self._send("xboard")
        self._send("protover 2")
        self._drain_features()
        self._send("force")          # we set positions, engine doesn't play
        atexit.register(self._hard_cleanup)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.quit()
        return False

    def _send(self, cmd):
        self.stdin.write(cmd + "\n")
        self.stdin.flush()

    def _drain_features(self, timeout_lines=200):
        # ADAPT: real code should use a select()/timeout loop; this just
        # reads until it sees "done=1" or a line count cap.
        for _ in range(timeout_lines):
            line = self.stdout.readline()
            if not line or "done=1" in line:
                break

    def quit(self):
        """Best-effort graceful shutdown, falling through to a hard
        kill -- this must never raise and must never leave the process
        (and its semaphore) alive, no matter what state it's in."""
        try:
            self._send("quit")
            self.proc.wait(timeout=5)
        except Exception:
            self._hard_cleanup()
        finally:
            atexit.unregister(self._hard_cleanup)

    def _hard_cleanup(self):
        if self.proc.poll() is None:
            try:
                self.proc.kill()
                self.proc.wait(timeout=5)
            except Exception:
                pass

    # -- DNA I/O -------------------------------------------------------

    def load_dna(self, flat_values, row_lengths):
        """Write flat_values (grouped back into g_EvalDNA's per-array
        rows) into a fresh .dna file and tell the engine to load it
        via `evaldna read <file>`. Verified live: on success the
        engine replies exactly `Loaded dna file "<path>"`."""
        rows = []
        i = 0
        for n in row_lengths:
            rows.append(",".join(str(int(round(v))) for v in flat_values[i:i + n]))
            i += n
        fname = self.dna_scratch_dir / f"cand_{random.randrange(10**9)}.dna"
        fname.write_text("\n".join(rows) + "\n")
        self._send(f"evaldna read {fname}")
        ack = self.stdout.readline()
        fname.unlink(missing_ok=True)
        return "Loaded dna file" in ack

    def dump_dna(self):
        """Parse the flat vector back out via `evaldna dump` (calls
        ExportEvalDNA under the hood) -- used once at startup to learn
        the baseline vector and each array's length/order.

        Verified live: ExportEvalDNA's result is one Trace() call
        containing embedded newlines, but on the pipe that arrives as
        one *separate* readline() per array -- the first prefixed with
        "EvalDNA: ", the rest bare -- followed by one blank line once
        all arrays (one per g_EvalDNA entry) have been emitted.
        """
        self._send("evaldna dump")
        first = self.stdout.readline()
        assert first.startswith("EvalDNA: "), first
        rows = [first[len("EvalDNA: "):].strip()]
        while True:
            line = self.stdout.readline()
            if not line or not line.strip():
                break
            rows.append(line.strip())
        flat = []
        row_lengths = []
        for row in rows:
            vals = [int(x) for x in row.split(",")]
            flat.extend(vals)
            row_lengths.append(len(vals))
        return flat, row_lengths

    # -- Static eval -----------------------------------------------------

    # command.c: EvalCommand -> Trace("Static eval: %s\n", ScoreToString(i))
    # util.c: ScoreToString -> "+2.34" / "-1.05" / "+MATE7" / "-MATE20"
    _EVAL_RE = re.compile(r"Static eval:\s*([+-])(\d+)\.(\d+)")
    _MATE_RE = re.compile(r"Static eval:\s*([+-])MATE(\d+)")

    def eval_fen(self, fen):
        """Return Eval()'s raw score in centipawns, from the
        side-to-move's perspective -- no search involved at all, this
        is the literal static eval used by the search's leaf nodes.

        Single-position convenience wrapper; prefer eval_fens_batched
        for anything scoring more than a handful of positions -- see
        its docstring for why the per-call version is slow at scale.
        """
        return self.eval_fens_batched([fen])[0]

    def eval_fens_batched(self, fens):
        """Score many FENs per round-trip instead of one.

        eval_fen()'s pattern -- write, flush, block on readline, repeat
        -- pays a full pipe-write/wakeup/dispatch/reply latency cycle
        per position, even though `eval` with no search is near-instant
        computation. Measured live: batching turns N of those
        round-trip latencies into effectively one, which is where
        nearly all the wall-clock time was going at corpus scale.

        Still a single already-running engine process -- this isn't
        about parallelism, just not idling between writes.
        """
        cmds = "".join(f"setboard {fen}\neval\n" for fen in fens)
        self.stdin.write(cmds)
        self.stdin.flush()

        scores = []
        for _ in fens:
            score = None
            for _ in range(50):
                line = self.stdout.readline()
                if not line:
                    break
                m = self._EVAL_RE.search(line)
                if m:
                    sign = -1 if m.group(1) == "-" else 1
                    score = sign * (int(m.group(2)) * 100 + int(m.group(3)))
                    break
                if self._MATE_RE.search(line):
                    # Mate scores are meaningless to the sigmoid fit --
                    # the caller should treat these as "skip this
                    # position" (None), same as eval_fen()'s contract.
                    break
            scores.append(score)
        return scores

    # -- Quiescence search -------------------------------------------------

    # command.c: QSearchCommand -> Trace("Qsearch score: %s\n", ...)
    # Same ScoreToString format as Static eval, different label.
    _QSEARCH_RE = re.compile(r"Qsearch score:\s*([+-])(\d+)\.(\d+)")
    _QSEARCH_MATE_RE = re.compile(r"Qsearch score:\s*([+-])MATE(\d+)")

    def qsearch_fens_batched(self, fens):
        """Like eval_fens_batched, but resolves captures/checks/promotions
        first (engine's own QSearch()) instead of scoring the position
        as-is. Side-to-move POV, same as eval_fens_batched -- directly
        comparable without flipping."""
        cmds = "".join(f"setboard {fen}\nqsearch\n" for fen in fens)
        self.stdin.write(cmds)
        self.stdin.flush()

        scores = []
        for _ in fens:
            score = None
            for _ in range(50):
                line = self.stdout.readline()
                if not line:
                    break
                m = self._QSEARCH_RE.search(line)
                if m:
                    sign = -1 if m.group(1) == "-" else 1
                    score = sign * (int(m.group(2)) * 100 + int(m.group(3)))
                    break
                if self._QSEARCH_MATE_RE.search(line):
                    break
            scores.append(score)
        return scores


class EnginePool:
    """N independent `typhoon` processes, splitting a batch of
    positions into shards scored in parallel across N cores.

    NOTE on history: per-position eval cost was measured at ~20ms
    early on and this class was originally justified as parallelizing
    "real compute inside Eval()". That measurement was itself an
    artifact of a since-fixed bug (see Engine's docstring) -- with
    that fixed, real steady-state cost is ~0.15ms/position, not 20ms.
    Multiple processes may still help throughput, but the case for
    this class is weaker now than when it was written; measure before
    trusting it, the same way n_workers=4/8 turned out to be unreliable
    earlier in development for reasons that had nothing to do with
    genuine CPU parallelism.

    Duck-types the same dump_dna/load_dna/eval_fens_batched/quit
    interface as Engine, so tune()/compute_error()/fit_k() work
    unchanged against either one.
    """

    def __init__(self, engine_path, dna_scratch_dir, n_workers=2):
        self.engines = [
            Engine(engine_path, f"{dna_scratch_dir}/worker{i}")
            for i in range(n_workers)
        ]
        self.pool = ThreadPoolExecutor(max_workers=n_workers)
        atexit.register(self._hard_cleanup)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.quit()
        return False

    def dump_dna(self):
        # All engines start with identical built-in constants; only
        # need one copy of the baseline vector/shape.
        return self.engines[0].dump_dna()

    def load_dna(self, flat_values, row_lengths):
        # Every worker must be scoring under the same candidate DNA,
        # so broadcast it to all of them before the next batch.
        list(self.pool.map(
            lambda e: e.load_dna(flat_values, row_lengths), self.engines
        ))

    def eval_fens_batched(self, fens):
        n = len(self.engines)
        shards = [fens[i::n] for i in range(n)]  # interleaved, not chunked,
        # so a batch dominated by one game's consecutive plies doesn't
        # pile onto a single worker
        shard_results = list(self.pool.map(
            lambda pair: pair[0].eval_fens_batched(pair[1]),
            zip(self.engines, shards),
        ))
        # Un-interleave back to the original order.
        scores = [None] * len(fens)
        for worker_idx, results in enumerate(shard_results):
            for j, score in enumerate(results):
                scores[worker_idx + j * n] = score
        return scores

    def qsearch_fens_batched(self, fens):
        n = len(self.engines)
        shards = [fens[i::n] for i in range(n)]
        shard_results = list(self.pool.map(
            lambda pair: pair[0].qsearch_fens_batched(pair[1]),
            zip(self.engines, shards),
        ))
        scores = [None] * len(fens)
        for worker_idx, results in enumerate(shard_results):
            for j, score in enumerate(results):
                scores[worker_idx + j * n] = score
        return scores

    def quit(self):
        try:
            list(self.pool.map(lambda e: e.quit(), self.engines))
        finally:
            self.pool.shutdown(wait=True)
            atexit.unregister(self._hard_cleanup)

    def _hard_cleanup(self):
        for e in self.engines:
            e._hard_cleanup()


# ---------------------------------------------------------------------------
# 2. Training set: quiet positions + game outcome labels
# ---------------------------------------------------------------------------
#
# Quiet-position filtering is done here in pure Python via a small SEE
# implementation rather than by driving the engine (e.g. `sd N` + `go`
# and comparing to `eval`). That route was tried live and rejected:
# it needs the opening book disabled, a time budget tuned so the
# search actually completes an iteration, and even then a genuinely
# tactical position that resolves into a drawn/insufficient-material
# endgame reports an eval swing of exactly 0 (the engine's endgame
# recognizer and Syzygy probe are *correct* to say that, but it makes
# "did the score change" useless as a tactic detector) -- plus it costs
# a full subprocess round-trip per position at corpus scale. A local
# SEE avoids all of that and is what production Texel-style tuners
# actually use for this.

PIECE_VALUES = {
    chess.PAWN: 100, chess.KNIGHT: 320, chess.BISHOP: 330,
    chess.ROOK: 500, chess.QUEEN: 900, chess.KING: 20000,
}


def _attackers_sorted_by_value(board, color, square):
    """Ascending list of piece values attacking `square` for `color`,
    lowest-value attacker first (the one SEE always uses next)."""
    values = []
    for sq in board.attackers(color, square):
        piece = board.piece_at(sq)
        values.append(PIECE_VALUES[piece.piece_type])
    values.sort()
    return values


def static_exchange_eval(board, move):
    """Classic swap-list SEE: net material value (in centipawns) the
    side to move nets by playing `move` and then trading off on the
    target square, assuming both sides always recapture with their
    least valuable attacker. Ignores pins/discoveries/x-rays through
    already-moved pieces -- a standard, accepted simplification for
    bulk position filtering (a full SEE needs to re-derive attackers
    after each simulated capture; this approximation is what most
    tuners use since it's cheap and rarely wrong by more than one
    trade).

    Reference: the "gain[]" swap-list algorithm described on the
    Chess Programming Wiki's SEE page.
    """
    square = move.to_square
    target = board.piece_at(square)
    if board.is_en_passant(move):
        first_captured_value = PIECE_VALUES[chess.PAWN]
    elif target is not None:
        first_captured_value = PIECE_VALUES[target.piece_type]
    else:
        return 0  # not a capture

    mover = board.piece_at(move.from_square)
    attacker_value = PIECE_VALUES[mover.piece_type]

    attackers = {
        chess.WHITE: _attackers_sorted_by_value(board, chess.WHITE, square),
        chess.BLACK: _attackers_sorted_by_value(board, chess.BLACK, square),
    }
    # The moving piece itself is one of the "attackers" of its own
    # from-square's list -- remove it since it's the one now occupying
    # the target square, about to potentially be recaptured.
    side = board.turn
    if attacker_value in attackers[side]:
        attackers[side].remove(attacker_value)

    gain = [first_captured_value]
    stm = not side  # opponent moves next, deciding whether to recapture
    captured_value = attacker_value  # value of the piece now sitting on `square`
    while attackers[stm]:
        next_attacker_value = attackers[stm].pop(0)
        gain.append(captured_value - gain[-1])
        captured_value = next_attacker_value
        stm = not stm

    # Back up the minimax: at each step a side only continues the
    # capture sequence if doing so improves on stopping.
    for i in range(len(gain) - 2, -1, -1):
        gain[i] = -max(-gain[i], gain[i + 1])
    return gain[0]


def has_hanging_material(board, threshold=50):
    """True if the side to move has a capture available that nets more
    than `threshold` centipawns by SEE -- i.e. this position is
    "loud": there's a free or winning capture sitting on the board, so
    its static eval shouldn't be trusted as a quiet training label."""
    for move in board.legal_moves:
        if not board.is_capture(move):
            continue
        if static_exchange_eval(board, move) >= threshold:
            return True
    return False


def filter_by_qsearch_agreement(engine, positions, max_diff_cp=30,
                                 batch_size=5000):
    """Second, stricter quiet-position gate on top of has_hanging_material's
    SEE check: keep only positions where the engine's own static Eval()
    is within `max_diff_cp` centipawns of its QSearch() -- i.e. resolving
    captures/checks/promotions barely moves the score, so the static eval
    used as this position's training label is actually representative of
    what the engine would settle on after searching.

    Unlike SEE (a Python approximation ignoring pins/discoveries/x-rays,
    see static_exchange_eval's docstring), this drives the real engine's
    real move generator and search logic via the `qsearch` command, so it
    catches quiet-looking-by-SEE positions that are still tactically loud
    for reasons SEE can't see. Costs one eval + one qsearch round trip per
    position, so run it on an already SEE-filtered, already-capped pool,
    not the raw PGN scan (that scan stays SEE-only for the reasons in
    load_training_positions' module comment: qsearch there would need the
    book disabled, a real time budget, and still return a useless 0 swing
    for tactics that resolve into recognized draws/endgames)."""
    kept = []
    n_dropped = 0
    for i in range(0, len(positions), batch_size):
        chunk = positions[i:i + batch_size]
        fens = [fen for fen, _ in chunk]
        evals = engine.eval_fens_batched(fens)
        qsearches = engine.qsearch_fens_batched(fens)
        for (fen, result), e, q in zip(chunk, evals, qsearches):
            if e is None or q is None:
                n_dropped += 1
                continue
            if abs(e - q) <= max_diff_cp:
                kept.append((fen, result))
            else:
                n_dropped += 1
    print(f"  qsearch-agreement filter: kept {len(kept)}/{len(positions)} "
          f"(dropped {n_dropped} where |eval-qsearch| > {max_diff_cp}cp "
          f"or either was a mate score)")
    return kept


def load_training_positions(pgn_path, max_positions=200_000, sample_every=8,
                             min_elo=2400, game_stride=1,
                             cache_dir="position_cache", force_rescan=False):
    """Cached wrapper around _load_training_positions_uncached(). The
    PGN scan + per-position SEE quiet-check is the single most
    expensive one-time cost in the whole pipeline (minutes, scaling
    with max_positions) and its output depends only on the arguments
    below plus the PGN file's contents -- so cache it keyed on both,
    rather than repeating the scan on every tune() run.

    Cache key includes the source file's mtime+size specifically so a
    replaced/updated TWIC file invalidates old cache entries
    automatically instead of silently serving stale positions.
    """
    st = Path(pgn_path).stat()
    key_material = (
        f"{pgn_path}|{st.st_mtime}|{st.st_size}|{max_positions}|"
        f"{sample_every}|{min_elo}|{game_stride}"
    )
    key = hashlib.sha256(key_material.encode()).hexdigest()[:16]
    cache_path = Path(cache_dir) / f"positions_{key}.tsv"

    if cache_path.exists() and not force_rescan:
        positions = []
        with open(cache_path) as f:
            for line in f:
                fen, result = line.rstrip("\n").split("\t")
                positions.append((fen, float(result)))
        print(f"  loaded {len(positions)} positions from cache {cache_path}")
        return positions

    positions = _load_training_positions_uncached(
        pgn_path, max_positions=max_positions, sample_every=sample_every,
        min_elo=min_elo, game_stride=game_stride,
    )
    cache_path.parent.mkdir(parents=True, exist_ok=True)
    with open(cache_path, "w") as f:
        for fen, result in positions:
            f.write(f"{fen}\t{result}\n")
    print(f"  cached {len(positions)} positions to {cache_path}")
    return positions


def _load_training_positions_uncached(pgn_path, max_positions=200_000,
                                       sample_every=8, min_elo=2400,
                                       game_stride=1):
    """Extract (fen, result_from_white_pov) pairs, skipping the opening,
    any position where a capture/check just happened, and any position
    with a hanging piece per has_hanging_material()'s SEE check.

    min_elo: skip games where either player is unrated or below this
    (TWIC-scale corpora span club players to super-GMs; blunder-laden
    games from weak players pollute the outcome labels without adding
    anything an eval function should learn from).

    game_stride: only keep 1 game out of every N seen (after the Elo
    filter) -- TWIC is in chronological order, so straight truncation
    at max_positions would bias the sample toward one era. Set this
    based on a quick game count for the corpus so the stride spans the
    whole file (e.g. 554135 games, want ~40k of them -> stride ~14).
    """
    positions = []
    game_index = 0
    # TWIC-scale PGN corpora are decades of weekly dumps from many
    # different editors/eras concatenated together and are not
    # uniformly UTF-8 (e.g. smart quotes under Windows-1252 in a
    # player name or comment) -- replace, don't crash, on a bad byte;
    # it's essentially never in board data that would affect a FEN.
    with open(pgn_path, encoding="utf-8", errors="replace") as fh:
        while len(positions) < max_positions:
            game = chess.pgn.read_game(fh)
            if game is None:
                break

            result = game.headers.get("Result", "*")
            if result == "1-0":
                r = 1.0
            elif result == "0-1":
                r = 0.0
            elif result == "1/2-1/2":
                r = 0.5
            else:
                continue

            try:
                white_elo = int(game.headers.get("WhiteElo", "0"))
                black_elo = int(game.headers.get("BlackElo", "0"))
            except ValueError:
                continue
            if white_elo < min_elo or black_elo < min_elo:
                continue

            game_index += 1
            if game_stride > 1 and (game_index % game_stride) != 0:
                continue

            board = game.board()
            ply = 0
            for move in game.mainline_moves():
                was_capture = board.is_capture(move)
                board.push(move)
                ply += 1
                if is_quiet_training_ply(board, ply, was_capture, sample_every):
                    positions.append((board.fen(), r))
    return positions


def is_quiet_training_ply(board, ply, was_capture, sample_every):
    """Shared quiet-position criteria used by both the TWIC PGN path
    and self-play game generation, so the two position sources are
    filtered identically: skip the opening, positions right after a
    capture, positions in check, and positions with a hanging piece."""
    return (
        ply > 10
        and ply % sample_every == 0
        and not was_capture
        and not board.is_check()
        and not has_hanging_material(board)
    )


# ---------------------------------------------------------------------------
# 2b. Self-play position generation (typhoon vs. itself, via xboard)
# ---------------------------------------------------------------------------
#
# Rationale (discussed at length): TWIC labels reflect what wins games
# for strong *humans*, which only approximately matches what wins
# games for typhoon's own search. Self-play sidesteps that -- the
# label is literally "did typhoon's own play, from this position, win
# the game" -- at the cost of needing to actually play out games
# instead of reusing an existing corpus.
#
# Diversity comes from the engine's own opening book: it makes
# weighted-random line selections (confirmed: `book move Bb5 [+2684
# =2901 -1737]`-style output), so successive games naturally diverge
# in the opening even though the search itself is otherwise
# deterministic. No artificial randomization is added here.

_MOVE_RE = re.compile(r"^move (\S+)")


def play_selfplay_game(engine_path, dna_scratch_dir, time_control_sec=1,
                        max_plies=200):
    """Play one full game of typhoon vs. itself via xboard's go/force
    turn-taking loop (verified live: each `go` plays exactly one move
    for whichever side is currently on move, updates the engine's own
    internal position, and a subsequent `go` after `force` correctly
    continues -- no need to feed the opponent's move back manually).

    Returns (list of (fen, ply) after each move, final_result) where
    final_result is 1.0/0.5/0.0 from White's perspective, determined
    locally via python-chess rather than trusting engine-side game-end
    detection.
    """
    with Engine(engine_path, dna_scratch_dir) as engine:
        engine._send(f"st {time_control_sec}")
        board = chess.Board()
        fens_by_ply = []

        for ply in range(1, max_plies + 1):
            engine._send("go")
            move_uci = None
            for _ in range(200):
                line = engine.stdout.readline()
                if not line:
                    break
                m = _MOVE_RE.match(line)
                if m:
                    move_uci = m.group(1)
                    break
            engine._send("force")
            if move_uci is None:
                break  # engine had nothing to play -- treat as game over

            try:
                move = chess.Move.from_uci(move_uci)
                if move not in board.legal_moves:
                    break
            except ValueError:
                break

            board.push(move)
            fens_by_ply.append(board.fen())
            if board.is_game_over():
                break

        outcome = board.outcome(claim_draw=True)
        if outcome is None or outcome.winner is None:
            result = 0.5
        else:
            result = 1.0 if outcome.winner == chess.WHITE else 0.0

    return fens_by_ply, result


def generate_selfplay_positions(engine_path, n_games, dna_scratch_dir,
                                 time_control_sec=1, sample_every=8,
                                 cache_dir="position_cache"):
    """Play n_games of self-play and extract quiet (fen, result) pairs
    using the identical filter TWIC positions go through. Cached the
    same way and for the same reason as load_training_positions --
    playing games is far more expensive per position than parsing an
    existing PGN, so it's even more worth not repeating on every run.
    """
    key = hashlib.sha256(
        f"selfplay|{n_games}|{time_control_sec}|{sample_every}".encode()
    ).hexdigest()[:16]
    cache_path = Path(cache_dir) / f"selfplay_{key}.tsv"
    if cache_path.exists():
        positions = []
        with open(cache_path) as f:
            for line in f:
                fen, result = line.rstrip("\n").split("\t")
                positions.append((fen, float(result)))
        print(f"  loaded {len(positions)} self-play positions from cache {cache_path}")
        return positions

    positions = []
    for g in range(n_games):
        fens_by_ply, result = play_selfplay_game(
            engine_path, f"{dna_scratch_dir}/game{g}", time_control_sec
        )
        board = chess.Board()
        for ply, fen in enumerate(fens_by_ply, start=1):
            board.set_fen(fen)
            prev_board = chess.Board(fens_by_ply[ply - 2]) if ply > 1 else chess.Board()
            was_capture = len(board.piece_map()) < len(prev_board.piece_map())
            if is_quiet_training_ply(board, ply, was_capture, sample_every):
                positions.append((fen, result))
        print(f"  game {g+1}/{n_games}: {len(fens_by_ply)} plies, "
              f"result={result}, {len(positions)} quiet positions so far")

    cache_path.parent.mkdir(parents=True, exist_ok=True)
    with open(cache_path, "w") as f:
        for fen, result in positions:
            f.write(f"{fen}\t{result}\n")
    print(f"  cached {len(positions)} self-play positions to {cache_path}")
    return positions


# ---------------------------------------------------------------------------
# 3. Error function: scaled-sigmoid MSE against game outcome
# ---------------------------------------------------------------------------

def eval_white_pov_batch(engine, fens):
    """Eval() returns a score relative to the side to move; flip each
    one back to White's perspective so it lines up with the
    game-result labels (1.0 = White won). Batched -- see
    Engine.eval_fens_batched for why this matters at corpus scale."""
    scores = engine.eval_fens_batched(fens)
    out = []
    for fen, score in zip(fens, scores):
        if score is None:
            out.append(None)
        else:
            stm_is_white = fen.split()[1] == "w"
            out.append(score if stm_is_white else -score)
    return out


def fit_k(engine, sample):
    """One-time calibration of the logistic scale K so that
    sigmoid(eval/K) tracks empirical win rate. Simple ternary search
    on K, minimizing MSE over `sample`."""
    evals = eval_white_pov_batch(engine, [fen for fen, _ in sample])
    results = [r for _, r in sample]

    def mse(k):
        err, n = 0.0, 0
        for e, r in zip(evals, results):
            if e is None:
                continue
            p = 1.0 / (1.0 + math.exp(-e / k))
            err += (r - p) ** 2
            n += 1
        return err / max(n, 1)

    lo, hi = 100.0, 2000.0
    for _ in range(30):
        m1 = lo + (hi - lo) / 3
        m2 = hi - (hi - lo) / 3
        if mse(m1) < mse(m2):
            hi = m2
        else:
            lo = m1
    return (lo + hi) / 2


def compute_error(engine, sample, k):
    evals = eval_white_pov_batch(engine, [fen for fen, _ in sample])
    err = 0.0
    n = 0
    for e, (_, r) in zip(evals, sample):
        if e is None:
            continue
        p = 1.0 / (1.0 + math.exp(-e / k))
        err += (r - p) ** 2
        n += 1
    return err / max(n, 1)


# ---------------------------------------------------------------------------
# 4. Symmetry-aware parameter reduction
# ---------------------------------------------------------------------------
#
# Many of g_EvalDNA's arrays are large (128/256-cell) mirrored
# location tables built from a handful of distinct values (e.g.
# PAWN_CENTRALITY_BONUS is 128 cells but only {0, -8, 5, 9} actually
# appear). Tuning every raw cell independently wastes compute 1:1 with
# how much it inflates the parameter count, AND it can overfit into
# board-asymmetric nonsense a batch of a few thousand positions has no
# business justifying. So: tune the distinct values per array, not
# the raw cells, and expand back to the full array before it's ever
# written to a .dna file or loaded into the engine.

def build_value_groups(flat_row):
    """Map a single g_EvalDNA array's cells to their distinct values.
    Returns (distinct_values, cell_to_group) where cell_to_group[i]
    is the index into distinct_values that cell i's original value
    came from."""
    value_to_id = {}
    cell_to_group = []
    for v in flat_row:
        if v not in value_to_id:
            value_to_id[v] = len(value_to_id)
        cell_to_group.append(value_to_id[v])
    distinct_values = list(value_to_id.keys())
    return distinct_values, cell_to_group


def expand_group(distinct_values, cell_to_group):
    """Inverse of build_value_groups: rebuild the full-length array
    from its distinct values, preserving whatever symmetry the
    original array had."""
    return [distinct_values[g] for g in cell_to_group]


def build_all_groups(flat, row_lengths):
    """Reduce a full g_EvalDNA flat vector to its per-array distinct
    values. Returns (param_vec, groups) where param_vec is the
    reduced tunable-parameter vector and groups holds, per array,
    (start index into param_vec, count of distinct values, cell_to_group)
    -- everything expand_all_groups() needs to reconstruct the full
    flat vector later."""
    i = 0
    param_vec = []
    groups = []
    for n in row_lengths:
        row = flat[i:i + n]
        i += n
        distinct_values, cell_to_group = build_value_groups(row)
        start = len(param_vec)
        param_vec.extend(distinct_values)
        groups.append((start, len(distinct_values), cell_to_group))
    return param_vec, groups


def build_frozen_mask(groups, frozen_arrays):
    """True/False per param_vec index: True means "do not perturb this
    parameter during coordinate descent." `frozen_arrays` is a set of
    g_EvalDNA array names (must match DNA_NAMES, one entry per group/
    array, same order as row_lengths) -- e.g. names dna_trend.py
    flagged as flip-flopping noise rather than a real trend. Raises on
    an unrecognized name so a typo doesn't silently freeze nothing."""
    unknown = frozen_arrays - set(DNA_NAMES)
    if unknown:
        raise ValueError(f"Unknown array name(s) in --freeze: {sorted(unknown)}")
    mask = [False] * sum(count for _, count, _ in groups)
    for name, (start, count, _cell_to_group) in zip(DNA_NAMES, groups):
        if name in frozen_arrays:
            for j in range(start, start + count):
                mask[j] = True
    return mask


def expand_all_groups(param_vec, groups):
    """Inverse of build_all_groups: rebuild the full flat g_EvalDNA
    vector (in the same per-array layout evaldna read/dump expects)
    from a reduced param_vec."""
    flat = []
    for start, count, cell_to_group in groups:
        distinct_values = param_vec[start:start + count]
        flat.extend(expand_group(distinct_values, cell_to_group))
    return flat


# ---------------------------------------------------------------------------
# 5. Local coordinate-descent tuner (the actual "Texel tuning" step)
# ---------------------------------------------------------------------------

def write_dna_file(path, flat, row_lengths):
    i = 0
    rows = []
    for n in row_lengths:
        rows.append(",".join(str(v) for v in flat[i:i + n]))
        i += n
    Path(path).write_text("\n".join(rows))


def tune(engine_path, pgn_path, dna_scratch_dir, out_path,
         batch_size=None, holdout_frac=0.1, full_eval_every=1, max_passes=6,
         n_workers=2, max_positions=200_000, min_elo=2400, frozen_arrays=None,
         game_stride=1, force_rescan=False, qsearch_max_diff=None):
    """batch_size=None (the default) sizes each pass's training batch as
    1/10th of the training pool (post-holdout-split) instead of a fixed
    count -- a fixed 2000 was fine at the 25k-position scale this was
    first tuned against, but silently stayed thin (~4 positions per
    tunable parameter) once runs started loading 600k positions,
    since nothing scaled the batch with the pool. Pass an explicit int
    to override."""
    import time as _time
    run_start = _time.time()

    engine = EnginePool(engine_path, dna_scratch_dir, n_workers=n_workers)
    try:
        _tune_body(engine, pgn_path, out_path, batch_size, holdout_frac,
                   full_eval_every, max_passes, max_positions, min_elo,
                   run_start, frozen_arrays=frozen_arrays,
                   game_stride=game_stride, force_rescan=force_rescan,
                   qsearch_max_diff=qsearch_max_diff)
    finally:
        # Graceful quit() is what actually releases each engine's
        # SysV semaphore (input.c) -- a bare kill() does not, since it
        # never runs the engine's own cleanup path. This is the fix
        # for the leak that cost hours of this session: always attempt
        # a real quit() here, even when _tune_body raised.
        engine.quit()


def _tune_body(engine, pgn_path, out_path, batch_size, holdout_frac,
               full_eval_every, max_passes, max_positions, min_elo, run_start,
               frozen_arrays=None, game_stride=1, force_rescan=False,
               qsearch_max_diff=None):
    import time as _time
    baseline_flat, row_lengths = engine.dump_dna()

    param_vec, groups = build_all_groups(baseline_flat, row_lengths)
    print(f"Reduced {len(baseline_flat)} raw DNA cells to "
          f"{len(param_vec)} distinct tunable parameters "
          f"({len(param_vec) / len(baseline_flat):.1%} of the original).")

    frozen_arrays = frozen_arrays or set()
    frozen_mask = build_frozen_mask(groups, frozen_arrays)
    n_frozen = sum(frozen_mask)
    if n_frozen:
        print(f"Freezing {n_frozen}/{len(param_vec)} parameters in "
              f"{len(frozen_arrays)} array(s): {sorted(frozen_arrays)}")

    print(f"Loading up to {max_positions} quiet positions (min_elo={min_elo}, "
          f"game_stride={game_stride})...")
    t0 = _time.time()
    all_positions = load_training_positions(
        pgn_path, max_positions=max_positions, min_elo=min_elo,
        game_stride=game_stride, force_rescan=force_rescan
    )

    if qsearch_max_diff is not None:
        print(f"Filtering by eval/qsearch agreement (max_diff={qsearch_max_diff}cp)...")
        t_qs = _time.time()
        all_positions = filter_by_qsearch_agreement(
            engine, all_positions, max_diff_cp=qsearch_max_diff
        )
        print(f"  filtered in {_time.time()-t_qs:.1f}s")

    random.shuffle(all_positions)
    print(f"  got {len(all_positions)} positions in {_time.time()-t0:.1f}s")

    # Hold back a slice that never enters a training batch, so pass-by-
    # pass error on it actually measures generalization instead of just
    # re-checking against (a superset of) the same positions the batch
    # was drawn from. Split before any sampling below touches the pool.
    n_holdout = int(len(all_positions) * holdout_frac)
    holdout_positions = all_positions[:n_holdout]
    train_positions = all_positions[n_holdout:]
    print(f"  holding back {len(holdout_positions)} positions for "
          f"validation, training on {len(train_positions)}")

    if batch_size is None:
        batch_size = max(1, len(train_positions) // 10)
        print(f"  batch_size not specified, defaulting to 1/10th of "
              f"the training pool: {batch_size}")

    calib_sample = train_positions[:2000]

    k = fit_k(engine, calib_sample)
    print(f"Calibrated K = {k:.1f}")

    def sample_batch():
        return random.sample(train_positions,
                              min(batch_size, len(train_positions)))

    def load(vec):
        engine.load_dna(expand_all_groups(vec, groups), row_lengths)

    load(param_vec)
    best_err = compute_error(engine, sample_batch(), k)
    print(f"Baseline error: {best_err:.5f}")

    # Track the best holdout-err checkpoint separately from out_path,
    # which every pass unconditionally overwrites with whatever that
    # pass ended up at -- observed live that a run can improve for
    # several passes and then regress (batch noise, not necessarily
    # true overfitting), silently clobbering the best DNA seen so far
    # with a worse one and losing it for good since there was no copy
    # anywhere else. best_holdout_path gets rewritten only when a pass
    # actually beats the best holdout err seen so far, so it survives
    # a late-run regression intact.
    best_holdout_err = float("inf")
    best_holdout_path = str(Path(out_path).with_suffix("")) + ".best_holdout.dna"
    if holdout_positions:
        best_holdout_err = compute_error(engine, holdout_positions, k)
        print(f"Baseline holdout err={best_holdout_err:.5f}")
        write_dna_file(best_holdout_path, expand_all_groups(param_vec, groups),
                        row_lengths)

    # Checkpoint immediately so out_path always reflects a valid,
    # loadable DNA file from the very first moment, not just at the
    # (possibly hours-away) end of the run.
    write_dna_file(out_path, expand_all_groups(param_vec, groups), row_lengths)

    step = 8  # start coarse, like Texel's original implementation
    for pass_no in range(max_passes):
        pass_start = _time.time()
        # CORRECTNESS FIX: hold ONE fixed batch for the entire pass
        # instead of resampling per trial. Verified live that this
        # mattered: probing the "converged" checkpoint from an earlier
        # run (which resampled per-trial) found real, often-large
        # improvements at ~25% of tested parameters that the run
        # itself never found -- not because there was nothing left to
        # improve, but because comparing each trial's error against a
        # best_err computed on a *different* random batch is an
        # apples-to-oranges comparison that resampling noise can
        # easily dominate at this batch size. Recomputing best_err
        # fresh against the same fixed batch at the top of each pass
        # keeps every comparison within a pass consistent; only
        # between passes does the batch (and therefore some
        # regularizing variety) change.
        pass_batch = sample_batch()
        load(param_vec)
        best_err = compute_error(engine, pass_batch, k)
        improved_this_pass = False
        for i in range(len(param_vec)):
            if frozen_mask[i]:
                continue
            improved = False
            for delta in (step, -step):
                trial = list(param_vec)
                trial[i] += delta
                load(trial)
                err = compute_error(engine, pass_batch, k)
                if err < best_err:
                    param_vec = trial
                    best_err = err
                    improved = True
                    improved_this_pass = True
                    break
            # Heartbeat every 25 params regardless of whether this one
            # improved -- the old `if improved and ...` gate went silent
            # for the whole rest of a pass once improvements dried up,
            # which reads as "stuck" on a multi-hour run even though
            # trials are still being tried at full speed.
            if i % 25 == 0:
                elapsed = _time.time() - run_start
                tag = "improved" if improved else "no change"
                print(f"pass {pass_no} param {i}/{len(param_vec)}: "
                      f"err={best_err:.5f} step={step} ({tag}) "
                      f"(elapsed {elapsed/3600:.2f}h)")

        # Checkpoint after every pass -- a pass can run for hours, and
        # this is the only point where losing partial progress to a
        # crash/kill would be expensive. Cheap relative to a pass.
        write_dna_file(out_path, expand_all_groups(param_vec, groups), row_lengths)
        print(f"pass {pass_no} done in {(_time.time()-pass_start)/3600:.2f}h, "
              f"checkpoint written to {out_path}")

        if not improved_this_pass:
            step = max(1, step // 2)
            if step == 1 and not improved_this_pass:
                break

        # Every pass: score against holdout_positions, which never
        # appears in any training batch. This is the actual
        # overfitting check -- a pass whose training error drops while
        # holdout error rises or flatlines is chasing batch noise, not
        # real signal.
        if holdout_positions:
            load(param_vec)
            holdout_err = compute_error(engine, holdout_positions, k)
            print(f"pass {pass_no}: holdout err={holdout_err:.5f}")
            if holdout_err < best_holdout_err:
                best_holdout_err = holdout_err
                write_dna_file(best_holdout_path,
                                expand_all_groups(param_vec, groups), row_lengths)
                print(f"pass {pass_no}: new best holdout err "
                      f"({best_holdout_err:.5f}), checkpoint written to "
                      f"{best_holdout_path}")

        # Periodically re-score the whole training pool (not just the
        # per-pass batch) as a secondary sanity check. full_eval_every
        # now defaults to 1 (every pass) specifically so this is
        # directly comparable pass-over-pass against holdout_err above
        # -- unlike the per-param err= lines (measured against a fresh
        # random pass_batch every pass), this uses a fixed slice of
        # train_positions, so its trend is the real apples-to-apples
        # "is training pulling away from holdout" signal.
        if pass_no % full_eval_every == 0:
            load(param_vec)
            full_err = compute_error(engine, train_positions[:20000], k)
            print(f"pass {pass_no}: full-train-sample err={full_err:.5f}")

    final_flat = expand_all_groups(param_vec, groups)
    engine.load_dna(final_flat, row_lengths)
    write_dna_file(out_path, final_flat, row_lengths)
    print(f"Wrote final tuned DNA to {out_path} "
          f"(total runtime {(_time.time()-run_start)/3600:.2f}h)")
    if holdout_positions:
        print(f"Best holdout err seen this run: {best_holdout_err:.5f} "
              f"(checkpoint: {best_holdout_path}) -- may differ from "
              f"{out_path}, which is whatever the final pass ended at, "
              f"not necessarily the best pass.")


if __name__ == "__main__":
    import sys
    # Optional trailing overrides for quick/dry runs:
    #   tune_eval_dna.py <engine> <pgn> [max_positions] [max_passes] \
    #       [batch_size] [holdout_frac] [n_workers] [--freeze A,B,C] \
    #       [--game-stride N] [--no-cache] [--qsearch-max-diff N] \
    #       [--no-pgn-elo-filter]
    # batch_size defaults to None (1/10th of the training pool, sized
    # after max_positions/holdout_frac are applied -- see tune()'s
    # docstring); pass 0 here to mean "use the default" if you want to
    # set holdout_frac/n_workers without pinning batch_size.
    #
    # --freeze takes a comma-separated list of g_EvalDNA array names
    # (exact spelling from dna_diff.py's DNA_NAMES, e.g. as printed by
    # dna_trend.py) whose parameters are excluded from coordinate
    # descent entirely -- they're dumped/loaded/checkpointed like
    # every other cell, just never perturbed. Use this to keep budget
    # off arrays dna_trend.py flagged as flip-flopping noise rather
    # than a real trend, without hand-editing the .dna file.
    #
    # --game-stride N keeps only 1 game in N (after the Elo filter)
    # while scanning the PGN, so a max_positions cap spans the whole
    # chronological corpus instead of only its earliest games -- see
    # load_training_positions()'s docstring. --no-cache forces a fresh
    # PGN scan even if a position_cache/ entry already matches these
    # exact args (e.g. to rule out a bad cached entry).
    #
    # --qsearch-max-diff N adds a second, stricter quiet-position gate
    # on top of the SEE-based one: after loading (and before the
    # train/holdout split), drop any position where the engine's own
    # static eval and QSearch() score disagree by more than N
    # centipawns -- see filter_by_qsearch_agreement's docstring. Off
    # by default since it costs one eval + one qsearch round trip per
    # already-loaded position.
    args = sys.argv[1:]
    game_stride = 1
    if "--game-stride" in args:
        idx = args.index("--game-stride")
        game_stride = int(args[idx + 1])
        del args[idx:idx + 2]
    force_rescan = "--no-cache" in args
    if force_rescan:
        args.remove("--no-cache")
    qsearch_max_diff = None
    if "--qsearch-max-diff" in args:
        idx = args.index("--qsearch-max-diff")
        qsearch_max_diff = int(args[idx + 1])
        del args[idx:idx + 2]
    # --no-pgn-elo-filter: some position sources (e.g. match_play.py's
    # match_games.pgn, baseline-vs-candidate self-play) have no
    # WhiteElo/BlackElo headers at all -- game.headers.get(...,"0")
    # then defaults to 0, which the default min_elo=2400 filters out
    # entirely, silently yielding zero positions. This forces
    # min_elo=0 so those games aren't dropped.
    no_pgn_elo_filter = "--no-pgn-elo-filter" in args
    if no_pgn_elo_filter:
        args.remove("--no-pgn-elo-filter")

    frozen_arrays = set()
    for flag in ("--freeze", "--freeze="):
        for a in list(args):
            if a == "--freeze":
                idx = args.index(a)
                frozen_arrays |= {n.strip() for n in args[idx + 1].split(",") if n.strip()}
                del args[idx:idx + 2]
                break
            if a.startswith("--freeze="):
                frozen_arrays |= {n.strip() for n in a[len("--freeze="):].split(",") if n.strip()}
                args.remove(a)
                break

    kwargs = {}
    if frozen_arrays:
        kwargs["frozen_arrays"] = frozen_arrays
    if game_stride != 1:
        kwargs["game_stride"] = game_stride
    if force_rescan:
        kwargs["force_rescan"] = force_rescan
    if qsearch_max_diff is not None:
        kwargs["qsearch_max_diff"] = qsearch_max_diff
    if no_pgn_elo_filter:
        kwargs["min_elo"] = 0
    if len(args) > 2:
        kwargs["max_positions"] = int(args[2])
    if len(args) > 3:
        kwargs["max_passes"] = int(args[3])
    if len(args) > 4 and args[4] and int(args[4]) != 0:
        kwargs["batch_size"] = int(args[4])
    if len(args) > 5 and args[5]:
        kwargs["holdout_frac"] = float(args[5])
    if len(args) > 6 and args[6]:
        kwargs["n_workers"] = int(args[6])
    tune(
        engine_path=args[0],
        pgn_path=args[1],
        dna_scratch_dir="/tmp/typhoon_tune_dna",
        out_path="tuned.dna",
        **kwargs,
    )