#!/usr/bin/env python
"""Write the problem files every lab uses (run once, on a machine with internet).

  math500.jsonl   MATH-500 test (500), gold = its 'answer' field         (the paper's benchmark)
  gsm8k.jsonl     GSM8K main test (1,319), gold = the number after ####   (the paper's benchmark)
  ifeval.jsonl    IFEval (541), graded by the official checkers          (the paper's benchmark)
  dev256.jsonl    128 GSM8K-train + 128 MATH-train practice problems the paper's agent never used:
                  where we measure every edit we try (same 50/50 mix as the agent's held-out sets)
  build128.jsonl  64 + 64 more unused practice problems: where the capstone builds its own edit
"""
import argparse
import json
import os
import random

ap = argparse.ArgumentParser()
ap.add_argument("--exp8", required=True, help="path to lm_understand-logs/exp8_math")
ap.add_argument("--out", required=True)
a = ap.parse_args()
os.makedirs(a.out, exist_ok=True)


def dump(name, rows):
    with open(os.path.join(a.out, f"{name}.jsonl"), "w") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    print(name, len(rows))


from datasets import load_dataset  # noqa: E402

dump("math500", [{"id": r["unique_id"], "source": "math500", "problem": r["problem"], "answer": r["answer"]}
                 for r in load_dataset("HuggingFaceH4/MATH-500", split="test")])
dump("gsm8k", [{"id": f"gsm8k_test_{i}", "source": "gsm8k", "problem": r["question"],
                "answer": r["answer"].split("####")[-1].strip().replace(",", "")}
               for i, r in enumerate(load_dataset("openai/gsm8k", "main", split="test"))])
dump("ifeval", [{"id": f"ifeval_{r['key']}", "source": "ifeval", "kind": "ifeval", "prompt": r["prompt"],
                 "instruction_id_list": r["instruction_id_list"], "kwargs": r["kwargs"]}
                for r in load_dataset("google/IFEval", split="train")])

used = set()
for fn in ("splits.json", "fresh_splits.json"):
    for v in json.load(open(os.path.join(a.exp8, "runs/p50c/work_notes", fn))).values():
        used |= set(v)


def practice(name):
    with open(os.path.join(a.exp8, "practice", f"{name}.jsonl")) as f:
        return [json.loads(line) for line in f if line.strip()]


rng = random.Random(0)
pools = {}
for name in ("gsm8k_train", "math_train"):
    rows = [r for r in practice(name) if r["id"] not in used]
    rng.shuffle(rows)
    pools[name] = rows
    print(name, "unused by the paper's agent:", len(rows))


def pick(rows):
    return [{"id": r["id"], "source": r["source"], "problem": r["problem"], "answer": r["answer"],
             **({"level": r["level"], "subject": r["subject"]} if "level" in r else {})} for r in rows]


dump("dev256", pick(pools["gsm8k_train"][:128] + pools["math_train"][:128]))
dump("build128", pick(pools["gsm8k_train"][128:192] + pools["math_train"][128:192]))