"""Builds the pool of verifiable training problems (math with reference answers, code with unit tests),
decontaminated against every prompt in the frozen eval suite (13-gram overlap on normalized text).

Output: data/train/pool_math.jsonl, data/train/pool_code.jsonl, data/train/dev.jsonl (held-out dev set used for
all intermediate decisions, so the test suite is only touched for the baseline and the final comparison).
"""
import glob
import json
import os
import random
import re

from datasets import load_dataset

ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")
OUT = os.path.join(ROOT, "data", "train")
os.makedirs(OUT, exist_ok=True)
rng = random.Random(7)
MATH_SUFFIX = "\n\nPlease reason step by step, and put your final answer within \\boxed{}."


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


def ngrams(s, n=13):
    w = words(s)
    return {" ".join(w[i:i + n]) for i in range(max(0, len(w) - n + 1))} or ({" ".join(w)} if w else set())


eval_grams = set()
eval_short = set()
for path in glob.glob(os.path.join(ROOT, "data", "eval", "*.jsonl")):
    with open(path, encoding="utf-8") as f:
        for l in f:
            r = json.loads(l)
            text = r.get("code_prompt") or r["prompt"].replace(MATH_SUFFIX, "")
            eval_grams |= ngrams(text)
            eval_short.add(" ".join(words(text))[:200])


def contaminated(text):
    if " ".join(words(text))[:200] in eval_short:
        return True
    return bool(ngrams(text) & eval_grams)


def last_boxed(text):
    idx = text.rfind("\\boxed")
    if idx < 0:
        return None
    i = text.find("{", idx)
    depth = 0
    for j in range(i, len(text)):
        depth += text[j] == "{"
        depth -= text[j] == "}"
        if depth == 0:
            return text[i + 1: j]
    return None


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


def add_math(src, problem, answer, level=None):
    global dropped
    if not answer or len(answer) > 60 or "\\text{" in answer and len(answer) > 25:
        return  # keep answers that a symbolic checker can verify reliably
    key = " ".join(words(problem))[:300]
    if key in seen:
        return
    seen.add(key)
    if contaminated(problem):
        dropped += 1
        return
    math_rows.append({"id": f"{src}/{len(math_rows)}", "task": "math", "src": src, "level": level,
                      "prompt": problem.strip() + MATH_SUFFIX, "answer": answer.strip()})


ds = load_dataset("DigitalLearningGmbH/MATH-lighteval", split="train")
for r in ds:
    lvl = int(r["level"].split()[-1]) if r["level"].split()[-1].isdigit() else None
    if lvl and lvl >= 3:
        add_math("math_train", r["problem"], last_boxed(r["solution"]), lvl)
n_math = len(math_rows)
ds = load_dataset("agentica-org/DeepScaleR-Preview-Dataset", split="train")
for r in ds:
    add_math("deepscaler", r["problem"], r["answer"])
print(f"math: {n_math} from MATH train, {len(math_rows) - n_math} from DeepScaleR, dropped {dropped} contaminated")

code_rows, dropped_code = [], 0
ds = load_dataset("google-research-datasets/mbpp", "full")
for split in ("train", "validation", "prompt"):
    for r in ds[split]:
        if contaminated(r["text"]):
            dropped_code += 1
            continue
        code_rows.append({
            "id": f"mbpp_train/{r['task_id']}", "task": "code_fn", "src": "mbpp_train",
            "prompt": r["text"].strip() + "\nYour code should pass this test:\n" + r["test_list"][0]
                      + "\n\nWrite the solution as a Python function in a single ```python code block.",
            "test": "\n".join(r["test_setup_code"].splitlines() + r["test_list"] + r.get("challenge_test_list", []))})
n_mbpp = len(code_rows)
ds = load_dataset("likaixin/TACO-verified", split="train")
for r in ds:
    q = r["question"]
    if len(q) > 6000 or contaminated(q):
        dropped_code += 1
        continue
    try:
        io = json.loads(r["input_output"])
    except Exception:
        continue
    if io.get("fn_name") or not io.get("inputs") or len(io["inputs"]) < 2:
        continue  # keep stdin/stdout problems with at least two tests
    if any(not isinstance(x, str) for x in io["inputs"] + io["outputs"]):
        continue
    if sum(len(x) for x in io["inputs"]) > 200_000:
        continue
    code_rows.append({
        "id": f"taco/{r['id'] if 'id' in r else len(code_rows)}", "task": "code_io", "src": "taco",
        "difficulty": r.get("difficulty"),
        "prompt": q.strip() + "\n\nWrite a Python 3 program that reads from standard input and writes to standard "
                              "output. Put the complete program in a single ```python code block.",
        "inputs": io["inputs"][:20], "outputs": io["outputs"][:20]})
print(f"code: {n_mbpp} MBPP, {len(code_rows) - n_mbpp} TACO, dropped {dropped_code} contaminated/long")

rng.shuffle(math_rows)
rng.shuffle(code_rows)
# dev: 300 math + 150 code problems, never trained on, used for every intermediate decision
dev = math_rows[:300] + code_rows[:150]
for name, rows in [("dev", dev), ("pool_math", math_rows[300:]), ("pool_code", code_rows[150:])]:
    with open(os.path.join(OUT, f"{name}.jsonl"), "w", encoding="utf-8") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    print(name, len(rows))