"""Export every scored eval answer into release/dataset/eval/ (Parquet) so the results tables can be audited.

  .venv/bin/python -I scripts/export_eval_data.py

  eval/benchmarks/<model>_loops<L>[_efficient].parquet   one row per sampled answer (GPQA, MATH-500, AIME, AMC, held-out)
  eval/ifeval/<model>_loops2.parquet                     one row per IFEval answer
  eval/problems.parquet                                  prompt + reference answer per problem, except GPQA

Models: dory1.0-4b (base), dory-beam-4b (v2, with reasoning_effort; `reasoning_mode` effort-low/medium/high or
normal for the untagged default), dory-beam-4b-v1 (the first release), dory-beam-4b-accurate, and ablations.
Base Dory and v1 were run with a 32k limit (16k for IFEval) and every answer that hit it was resumed to 64k by
scripts/spawn_continue.py (`resumed_tokens` > 0). v2 and Accurate ran with 64k from the start. The ablations rarely hit
32k and stay there; `max_output_tokens` says which limit applies.

GPQA's authors ask that its questions not appear online in plain text (to keep them out of training data), so GPQA
rows keep their grades, extracted letters and token counts but drop the question and the model's text.
"""

import json
from pathlib import Path

import pyarrow as pa
import pyarrow.parquet as pq

ROOT = Path(__file__).resolve().parent.parent
RES = ROOT / "data" / "eval" / "results"
OUT = ROOT / "release" / "dataset" / "eval"

# local result file -> (model name in the dataset, loops, reasoning mode)
RUNS = {
    "base_loops1-64k": ("dory1.0-4b", 1, "normal"), "base_loops2-64k": ("dory1.0-4b", 2, "normal"),
    "base_loops4-64k": ("dory1.0-4b", 4, "normal"), "base_loops2eff": ("dory1.0-4b", 2, "efficient"),
    "mix050_loops1-64k": ("dory-beam-4b-v1", 1, "normal"), "mix050_loops2-64k": ("dory-beam-4b-v1", 2, "normal"),
    "mix050_loops4-64k": ("dory-beam-4b-v1", 4, "normal"),
    "rE-t050_loops2-default-64k": ("dory-beam-4b", 2, "normal"),
    "rE-t050_loops2-low-64k": ("dory-beam-4b", 2, "effort-low"),
    "rE-t050_loops2-medium-64k": ("dory-beam-4b", 2, "effort-medium"),
    "rE-t050_loops2-high-64k": ("dory-beam-4b", 2, "effort-high"),
    "rE-t050_loops4-high-64k": ("dory-beam-4b", 4, "effort-high"),
    "mix030_loops2-64k": ("dory-beam-4b-accurate", 2, "normal"),
    "mix030_loops4-64k": ("dory-beam-4b-accurate", 4, "normal"),
    "mix070_loops2": ("blend-70-30", 2, "normal"),
    "rA_loops1": ("run-a-sft", 1, "normal"), "rA_loops2": ("run-a-sft", 2, "normal"),
    "rA_loops4": ("run-a-sft", 4, "normal"),
    "rA400_loops2": ("run-a-step400", 2, "normal"), "rA800_loops2": ("run-a-step800", 2, "normal"),
    "rA2_loops1": ("run-a2-random-trace", 1, "normal"), "rA2_loops2": ("run-a2-random-trace", 2, "normal"),
    "rA2_loops4": ("run-a2-random-trace", 4, "normal"),
}
IFEVAL = {"base_ifeval_loops2-64k": "dory1.0-4b", "mix050_ifeval_loops2-64k": "dory-beam-4b-v1",
          "rE-t050_ifeval_loops2-64k": "dory-beam-4b"}


def limit(tag):
    return 65536 if tag.endswith("-64k") else 16384 if "ifeval" in tag else 32768


def main():
    problems = {json.loads(l)["id"]: json.loads(l) for l in open(ROOT / "data" / "evalsets.jsonl")}
    for f in list((OUT / "benchmarks").glob("*.parquet")) + list((OUT / "ifeval").glob("*.parquet")):
        f.unlink()  # file names changed between releases; start clean
    (OUT / "benchmarks").mkdir(parents=True, exist_ok=True)
    (OUT / "ifeval").mkdir(parents=True, exist_ok=True)

    for tag, (model, loops, mode) in RUNS.items():
        rows = []
        for line in open(RES / f"{tag}.jsonl"):
            r = json.loads(line)
            gpqa = r["bench"] == "gpqa"
            rows.append({"model": model, "loops": loops, "reasoning_mode": mode, "bench": r["bench"],
                         "problem_id": r["id"], "sample": r["sample"], "correct": r["correct"],
                         "extracted_answer": r["pred"], "reference_answer": problems[r["id"]]["answer"],
                         "output_tokens": r["tokens"], "finish_reason": r["finish"],
                         "max_output_tokens": limit(tag), "resumed_tokens": r.get("continued_tokens", 0),
                         "thinking_closed": r["thinking_closed"], "text": None if gpqa else r["text"]})
        name = f"{model}_loops{loops}{'' if mode == 'normal' else '_' + mode}.parquet"
        pq.write_table(pa.Table.from_pylist(rows), OUT / "benchmarks" / name, compression="zstd")
        print(f"{name}: {len(rows)} rows")

    for tag, model in IFEVAL.items():
        rows = []
        for line in open(RES / f"{tag}.jsonl"):
            r = json.loads(line)
            rows.append({"model": model, "loops": 2, "key": r["key"], "sample": r["sample"],
                         "strict_prompt": r["strict_prompt"], "strict_instructions": r["strict_inst"],
                         "loose_prompt": r["loose_prompt"], "loose_instructions": r["loose_inst"],
                         "output_tokens": r["tokens"], "finish_reason": r["finish"],
                         "max_output_tokens": limit(tag), "resumed_tokens": r.get("continued_tokens", 0),
                         "text": r["text"]})
        pq.write_table(pa.Table.from_pylist(rows), OUT / "ifeval" / f"{model}_loops2.parquet", compression="zstd")
        print(f"ifeval/{model}_loops2.parquet: {len(rows)} rows")

    rows = [{"bench": p["bench"], "problem_id": p["id"], "prompt": p["prompt"], "reference_answer": p["answer"],
             "answer_type": p["answer_type"], "samples": p["samples"]} for p in problems.values() if p["bench"] != "gpqa"]
    pq.write_table(pa.Table.from_pylist(rows), OUT / "problems.parquet", compression="zstd")
    print(f"problems.parquet: {len(rows)} rows")


if __name__ == "__main__":
    main()