"""Builds the fixed evaluation suite (held-out test splits only) as JSONL files in data/eval.

Every item: {"id", "task", "prompt", ...grader fields}. The suite is frozen once built; both the
original and the modified models are evaluated on exactly these files with the same sampler.
"""
import json
import os
import random

from datasets import load_dataset

OUT = os.path.join(os.path.dirname(__file__), "..", "data", "eval")
os.makedirs(OUT, exist_ok=True)
rng = random.Random(1234)

MATH_SUFFIX = "\n\nPlease reason step by step, and put your final answer within \\boxed{}."


def dump(name, rows):
    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))


# GSM8K test
ds = load_dataset("openai/gsm8k", "main", split="test")
dump("gsm8k", [
    {"id": f"gsm8k/{i}", "task": "gsm8k", "prompt": r["question"] + MATH_SUFFIX,
     "answer": r["answer"].split("####")[-1].strip().replace(",", "")}
    for i, r in enumerate(ds)
])

# MATH-500
ds = load_dataset("HuggingFaceH4/MATH-500", split="test")
dump("math500", [
    {"id": f"math500/{r['unique_id']}", "task": "math500", "prompt": r["problem"] + MATH_SUFFIX,
     "answer": r["answer"], "level": r["level"], "subject": r["subject"]}
    for r in ds
])

# AIME 2024 + 2025 (competition problems, integer answers)
rows = []
ds = load_dataset("HuggingFaceH4/aime_2024", split="train")
for r in ds:
    rows.append({"id": f"aime24/{r['id']}", "task": "aime", "prompt": r["problem"] + MATH_SUFFIX, "answer": str(r["answer"])})
for cfg in ["AIME2025-I", "AIME2025-II"]:
    ds = load_dataset("opencompass/AIME2025", cfg, split="test")
    for i, r in enumerate(ds):
        rows.append({"id": f"aime25/{cfg}/{i}", "task": "aime", "prompt": r["question"] + MATH_SUFFIX, "answer": str(r["answer"])})
dump("aime", rows)

# HumanEval+ and MBPP+ (extended hidden tests from EvalPlus)
ds = load_dataset("evalplus/humanevalplus", split="test")
dump("humaneval", [
    {"id": r["task_id"], "task": "humaneval",
     "prompt": "Complete the following Python function. Return the full function implementation "
               "(including the signature and any needed imports) in a single ```python code block.\n\n```python\n"
               + r["prompt"] + "```",
     "code_prompt": r["prompt"], "entry_point": r["entry_point"], "test": r["test"]}
    for r in ds
])
ds = load_dataset("evalplus/mbppplus", split="test")
dump("mbpp", [
    {"id": f"mbpp/{r['task_id']}", "task": "mbpp",
     "prompt": r["prompt"].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": r["test"], "test_imports": r.get("test_imports", [])}
    for r in ds
])

# IFEval
ds = load_dataset("google/IFEval", split="train")
dump("ifeval", [
    {"id": f"ifeval/{r['key']}", "task": "ifeval", "prompt": r["prompt"],
     "instruction_id_list": r["instruction_id_list"], "kwargs": r["kwargs"]}
    for r in ds
])

# MMLU-Pro: 100 random questions per category (14 categories), 10 options
ds = load_dataset("TIGER-Lab/MMLU-Pro", split="test")
by_cat = {}
for r in ds:
    by_cat.setdefault(r["category"], []).append(r)
rows = []
for cat in sorted(by_cat):
    for r in rng.sample(by_cat[cat], 100):
        letters = "ABCDEFGHIJ"[: len(r["options"])]
        opts = "\n".join(f"{l}. {o}" for l, o in zip(letters, r["options"]))
        rows.append({"id": f"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"]})
dump("mmlupro", rows)

# BIG-Bench Hard: 20 per task, all 27 tasks (multi-step logical reasoning)
BBH_TASKS = ["boolean_expressions", "causal_judgement", "date_understanding", "disambiguation_qa", "dyck_languages",
             "formal_fallacies", "geometric_shapes", "hyperbaton", "logical_deduction_five_objects",
             "logical_deduction_seven_objects", "logical_deduction_three_objects", "movie_recommendation",
             "multistep_arithmetic_two", "navigate", "object_counting", "penguins_in_a_table",
             "reasoning_about_colored_objects", "ruin_names", "salient_translation_error_detection", "snarks",
             "sports_understanding", "temporal_sequences", "tracking_shuffled_objects_five_objects",
             "tracking_shuffled_objects_seven_objects", "tracking_shuffled_objects_three_objects", "web_of_lies",
             "word_sorting"]
rows = []
for t in BBH_TASKS:
    ds = load_dataset("lukaemon/bbh", t, split="test")
    idx = rng.sample(range(len(ds)), 20)
    for i in idx:
        r = ds[i]
        rows.append({"id": f"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"]})
dump("bbh", rows)