"""Build a decontaminated, verifiable prompt pool for Beam -> Dory distillation.

Sources (train splits only):
  math    zwhe99/DeepMath-103K           difficulty >= 5, closed-form final answers
  science nvidia/OpenScienceReasoning-2  multiple choice, single-letter answers

Decontamination: drop any prompt sharing a 13-gram with MATH-500, AIME 2024/2025,
AMC 2023 or GPQA-style eval questions we can access without auth.
"""

import argparse
import hashlib
import json
import random
import re
from pathlib import Path

import duckdb

ROOT = Path(__file__).resolve().parent.parent
SOURCES = ROOT / "data" / "sources"

EVAL_SETS = [
    ("read_json_auto('hf://datasets/HuggingFaceH4/MATH-500/test.jsonl')", "problem"),
    ("'hf://datasets/HuggingFaceH4/aime_2024/**/*.parquet'", "problem"),
    ("'hf://datasets/yentinglin/aime_2025/**/*.parquet'", "problem"),
    ("'hf://datasets/math-ai/amc23/**/*.parquet'", "question"),
]

MCQ_PREFIX = (
    "Solve the following problem. Make sure to put the answer (and only answer) "
    "inside \\boxed{}.\n\n"
)
MATH_SUFFIX = "\n\nPut your final answer within \\boxed{}."


def norm_words(text):
    return re.sub(r"[^a-z0-9 ]+", " ", text.lower()).split()


def ngrams(text, n=13):
    w = norm_words(text)
    return {" ".join(w[i : i + n]) for i in range(max(0, len(w) - n + 1))}


def eval_ngrams(con):
    grams = set()
    for src, col in EVAL_SETS:
        repo = src.split("datasets/")[1].split("/", 2)[:2]
        repo = "/".join(repo)
        try:
            rows = con.execute(f"SELECT {col} FROM {src}").fetchall()
        except Exception as e:  # missing/gated eval set: report, keep going
            print(f"  decontam: skipped {repo}: {str(e)[:120]}")
            continue
        for (q,) in rows:
            grams |= ngrams(q or "")
        print(f"  decontam: {repo}: {len(rows)} questions")
    return grams


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--math", type=int, default=6000)
    ap.add_argument("--science", type=int, default=4000)
    ap.add_argument("--seed", type=int, default=20261022)
    ap.add_argument("--out", default=str(ROOT / "data" / "pool.jsonl"))
    args = ap.parse_args()

    con = duckdb.connect()
    con.execute("INSTALL httpfs; LOAD httpfs;")
    rng = random.Random(args.seed)

    print("loading eval sets for decontamination")
    bad = eval_ngrams(con)

    print("loading DeepMath-103K (difficulty >= 5)")
    math_rows = con.execute(
        """SELECT question, final_answer, difficulty, topic
           FROM 'hf://datasets/zwhe99/DeepMath-103K/**/*.parquet'
           WHERE difficulty >= 5"""
    ).fetchall()
    print(f"  {len(math_rows)} candidates")

    print("loading OpenScienceReasoning-2 (single-letter MCQ)")
    sci_rows = con.execute(
        """SELECT input, expected_answer
           FROM 'hf://datasets/nvidia/OpenScienceReasoning-2/**/*.parquet'
           WHERE length(expected_answer) = 1 AND expected_answer BETWEEN 'A' AND 'J'"""
    ).fetchall()
    print(f"  {len(sci_rows)} candidates")

    pool, seen, dropped = [], set(), 0

    def add(item):
        nonlocal dropped
        key = hashlib.sha1(" ".join(norm_words(item["prompt"])).encode()).hexdigest()
        if key in seen:
            return False
        if ngrams(item["prompt"]) & bad:
            dropped += 1
            return False
        seen.add(key)
        item["id"] = f"{item['source']}-{key[:12]}"
        pool.append(item)
        return True

    # Math: stratify by difficulty band so harder problems aren't swamped.
    bands = {}
    for q, a, d, t in math_rows:
        if not a or len(a) > 80:
            continue
        bands.setdefault(min(int(d), 9), []).append((q, a, d, t))
    per_band = args.math // len(bands)
    for band, rows in sorted(bands.items()):
        rng.shuffle(rows)
        taken = 0
        for q, a, d, t in rows:
            if taken >= per_band:
                break
            taken += add(
                {
                    "source": "deepmath",
                    "domain": "math",
                    "prompt": q.strip() + MATH_SUFFIX,
                    "answer": a,
                    "answer_type": "math",
                    "difficulty": d,
                    "topic": t,
                }
            )
        print(f"  math band {band}: {taken}")

    rng.shuffle(sci_rows)
    taken = 0
    for q, a in sci_rows:
        if taken >= args.science:
            break
        prompt = q if q.startswith(MCQ_PREFIX) else MCQ_PREFIX + q
        taken += add(
            {
                "source": "openscience2",
                "domain": "science",
                "prompt": prompt.strip(),
                "answer": a,
                "answer_type": "mcq",
            }
        )
    print(f"  science: {taken}")

    rng.shuffle(pool)
    with open(args.out, "w") as f:
        for item in pool:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")
    print(f"wrote {len(pool)} prompts to {args.out} (decontam dropped {dropped})")


if __name__ == "__main__":
    main()