"""Stage 2: turn graded answers into Jedi's verifier dataset.

One example per (task, agent answer): Jedi reads `judge_view(...)` and must reply `accept` (the answer was graded
correct) or `escalate` (graded wrong). Splits are by task, never by answer, so no task leaks across splits:
  * test - the in-domain Shaktimaan-routing test tasks (text + code) plus 10% of the new-domain tasks;
  * ood  - every BIG-Bench Hard task (never trained on: measures generalisation to unseen task families);
  * val  - 6% of the remaining tasks, used only to calibrate confidence (Platt scaling);
  * train - the rest. Wrong answers are the minority, so they are oversampled in train to ~40% of examples.
Extra samples from the borderline pass (gen.py BORDERLINE=1) are used for train only, so held-out splits keep the
same, unselected sampling for every task.
"""
import json, random
from collections import Counter
from datasets import Dataset, DatasetDict, load_dataset
from common import JEDI_SYSTEM, judge_view

random.seed(0)
WRONG_SHARE = 0.40
BASE_SAMPLES = {"glm-5.3-flash": 2, "qwen3.8-27b": 1, "kimi-vl": 1}


def main():
    tasks = {t["id"]: t for t in map(json.loads, open("data/tasks.jsonl"))}
    answers = [a for a in map(json.loads, open("data/answers.jsonl")) if not a.get("error") and not a.get("skipped")]
    shak_test = set(load_dataset("podstack/Shaktimaan-routing", split="test")["id"])
    split = {}
    for tid, t in tasks.items():
        if t["source"] == "bbh":
            split[tid] = "ood"
        elif tid in shak_test or (t["source"] in ("aqua", "csqa") and random.random() < 0.10):
            split[tid] = "test"
        else:
            split[tid] = "val" if random.random() < 0.06 else "train"
    rows = {k: [] for k in ("train", "val", "test", "ood")}
    for a in answers:
        t = tasks[a["task_id"]]
        if split[t["id"]] != "train" and a["sample"] >= BASE_SAMPLES.get(a["agent"], 1):
            continue
        verdict = "accept" if a["correct"] else "escalate"
        view = judge_view(t["prompt"], a["text"], a["agent"], t["kind"], a["finish"])
        rows[split[t["id"]]].append(dict(
            task_id=t["id"], source=t["source"], kind=t["kind"], agent=a["agent"], sample=a["sample"],
            correct=bool(a["correct"]), finish=a["finish"] or "", completion_tokens=a["completion_tokens"],
            latency=a["latency"], pred=a["pred"], gold=t["gold"], judge_input=view, verdict=verdict,
            prompt=[{"role": "system", "content": JEDI_SYSTEM}, {"role": "user", "content": view}],
            completion=[{"role": "assistant", "content": verdict}]))
    train = rows["train"]
    wrong = [r for r in train if not r["correct"]]
    need = int(WRONG_SHARE * (len(train) - len(wrong)) / (1 - WRONG_SHARE)) - len(wrong)
    train += random.choices(wrong, k=max(0, need))
    random.shuffle(train)
    dd = DatasetDict({k: Dataset.from_list(v) for k, v in rows.items()})
    dd.save_to_disk("data/jedi_judge")
    print(dd)
    for k, v in rows.items():
        print(f"{k:5s} tasks={len({r['task_id'] for r in v}):5d} answers={len(v):6d} "
              f"wrong={sum(not r['correct'] for r in v) / len(v):.1%} "
              f"agents={dict(Counter(r['agent'] for r in v))}")
    print("solve rate by source x agent (all answers):")
    by = Counter((tasks[a["task_id"]]["source"], a["agent"]) for a in answers)
    ok = Counter((tasks[a["task_id"]]["source"], a["agent"]) for a in answers if a["correct"])
    for src in sorted({s for s, _ in by}):
        print(f"  {src:10s} " + " ".join(f"{ag}={ok[src, ag] / by[src, ag]:.2f}" for s, ag in sorted(by) if s == src))


if __name__ == "__main__":
    main()