#!/usr/bin/env python3
"""Sample random (fen) positions from one or more PGN pools, at random
plies anywhere in the game (not filtered for quiet/tactical), for the
margin-safety recalibration described in board_representation/EVAL.md
section 9. Unlike tune_eval_dna.py's quiet-position pools, this
deliberately does NOT filter on SEE/hanging-material/qsearch-agreement --
Eval() runs at every node of the real search tree, quiet or not, so the
calibration sample should look like "all chess positions", not just the
quiet subset used for Texel tuning.

Uses the same random-byte-offset-seek trick as match_play.py's
sample_openings() to avoid a full linear scan of multi-GB pools.

Output: one FEN per source game (one random ply from openings-excluded
range), written as alternating "setboard <fen>" / "go" lines matching
tests/twic_sample.ep_'s format, suitable for a `script <file>` run
against a CALIBRATE_MARGIN_SAFETY build.
"""
import argparse
import hashlib
import io
import random
import sys
from pathlib import Path

import chess
import chess.pgn


def iter_random_game_texts(pgn_path, rng, max_attempts):
    """Yield raw PGN text blocks for randomly-located games in pgn_path,
    same seek-scan-extract approach as match_play.py's sample_openings."""
    st = Path(pgn_path).stat()
    file_size = st.st_size
    attempts = 0
    with open(pgn_path, "rb") as fb:
        while attempts < max_attempts:
            attempts += 1
            offset = rng.randint(0, max(file_size - 1, 0))
            fb.seek(offset)
            fb.readline()
            start = fb.tell()
            line = fb.readline()
            while line and not line.startswith(b"[Event "):
                start = fb.tell()
                line = fb.readline()
            if not line:
                continue
            fb.seek(start)
            fb.readline()
            end = fb.tell()
            line = fb.readline()
            while line and not line.startswith(b"[Event "):
                end = fb.tell()
                line = fb.readline()
            fb.seek(start)
            raw = fb.read(end - start).decode("utf-8", errors="replace")
            yield raw


PIECE_VALUES = {chess.KNIGHT: 3, chess.BISHOP: 3, chess.ROOK: 5, chess.QUEEN: 9}


def material_bucket(board):
    """Same bucketing as eval.c's _MaterialBucket: combined non-pawn
    army value (pawn units) / 8, capped at 7 (0 = bare kings, 7 = full
    army)."""
    total = 0
    for pt, v in PIECE_VALUES.items():
        total += v * len(board.pieces(pt, chess.WHITE))
        total += v * len(board.pieces(pt, chess.BLACK))
    return min(7, total // 8)


def sample_positions(pgn_paths, n, min_ply, seed, cache_dir, stratify=False):
    key_parts = []
    for p in pgn_paths:
        st = Path(p).stat()
        key_parts.append(f"{p}|{st.st_mtime}|{st.st_size}")
    key = hashlib.sha256(
        f"{'|'.join(key_parts)}|{n}|{min_ply}|{seed}|{stratify}".encode()
    ).hexdigest()[:16]
    cache_path = Path(cache_dir) / f"eval_calib_{key}.fens"
    if cache_path.exists():
        fens = cache_path.read_text().splitlines()
        print(f"  loaded {len(fens)} positions from cache {cache_path}",
              file=sys.stderr)
        return fens

    rng = random.Random(seed)
    fens = []
    # Stratified mode needs many more discarded candidates (most draws
    # land in the naturally-common middlegame buckets, not the rare
    # near-endgame ones we're specifically trying to fill), so give it
    # a much larger attempt budget than the unstratified path.
    max_attempts = (n * 200 + 2000) if stratify else (n * 30 + 500)
    attempts = 0
    # Round-robin across pools so a large pool (twic_filtered) doesn't
    # drown out a smaller one (any single lichess_elite month) purely by
    # byte-offset odds; weight by requested count via cycling.
    pool_iters = [iter_random_game_texts(p, rng, max_attempts) for p in pgn_paths]
    idx = 0

    # Unstratified: just take the first n candidates. Stratified: fill
    # 8 equal-size buckets (by combined material) up to quota each,
    # discarding candidates for buckets that are already full -- this
    # is what actually fixes cross-sectional coverage (uniform-random-
    # ply sampling is heavily skewed toward the middlegame/opening
    # buckets simply because games spend more plies there than grinding
    # out endgames; more positions alone doesn't fix that skew, only
    # explicit per-bucket quotas do).
    num_buckets = 8
    quota = -(-n // num_buckets) if stratify else 0  # ceil division
    bucket_counts = [0] * num_buckets

    while len(fens) < n and attempts < max_attempts:
        raw = next(pool_iters[idx % len(pool_iters)], None)
        idx += 1
        attempts += 1
        if raw is None:
            continue
        game = chess.pgn.read_game(io.StringIO(raw))
        if game is None:
            continue
        moves = list(game.mainline_moves())
        if len(moves) < min_ply + 1:
            continue
        # Pick one random ply strictly after min_ply and strictly before
        # the last move (avoid checkmate/stalemate terminal boards, which
        # `go` can't search meaningfully).
        target_ply = rng.randint(min_ply, len(moves) - 1)
        board = game.board()
        for i, mv in enumerate(moves):
            if i == target_ply:
                break
            board.push(mv)
        if board.is_game_over():
            continue
        if stratify:
            b = material_bucket(board)
            if bucket_counts[b] >= quota:
                continue
            bucket_counts[b] += 1
        fens.append(board.fen())
        if len(fens) % 2000 == 0:
            print(f"  sampled {len(fens)}/{n}...", file=sys.stderr)

    if stratify:
        print(f"  final bucket counts: {bucket_counts}", file=sys.stderr)

    Path(cache_dir).mkdir(parents=True, exist_ok=True)
    cache_path.write_text("\n".join(fens) + "\n")
    print(f"  wrote {len(fens)} positions to cache {cache_path}", file=sys.stderr)
    return fens


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("pgn_paths", nargs="+", help="one or more PGN pool files")
    ap.add_argument("--n", type=int, default=20000)
    ap.add_argument("--min-ply", type=int, default=4,
                     help="skip book-ish very early plies")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--cache-dir", default="eval_tune/position_cache")
    ap.add_argument("--out", required=True, help="output .ep_ file")
    ap.add_argument("--stratify", action="store_true",
                     help="fill 8 equal-size combined-material buckets "
                          "up to quota each, instead of taking positions "
                          "in natural (middlegame-skewed) draw order")
    args = ap.parse_args()

    fens = sample_positions(args.pgn_paths, args.n, args.min_ply, args.seed,
                             args.cache_dir, stratify=args.stratify)
    with open(args.out, "w") as f:
        for fen in fens:
            f.write(f"setboard {fen}\n")
            f.write("go\n")
    print(f"Wrote {len(fens)} positions ({2 * len(fens)} lines) to {args.out}")


if __name__ == "__main__":
    main()