# E20 data prep: MATH-500 items, MATH train items, and the train/test leakage check.
import json, re, random, sys
sys.path.insert(0, "."); from math_grade import last_boxed
from datasets import load_dataset
def norm(t): return re.sub(r"\s+", " ", t.lower()).strip()
def shingles(t, n=8):
    w = re.findall(r"\w+|[^\w\s]", t.lower()); return {" ".join(w[i:i+n]) for i in range(max(1, len(w)-n+1))}
te = load_dataset("HuggingFaceH4/MATH-500", split="test")
with open("data/math500.jsonl", "w", encoding="utf-8") as f:
    for k, ex in enumerate(te):
        f.write(json.dumps({"item": k, "unique_id": ex["unique_id"], "problem": ex["problem"], "gold": ex["answer"],
                            "level": ex["level"], "subject": ex["subject"]}) + "\n")
tr = []
for cfg in ["algebra", "counting_and_probability", "geometry", "intermediate_algebra", "number_theory", "prealgebra", "precalculus"]:
    for ex in load_dataset("EleutherAI/hendrycks_math", cfg, split="train"):
        g = last_boxed(ex["solution"])
        if g: tr.append({"problem": ex["problem"], "gold": g, "level": ex["level"], "subject": ex["type"], "solution": ex["solution"]})
print("train with boxed answer:", len(tr))
# leakage check: exact (whitespace/case normalized) and near-duplicate (8-gram shingle Jaccard >= 0.5)
tset = {norm(e["problem"]): i for i, e in enumerate(te)}
exact = [e["problem"][:80] for e in tr if norm(e["problem"]) in tset]
idx = {}
tsh = [shingles(e["problem"]) for e in te]
for i, s in enumerate(tsh):
    for g in s: idx.setdefault(g, []).append(i)
near = []; maxj = 0
for e in tr:
    s = shingles(e["problem"]); cnt = {}
    for g in s:
        for i in idx.get(g, []): cnt[i] = cnt.get(i, 0) + 1
    for i, c in cnt.items():
        j = c / len(s | tsh[i]); maxj = max(maxj, j)
        if j >= 0.5: near.append((e["problem"][:80], te[i]["problem"][:80], round(j, 3)))
rep = {"train_n": len(tr), "test_n": len(te), "exact_overlap": len(exact), "near_dup_pairs_j>=0.5": len(near),
       "max_jaccard_8gram": round(maxj, 3), "examples_exact": exact[:10], "examples_near": near[:10]}
json.dump(rep, open("logs/leakage_check.json", "w"), indent=1); print(rep)
bad = {n[0] for n in near} | set(exact)
clean = [e for e in tr if e["problem"][:80] not in bad]
random.Random(20).shuffle(clean)
N = int(sys.argv[1]) if len(sys.argv) > 1 else 1500
with open("data/train_subset.jsonl", "w", encoding="utf-8") as f:
    for k, e in enumerate(clean[:N]): f.write(json.dumps(dict(e, item=k)) + "\n")
print("subset", N, "dropped for leakage", len(tr) - len(clean))