"""Export the reasoning-effort data into release/dataset/data/ (Parquet).

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

  effort_sft.parquet     the training set behind Dory-Beam-4B's reasoning_effort (scripts/build_effort_dataset.py)
  effort_traces.parquet  every Beam generation at low, high and xhigh effort, correct or not, with its grade
"""

import glob
import json
import sys
from pathlib import Path

import pyarrow as pa
import pyarrow.parquet as pq

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "scripts"))
from verify import check  # noqa: E402

TEACHER = {"beam-api": "api", "beam-nvfp4-selfhost-sglang": "selfhost-nvfp4", "beam-fp8-selfhost-sglang": "selfhost-fp8"}
OUT = ROOT / "release" / "dataset" / "data"


def main():
    eval_ids = {l.strip() for l in open(ROOT / "data" / "eval_ids.txt") if l.strip()}
    sft = [json.loads(l) for l in open(ROOT / "data" / "final_effort" / "sft_effort.jsonl")]
    used = {}  # (prompt id, Beam effort, reasoning text) -> level it was trained as
    rows = []
    for r in sft:
        assert r["prompt_id"] not in eval_ids
        used[(r["prompt_id"], r["beam_effort"], r["messages"][1]["reasoning_content"])] = r["effort"]
        rows.append({"id": r["id"], "prompt_id": r["prompt_id"], "reasoning_effort": r["effort"],
                     "beam_effort": r["beam_effort"], "tagged": r["tagged"], "domain": r["domain"],
                     "teacher": TEACHER[r["teacher"]], "messages": r["messages"], "gold": r["gold"]})
    pq.write_table(pa.Table.from_pylist(rows), OUT / "effort_sft.parquet", compression="zstd")
    print(f"effort_sft.parquet: {len(rows)}")

    rows = []
    for path in sorted(glob.glob(str(ROOT / "data" / "selfhost_effort" / "*.jsonl"))):
        for line in open(path):
            r = json.loads(line)
            if "error" in r:
                continue
            assert r["id"] not in eval_ids
            ok, pred = check(r["content"], r["gold"], r["answer_type"])
            rows.append({
                "id": r["id"], "domain": r["domain"], "source": r["source"], "topic": r.get("topic"),
                "difficulty": r["difficulty"], "teacher": TEACHER[r["teacher"]], "reasoning_effort": r["effort"],
                "prompt": r["prompt"], "reasoning": r["reasoning"], "content": r["content"], "gold": r["gold"],
                "answer_type": r["answer_type"], "pred": pred,
                "correct": bool(ok and r["finish_reason"] == "stop" and r["reasoning"].strip()),
                "finish_reason": r["finish_reason"], "completion_tokens": r["usage"].get("completion_tokens"),
                "trained_as": used.get((r["id"], r["effort"], r["reasoning"])),
            })
    pq.write_table(pa.Table.from_pylist(rows), OUT / "effort_traces.parquet", compression="zstd")
    n_used = sum(r["trained_as"] is not None for r in rows)
    print(f"effort_traces.parquet: {len(rows)} ({n_used} used in effort_sft)")


if __name__ == "__main__":
    main()