#!/usr/bin/env python3
"""Mine middlegame checkmate positions + matched controls for king-safety
defect regression. See CLAUDE.md / conversation context: _EvalKing's
counted "defects" are summed into one scalar and run through one lookup
table with no per-defect weighting; this pulls training data to fit
relative weights instead of hand-picking them.

Output: a CSV of (fen, label, mated_color, n_back, game_id) rows. FEN
extraction only -- no engine call here. The actual per-defect feature
values get pulled from these FENs by the engine-side dump command
(separate step).
"""
import argparse
import chess
import chess.pgn
import csv
import random
import sys

# has_hanging_material (Python SEE, same swap-list algorithm already
# validated against 800k TWIC positions) -- reused rather than
# reimplemented, so "quiet" means the identical thing here as it does
# for the generic TWIC pool. 2026-09-15: this project's original
# mine_king_mates.py design deliberately did NOT filter for quietness
# (a position N plies before mate is *supposed* to capture real,
# unresolved danger, including in-check/mid-tactic cases) -- fine for
# the mate-mining regression, which never fed FENs to a static-eval-vs-
# outcome fit, but wrong for Texel tuning, which needs every training
# position's static eval to be meaningful on its own. Counters below
# report the actual quiet-survival rate directly -- expected to be much
# lower than TWIC's ~93-94% given near-mate positions are inherently
# more likely to be mid-tactic by construction.
sys.path.insert(0, ".")
from tune_eval_dna import has_hanging_material

g_uPositiveCandidates = 0
g_uPositiveQuiet = 0
g_uControlCandidates = 0
g_uControlQuiet = 0


def is_quiet_position(board):
    return (not board.is_check()) and (not has_hanging_material(board))

# Matches eval.c's DO_KING_SAFETY_THRESHOLD minus the VALUE_KING baseline
# _EvalKing adds internally: real threshold is "enemy has at least a
# rook+bishop of non-pawn, non-king material" before king safety is even
# scored at all. No point mining data _EvalKing would ignore.
MATING_MATERIAL_THRESHOLD_CP = 800

PIECE_VALUES_CP = {
    chess.PAWN: 0,  # excluded: "non-pawn material" by definition
    chess.KNIGHT: 300,
    chess.BISHOP: 300,
    chess.ROOK: 500,
    chess.QUEEN: 975,
    chess.KING: 0,  # excluded: baseline, not a mating-material contributor
}


def non_pawn_material_cp(board, color):
    total = 0
    for piece_type in (chess.KNIGHT, chess.BISHOP, chess.ROOK, chess.QUEEN):
        total += len(board.pieces(piece_type, color)) * PIECE_VALUES_CP[piece_type]
    return total


# Mirrors _EvalKing's own middlegame/endgame boundary exactly (eval.c:
# "u = pos->uNonPawnCount[WHITE][0] + pos->uNonPawnCount[BLACK][0]; if (u <
# 8)") -- total non-pawn pieces, including both kings, across both sides.
# Below this, _EvalKing itself stops caring about ring-square/file-storm
# defects and switches to endgame-only terms (centralize king, escort
# passer). Checking the attacker's material alone (non_pawn_material_cp)
# is NOT sufficient on its own: a lone queen mating a bare king trivially
# clears that floor but is a basic mating technique, not a middlegame
# king-safety collapse -- caught by manually sampling FENs from the mined
# set (several K+Q-vs-lone-K endings slipped through material-only
# filtering).
def total_non_pawn_piece_count(board):
    count = 2  # both kings
    for color in (chess.WHITE, chess.BLACK):
        for piece_type in (chess.KNIGHT, chess.BISHOP, chess.ROOK, chess.QUEEN):
            count += len(board.pieces(piece_type, color))
    return count


MIDDLEGAME_MIN_TOTAL_NON_PAWN_PIECES = 8

