"""Assemble the evaluation suite into data/evalsets.jsonl (one row per problem).

  gpqa     fingertap/GPQA-Diamond (198, letter answers)
  math500  HuggingFaceH4/MATH-500 (500)
  aime24   HuggingFaceH4/aime_2024 (30)
  aime25   math-ai/aime25 (30)
  amc23    math-ai/amc23 (40)
  heldout  our 300 held-out pool prompts (data/eval_ids.txt), never trained on
"""

import json
from pathlib import Path

import duckdb

ROOT = Path(__file__).resolve().parent.parent
MCQ_PREFIX = ("Solve the following problem. Make sure to put the answer (and only answer) "
              "inside \\boxed{}.\n\n")
MATH_SUFFIX = "\n\nPut your final answer within \\boxed{}."
# GPQA options are listed as "A. ..."; ask for the letter explicitly (as standard GPQA evals do),
# otherwise models box the option's value instead of its letter.
GPQA_PREFIX = ("Answer the following multiple choice question. Put only the letter of the correct "
               "option (A, B, C or D) inside \\boxed{}.\n\n")
SAMPLES = {"gpqa": 2, "math500": 1, "aime24": 8, "aime25": 8, "amc23": 4, "heldout": 1}


def main():
    con = duckdb.connect()
    con.execute("INSTALL httpfs; LOAD httpfs;")
    rows = []

    def q(sql):
        return con.execute(sql).fetchall()

    for i, (question, answer) in enumerate(q("SELECT question, answer FROM 'hf://datasets/fingertap/GPQA-Diamond/**/*.parquet'")):
        rows.append({"bench": "gpqa", "id": f"gpqa-{i}", "prompt": GPQA_PREFIX + question.strip(),
                     "answer": answer.strip(), "answer_type": "mcq"})
    for i, (problem, answer) in enumerate(q("SELECT problem, answer FROM read_json_auto('hf://datasets/HuggingFaceH4/MATH-500/test.jsonl')")):
        rows.append({"bench": "math500", "id": f"math500-{i}", "prompt": problem.strip() + MATH_SUFFIX,
                     "answer": str(answer), "answer_type": "math"})
    for bench, sql in (("aime24", "SELECT problem, answer FROM 'hf://datasets/HuggingFaceH4/aime_2024/**/*.parquet'"),
                       ("aime25", "SELECT problem, answer FROM read_json_auto('hf://datasets/math-ai/aime25/test.jsonl')"),
                       ("amc23", "SELECT question, answer FROM 'hf://datasets/math-ai/amc23/**/*.parquet'")):
        for i, (problem, answer) in enumerate(q(sql)):
            rows.append({"bench": bench, "id": f"{bench}-{i}", "prompt": problem.strip() + MATH_SUFFIX,
                         "answer": str(answer).strip(), "answer_type": "math"})

    eval_ids = [l.strip() for l in open(ROOT / "data" / "eval_ids.txt") if l.strip()]
    pool = {json.loads(l)["id"]: json.loads(l) for l in open(ROOT / "data" / "pool.jsonl")}
    for pid in eval_ids:
        p = pool[pid]
        rows.append({"bench": "heldout", "id": pid, "prompt": p["prompt"], "answer": p["answer"],
                     "answer_type": p["answer_type"], "domain": p["domain"]})

    for r in rows:
        r["samples"] = SAMPLES[r["bench"]]
    out = ROOT / "data" / "evalsets.jsonl"
    with open(out, "w") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    from collections import Counter
    print(Counter(r["bench"] for r in rows), f"-> {out}")


if __name__ == "__main__":
    main()