"""General-skills dev set, disjoint from the test suite: MMLU-Pro (20/category, other questions), BBH (10/task,
other examples), GSM8K train (200). Used to catch regressions outside the trained domains during development."""
import json, os, random
from datasets import load_dataset
ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")
rng = random.Random(99)
test_ids = set()
for f in ("mmlupro", "bbh"):
    test_ids |= {json.loads(l)["id"] for l in open(os.path.join(ROOT, "data/eval", f + ".jsonl"), encoding="utf-8")}
rows = []
ds = load_dataset("TIGER-Lab/MMLU-Pro", split="test")
by = {}
for r in ds:
    if f"mmlupro/{r['question_id']}" not in test_ids:
        by.setdefault(r["category"], []).append(r)
for cat in sorted(by):
    for r in rng.sample(by[cat], 20):
        letters = "ABCDEFGHIJ"[: len(r["options"])]
        opts = "\n".join(f"{l}. {o}" for l, o in zip(letters, r["options"]))
        rows.append({"id": f"dev_mmlupro/{r['question_id']}", "task": "mmlupro", "category": cat,
                     "prompt": f"{r['question']}\n\n{opts}\n\nAnswer with the letter of the correct option. "
                               f"End your response with 'The answer is X.' where X is the letter.", "answer": r["answer"]})
import importlib.util
spec = importlib.util.spec_from_file_location("pe", os.path.join(ROOT, "scripts/prep_eval.py"))
BBH_TASKS = [l.strip().strip('",') for l in open(os.path.join(ROOT, "scripts/prep_eval.py"), encoding="utf-8").read().split("BBH_TASKS = [")[1].split("]")[0].replace("\n", " ").split(",")]
BBH_TASKS = [t.strip().strip('"') for t in BBH_TASKS if t.strip()]
for t in BBH_TASKS:
    ds = load_dataset("lukaemon/bbh", t, split="test")
    idx = [i for i in range(len(ds)) if f"bbh/{t}/{i}" not in test_ids]
    for i in rng.sample(idx, 10):
        r = ds[i]
        rows.append({"id": f"dev_bbh/{t}/{i}", "task": "bbh", "subtask": t,
                     "prompt": r["input"] + "\n\nThink step by step, then end your response with 'The answer is: <answer>'.",
                     "answer": r["target"]})
ds = load_dataset("openai/gsm8k", "main", split="train")
for i in rng.sample(range(len(ds)), 200):
    r = ds[i]
    rows.append({"id": f"dev_gsm8k/{i}", "task": "gsm8k",
                 "prompt": r["question"] + "\n\nPlease reason step by step, and put your final answer within \boxed{}.",
                 "answer": r["answer"].split("####")[-1].strip().replace(",", "")})
with open(os.path.join(ROOT, "data/train/dev_general.jsonl"), "w", encoding="utf-8") as f:
    for r in rows:
        f.write(json.dumps(r, ensure_ascii=False) + "\n")
print(len(rows), len(BBH_TASKS))