# 2026-09-15, by direct instruction: a promotion within the last N plies
# of the game means this "mate" is really an endgame pawn-race/technique
# conversion (queen a pawn, then mate the resulting weak/bare king) --
# not a middlegame king-safety collapse. We want positions where one
# side mounts a real attack on a king still attempting fortification/
# defense, not a K+P-and-a-few-pieces endgame decided by promotion.
# Matches the n_backs cap (14) -- if we'd sample as far back as 14 plies
# before mate, a promotion inside that same trailing window taints every
# n_back for this game, not just the ones after the promotion.
LATE_PROMOTION_EXCLUSION_WINDOW = 14


def is_middlegame(board):
    return total_non_pawn_piece_count(board) >= MIDDLEGAME_MIN_TOTAL_NON_PAWN_PIECES


def mine_game(game, n_backs, min_ply_from_start=20,
              promotion_exclusion_window=LATE_PROMOTION_EXCLUSION_WINDOW):
    """Yields (fen, mated_color, n_back) for a game ending in checkmate."""
    board = game.board()
    moves = list(game.mainline_moves())
    if not moves:
        return
    # Replay to the end to confirm checkmate and find the mated side.
    for mv in moves:
        board.push(mv)
    if not board.is_checkmate():
        return
    mated_color = board.turn  # side to move at mate = side that got mated
    total_plies = len(moves)

    # Cheap check -- Move objects already carry .promotion, no replay
    # needed. Skips the WHOLE game's positives, not just n_backs past
    # the promotion, since a late-game promotion taints the game's
    # suitability as a middlegame king-safety example regardless of
    # which n_back we're sampling.
    if any(mv.promotion is not None
           for mv in moves[max(0, total_plies - promotion_exclusion_window):]):
        return

    for n_back in n_backs:
        idx = total_plies - n_back
        if idx < min_ply_from_start:
            continue
        replay = game.board()
        for mv in moves[:idx]:
            replay.push(mv)
        # Only meaningful if the mated side's king safety was even "in
        # scope" for _EvalKing at this point -- attacker must have
        # mating material.
        attacker = not mated_color
        if non_pawn_material_cp(replay, attacker) < MATING_MATERIAL_THRESHOLD_CP:
            continue
        if not is_middlegame(replay):
            continue
        global g_uPositiveCandidates, g_uPositiveQuiet
        g_uPositiveCandidates += 1
        if not is_quiet_position(replay):
            continue
        g_uPositiveQuiet += 1
        yield replay.fen(), mated_color, n_back


def is_draw(game):
    return game.headers.get("Result", "") == "1/2-1/2"


