"""Per-item outputs for the audit: every test question with the full probability vector.

    PYTHONPATH=<decisions worktree> python scripts/decisions/eval_items.py --checkpoint CKPT --name v3 \
        --test typed_en=PATH --test tasksource=PATH ... --out DIR [--engram-off]

Same reader and prompt as train_decisions.py (decisions.Reader, prefill num_logits=1, bf16), one
question per prompt. Writes DIR/<name>[-engram_off]__<test>.jsonl with one line per question:
    {test, id, qid, type, grupo, keys, probs, pred, gold, soft}
`gold` is the list of accepted keys, `soft` the teacher distribution when the data has one.
"""
from __future__ import annotations

import argparse
import contextlib
import json
import sys
import time
from pathlib import Path

import torch

sys.path.insert(0, str(Path(__file__).resolve().parent))
from decisions import Reader, read_jsonl  # noqa: E402


@torch.no_grad()
def run(reader, cases, test, fh):
    n = 0
    for c in cases:
        lang = c.get("lang", "es")
        for qid, q in c["questions"].items():
            ids, keys = reader.prompt(c["state"], q, lang)
            p = torch.softmax(reader.logits_eval(ids, len(keys)), -1).tolist()
            soft = (c.get("suave") or {}).get(qid)
            fh.write(json.dumps({"test": test, "id": c["id"], "qid": qid, "type": q["type"], "grupo": c.get("grupo"),
                                 "keys": keys, "probs": p, "pred": keys[max(range(len(p)), key=p.__getitem__)],
                                 "gold": c["gold"][qid], "soft": soft, "prompt_tokens": len(ids)},
                                ensure_ascii=False) + "\n")
            n += 1
    return n


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--checkpoint", type=Path, required=True)
    ap.add_argument("--name", required=True)
    ap.add_argument("--test", action="append", default=[], help="name=path.jsonl")
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--engram-off", action="store_true")
    a = ap.parse_args()
    from mini_v41.inference import InferenceModel
    im = InferenceModel.from_checkpoint(a.checkpoint, device="cuda:0", dtype="bf16")
    im.model.eval()
    reader = Reader(im)
    ctx = contextlib.nullcontext()
    if a.engram_off:
        from mini_v41.engram import engram_disabled
        ctx = engram_disabled(im.model)
    a.out.mkdir(parents=True, exist_ok=True)
    tag = a.name + ("-engram_off" if a.engram_off else "")
    with ctx:
        for spec in a.test:
            test, path = spec.split("=", 1)
            t0 = time.time()
            with open(a.out / ("%s__%s.jsonl" % (tag, test)), "w", encoding="utf-8") as fh:
                n = run(reader, read_jsonl(path), test, fh)
            print("%s %s %d questions %.0fs" % (tag, test, n, time.time() - t0), flush=True)


if __name__ == "__main__":
    main()