#!/usr/bin/env python3
"""Fetch seeded reasoning-eval samples -> data/*.jsonl. Network only (no model server)."""
import json, random, os, sys
from datasets import load_dataset

OUT = os.path.join(os.path.dirname(__file__), "data")
os.makedirs(OUT, exist_ok=True)
SEED = 1234

# (name, n) — directional defaults; override via argv: prep_data.py gsm8k=100 math=50 ...
DEFAULTS = {"gsm8k": 30, "math": 20, "bbh": 20, "gpqa": 15}
overrides = dict(kv.split("=") for kv in sys.argv[1:] if "=" in kv)
N = {k: int(overrides.get(k, v)) for k, v in DEFAULTS.items()}

def save(name, rows):
    p = os.path.join(OUT, f"{name}.jsonl")
    with open(p, "w") as f:
        for r in rows: f.write(json.dumps(r) + "\n")
    print(f"  saved {len(rows):>4} -> {p}")

def sample(ds, n):
    idx = list(range(len(ds)))
    random.Random(SEED).shuffle(idx)
    return [ds[i] for i in idx[:n]]

print("Fetching eval samples (seed=%d)..." % SEED)

# GSM8K — answer after '#### '
try:
    ds = load_dataset("openai/gsm8k", "main", split="test")
    rows = [{"id": f"gsm8k-{i}", "category": "gsm8k",
             "question": r["question"],
             "answer": r["answer"].split("####")[-1].strip().replace(",", "")}
            for i, r in enumerate(sample(ds, N["gsm8k"]))]
    save("gsm8k", rows)
except Exception as e:
    print(f"  GSM8K FAILED: {e}")

# MATH-500 — clean ungated test set with explicit 'answer' field
try:
    ds = load_dataset("HuggingFaceH4/MATH-500", split="test")
    rows = [{"id": f"math-{i}", "category": "math",
             "question": r["problem"], "answer": str(r["answer"]).strip(),
             "level": r.get("level"), "type": r.get("subject")}
            for i, r in enumerate(sample(ds, N["math"]))]
    save("math", rows)
except Exception as e:
    print(f"  MATH FAILED: {e}")

# BBH — use a multi-step subset (causal_judgement is y/n; use 'logical_deduction_three_objects' MC)
try:
    ds = load_dataset("lukaemon/bbh", "logical_deduction_three_objects", split="test")
    rows = [{"id": f"bbh-{i}", "category": "bbh",
             "question": r["input"], "answer": r["target"].strip("()").strip()}
            for i, r in enumerate(sample(ds, N["bbh"]))]
    save("bbh", rows)
except Exception as e:
    print(f"  BBH FAILED: {e} (may need a different config name)")

# MMLU-Pro — ungated, 10-option hard knowledge+reasoning (harder than MMLU/GPQA-easy)
try:
    ds = load_dataset("TIGER-Lab/MMLU-Pro", split="test")
    # focus on reasoning-heavy categories
    hard = ds.filter(lambda r: r["category"] in
                     {"physics", "engineering", "math", "chemistry", "law"})
    rows = []
    for i, r in enumerate(sample(hard, N["gpqa"])):
        labels = "ABCDEFGHIJ"
        q = r["question"] + "\n" + "\n".join(
            f"{labels[k]}) {opt}" for k, opt in enumerate(r["options"]))
        rows.append({"id": f"mmlupro-{i}", "category": "mmlu_pro",
                     "question": q, "answer": r["answer"], "subject": r["category"]})
    save("mmlu_pro", rows)
except Exception as e:
    print(f"  MMLU-Pro SKIPPED: {e}")

print("Done. Files in", OUT)