#!/usr/bin/env python3
"""Build a subsampled test set for cheap in-training monitoring evals.

The full Sudoku-Extreme test set is 422,786 puzzles = 551 batches of 768, and each
eval batch runs the full ACT loop (16 steps x 21 network applications), so a full
eval costs ~3h on halfmind. A random subsample of ~12k puzzles gives +/-0.3%
standard error at p~0.87 — ample for tracking progress — at ~5 min per eval.

The final reported number still comes from a full-test-set eval via
scripts/eval_checkpoint.py.

Usage:
  python scripts/make_test_subset.py --src data/sudoku-extreme-1k-aug-1000 \
      --dst data/sudoku-extreme-testsub-12k --n 12288 --seed 0
"""

import argparse
import json
import os

import numpy as np


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", required=True, help="dataset dir containing test/")
    ap.add_argument("--dst", required=True)
    ap.add_argument("--n", type=int, default=12288, help="puzzles to keep (multiple of batch size is tidiest)")
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()

    src_test = os.path.join(args.src, "test")
    dst_test = os.path.join(args.dst, "test")
    os.makedirs(dst_test, exist_ok=True)

    inputs = np.load(os.path.join(src_test, "all__inputs.npy"), mmap_mode="r")
    labels = np.load(os.path.join(src_test, "all__labels.npy"), mmap_mode="r")
    total = inputs.shape[0]
    n = min(args.n, total)

    rng = np.random.default_rng(args.seed)
    idx = np.sort(rng.choice(total, size=n, replace=False))

    np.save(os.path.join(dst_test, "all__inputs.npy"), np.asarray(inputs[idx]))
    np.save(os.path.join(dst_test, "all__labels.npy"), np.asarray(labels[idx]))
    # one example per puzzle, one puzzle per group (matches how the builder emits Sudoku)
    np.save(os.path.join(dst_test, "all__puzzle_identifiers.npy"), np.zeros(n, dtype=np.int32))
    np.save(os.path.join(dst_test, "all__puzzle_indices.npy"), np.arange(n + 1, dtype=np.int32))
    np.save(os.path.join(dst_test, "all__group_indices.npy"), np.arange(n + 1, dtype=np.int32))

    with open(os.path.join(src_test, "dataset.json")) as f:
        meta = json.load(f)
    meta.update(total_groups=n, total_puzzles=n, mean_puzzle_examples=1)
    with open(os.path.join(dst_test, "dataset.json"), "w") as f:
        json.dump(meta, f)

    with open(os.path.join(args.dst, "subset_provenance.json"), "w") as f:
        json.dump({"src": args.src, "n": n, "of_total": int(total), "seed": args.seed}, f, indent=2)

    print(f"wrote {n} of {total} test puzzles to {dst_test} (seed {args.seed})")


if __name__ == "__main__":
    main()