#!/usr/bin/env python3
"""Stratified game-level train/holdout split over a games.jsonl
(build_games.py output).

Splits by GAME, not by individual position -- every position from one
game lands on the same side of the split, avoiding the leakage a flat
position-level shuffle would have (adjacent plies / recurring
middlegame structures from the same game are highly correlated, not
independent samples; see eval_tune/tune_eval_dna.py's own holdout
split, which doesn't do this).

Stratified on two axes so the holdout's composition mirrors the full
pool's rather than a flat 10% sample risking (by chance) too many
decisive games, too few draws, or skewing toward short/long games:
  - result: White win / Black win / Draw
  - length: quartile bin of the game's annotated ply count (quartiles
    of the actual corpus, not guessed fixed cutoffs)
10% (configurable) is sampled independently within each of the
resulting (result, length-quartile) strata, not 10% of the pool
overall.

Usage:
    python3 split_holdout.py --in runs/sd10_twic_2M/games.jsonl \
        --train-out runs/sd10_twic_2M/games.train.jsonl \
        --holdout-out runs/sd10_twic_2M/games.holdout.jsonl \
        --holdout-frac 0.1 --seed 0
"""
import argparse
import json
import math
import random
from collections import defaultdict


def result_bucket(result):
    if result == "1-0":
        return "white_win"
    if result == "0-1":
        return "black_win"
    if result == "1/2-1/2":
        return "draw"
    return "unknown"


def quartile_bin(length, cuts):
    for i, c in enumerate(cuts):
        if length <= c:
            return i
    return len(cuts)


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--in", dest="infile", required=True)
    ap.add_argument("--train-out", required=True)
    ap.add_argument("--holdout-out", required=True)
    ap.add_argument("--holdout-frac", type=float, default=0.1)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()

    games = []
    with open(args.infile) as f:
        for line in f:
            games.append(json.loads(line))
    print(f"Loaded {len(games)} games from {args.infile}")

    lengths = sorted(len(g["positions"]) for g in games)
    def pct(p):
        idx = min(int(p * len(lengths)), len(lengths) - 1)
        return lengths[idx]
    cuts = [pct(0.25), pct(0.50), pct(0.75)]
    print(f"Length quartile cutpoints (annotated plies): {cuts}")

    strata = defaultdict(list)
    for g in games:
        key = (result_bucket(g["result"]),
               quartile_bin(len(g["positions"]), cuts))
        strata[key].append(g)

    rng = random.Random(args.seed)
    train_games = []
    holdout_games = []
    print(f"\n{'stratum':<24} {'total':>7} {'holdout':>8} {'train':>8}")
    for key in sorted(strata.keys()):
        bucket = strata[key]
        rng.shuffle(bucket)
        n_holdout = math.ceil(len(bucket) * args.holdout_frac)
        holdout_games.extend(bucket[:n_holdout])
        train_games.extend(bucket[n_holdout:])
        label = f"{key[0]}/q{key[1]}"
        print(f"{label:<24} {len(bucket):>7} {n_holdout:>8} "
              f"{len(bucket) - n_holdout:>8}")

    print(f"\n{'TOTAL':<24} {len(games):>7} {len(holdout_games):>8} "
          f"{len(train_games):>8}")

    with open(args.train_out, "w") as f:
        for g in train_games:
            f.write(json.dumps(g) + "\n")
    with open(args.holdout_out, "w") as f:
        for g in holdout_games:
            f.write(json.dumps(g) + "\n")

    n_train_pos = sum(len(g["positions"]) for g in train_games)
    n_holdout_pos = sum(len(g["positions"]) for g in holdout_games)
    print(f"\nWrote {args.train_out}: {len(train_games)} games, "
          f"{n_train_pos} positions")
    print(f"Wrote {args.holdout_out}: {len(holdout_games)} games, "
          f"{n_holdout_pos} positions")


if __name__ == "__main__":
    main()