def mine_controls(game, n_needed, rng, min_ply=20, max_ply_from_end=15):
    """Yields (fen, color) control samples from a non-mate game (or a mate
    game's earlier-game portion), gated on the same material threshold,
    kept well clear of the game's ending so a real collapse (if any) isn't
    accidentally captured as a "stable" control.

    Draw games should be called with a larger max_ply_from_end than
    decisive/resignation games (see is_draw callers in main()) -- a draw
    only vouches for a king position if the position stayed quiet for a
    real stretch *after* the sample point too, not just "the game
    eventually drew" (e.g. a middlegame attack that later got defused into
    a simplified drawn endgame doesn't vouch for the earlier position).

    min_ply must be large enough that development is realistically done
    (move ~10, ply 20) -- NOT just "not literally move 1". A large chunk
    of the pool's draws are short draws-by-agreement close to
    filter_pgn.py's own 40-ply floor; combined with a big max_ply_from_end
    and a too-small min_ply, the only plies satisfying BOTH constraints on
    a short game land in the early opening, where minorsathome/development
    features are legitimately still nonzero for completely ordinary
    reasons -- not a real king-safety defect. Confirmed as the actual bug
    behind an earlier version's controls showing minorsathome>=2 in ~85%
    of samples: the engine's own computation was verified correct on a
    real developed middlegame FEN, so the contamination was 100% a
    sampling-window problem, not an eval bug.

    A game too short to satisfy both min_ply and max_ply_from_end is
    skipped entirely rather than forced into an early sample -- fewer,
    honest control samples beat many contaminated ones."""
    moves = list(game.mainline_moves())
    total_plies = len(moves)
    end_guard = max_ply_from_end
    if total_plies < min_ply + end_guard:
        return
    candidates = list(range(min_ply, max(min_ply, total_plies - end_guard)))
    if not candidates:
        return
    # BUG #1 (found while investigating why controls skewed early-game):
    # shuffling `candidates` then iterating `sorted(candidates)` completely
    # undid the shuffle, so every game always visited its window in
    # ascending ply order and grabbed the *earliest* plies satisfying
    # per_game_budget -- systematically starving late-middlegame coverage
    # in any game with a wide candidate window, regardless of min_ply.
    #
    # BUG #2 (found immediately after "fixing" #1 with a random subsample
    # + sort): stopping the ascending-order loop as soon as `produced >=
    # n_needed` still biases toward whichever subsample members happen to
    # be smallest, since ascending iteration always finds the low end of
    # the subsample first. Measured: even after subsampling, sampled plies
    # still clustered at the 20th-40th percentile of game length instead
    # of spreading across the whole game.
    #
    # Fix for both: replay ascending through an oversampled random subset
    # ONCE (required -- python-chess boards are only walkable forward) to
    # collect every candidate result that passes the material/middlegame
    # filters, then pick which ones to actually keep via `rng.sample` on
    # that result LIST -- decoupling "must replay in order" from "which
    # results count" is what actually gives uniform ply coverage.
    num_to_try = min(len(candidates), max(n_needed * 4, 20))
    sample_indices = sorted(rng.sample(candidates, num_to_try))
    replay = game.board()
    ply = 0
    valid_results = []
    for idx in sample_indices:
        while ply < idx:
            replay.push(moves[ply])
            ply += 1
        if not is_middlegame(replay):
            continue
        global g_uControlCandidates, g_uControlQuiet
        g_uControlCandidates += 1
        if not is_quiet_position(replay):
            continue
        g_uControlQuiet += 1
        for color in (chess.WHITE, chess.BLACK):
            attacker = not color
            if non_pawn_material_cp(replay, attacker) >= MATING_MATERIAL_THRESHOLD_CP:
                valid_results.append((replay.fen(), color))

    if len(valid_results) > n_needed:
        valid_results = rng.sample(valid_results, n_needed)
    for fen, color in valid_results:
        yield fen, color


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("pgn_file")
    ap.add_argument("out_csv")
    ap.add_argument("--n-backs", type=int, nargs="+",
                     default=[4, 6, 8, 10, 12, 14],
                     help="Plies before mate to sample -- capped at 14 by "
                          "direct instruction (anything farther back is too "
                          "loosely connected to the actual mating attack).")
    ap.add_argument("--max-games", type=int, default=None,
                     help="Stop after scanning this many games (0/None = all)")
    ap.add_argument("--control-ratio", type=float, default=1.0,
                     help="Controls per positive, roughly")
    ap.add_argument("--draw-only-controls", action="store_true",
                     help="Don't fall back to decisive/resignation games "
                          "for controls even if draws don't fill the quota")
    ap.add_argument("--draw-end-guard", type=int, default=24,
                     help="Plies of quiet play required *after* a draw "
                          "control sample, stricter than decisive games "
                          "since a draw only vouches for a position that "
                          "stayed sound past the sample point too")
    ap.add_argument("--decisive-end-guard", type=int, default=15)
    ap.add_argument("--min-ply", type=int, default=20,
                     help="Earliest ply eligible for any sample (positive "
                          "or control) -- must be well past normal opening "
                          "development (ply 20 = move 10), not just >0. "
                          "See mine_controls' docstring for why this "
                          "matters a lot for short draws.")
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--byte-start", type=int, default=0,
                     help="Seek here before reading (shard support -- "
                          "must be a real game-start offset, e.g. from "
                          "shard_mine.py, not an arbitrary byte)")
    ap.add_argument("--byte-end", type=int, default=None,
                     help="Stop once past this offset (shard support)")
    args = ap.parse_args()

    rng = random.Random(args.seed)
    positives = []
    draw_pool_games = []
    decisive_pool_games = []

    with open(args.pgn_file, encoding="utf-8", errors="replace") as f:
        if args.byte_start:
            f.seek(args.byte_start)
        n_scanned = 0
        while True:
            if args.max_games and n_scanned >= args.max_games:
                break
            if args.byte_end is not None and f.tell() >= args.byte_end:
                break
            game = chess.pgn.read_game(f)
            if game is None:
                break
            n_scanned += 1
            if n_scanned % 20000 == 0:
                print(f"...scanned {n_scanned} games, "
                      f"{len(positives)} positives so far", file=sys.stderr)
            try:
                found_any = False
                for fen, mated_color, n_back in mine_game(
                        game, args.n_backs, min_ply_from_start=args.min_ply):
                    positives.append((fen, 1, chess.COLOR_NAMES[mated_color],
                                       n_back, n_scanned))
                    found_any = True
                if not found_any:
                    if is_draw(game):
                        draw_pool_games.append(game)
                    else:
                        decisive_pool_games.append(game)
            except Exception as e:
                print(f"skip game {n_scanned}: {e}", file=sys.stderr)
                continue

    print(f"Scanned {n_scanned} games -> {len(positives)} positive samples "
          f"from checkmate games. Control pools: {len(draw_pool_games)} "
          f"draws, {len(decisive_pool_games)} decisive/other.",
          file=sys.stderr)

    n_controls_needed = int(len(positives) * args.control_ratio)
    controls = []
    per_game_budget = 2

    # Draws first -- a position two comparable players couldn't crack is
    # stronger "this was sound" evidence than a resignation, where the
    # king position may simply never have been tested by that opponent.
    rng.shuffle(draw_pool_games)
    for game in draw_pool_games:
        if len(controls) >= n_controls_needed:
            break
        try:
            for fen, color in mine_controls(game, per_game_budget, rng,
                                             min_ply=args.min_ply,
                                             max_ply_from_end=args.draw_end_guard):
                controls.append((fen, 0, chess.COLOR_NAMES[color], -1, None, "draw"))
        except Exception:
            continue

    n_from_draws = len(controls)
    print(f"Collected {n_from_draws} control samples from draws.",
          file=sys.stderr)

    if not args.draw_only_controls and len(controls) < n_controls_needed:
        rng.shuffle(decisive_pool_games)
        for game in decisive_pool_games:
            if len(controls) >= n_controls_needed:
                break
            try:
                for fen, color in mine_controls(game, per_game_budget, rng,
                                                 min_ply=args.min_ply,
                                                 max_ply_from_end=args.decisive_end_guard):
                    controls.append((fen, 0, chess.COLOR_NAMES[color], -1, None, "decisive"))
            except Exception:
                continue
        print(f"Collected {len(controls) - n_from_draws} additional control "
              f"samples from decisive/resignation games (quota not met by "
              f"draws alone).", file=sys.stderr)

    print(f"Collected {len(controls)} control samples total.", file=sys.stderr)

    with open(args.out_csv, "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["fen", "label", "mated_color_or_sampled_color", "n_back",
                     "game_index", "control_source"])
        for row in positives:
            w.writerow(row + (None,))
        for row in controls:
            w.writerow(row)

    print(f"Wrote {len(positives) + len(controls)} rows to {args.out_csv}",
          file=sys.stderr)
    pos_rate = (g_uPositiveQuiet / g_uPositiveCandidates
                if g_uPositiveCandidates else 0.0)
    ctl_rate = (g_uControlQuiet / g_uControlCandidates
                if g_uControlCandidates else 0.0)
    print(f"Quiet-filter yield: positives {g_uPositiveQuiet}/"
          f"{g_uPositiveCandidates} ({pos_rate:.1%}), controls "
          f"{g_uControlQuiet}/{g_uControlCandidates} ({ctl_rate:.1%})",
          file=sys.stderr)


if __name__ == "__main__":
    main()