"""Accuracy / NLL / Brier / ECE of a checkpoint on a JSONL with gold labels (and agreement with teacher).
If the data has `box` questions with a gold box, also reports a "box" block: mean IoU, IoU>=0.5 / >=0.75 rates,
mean confidence and ECE(15) of the confidence vs the event IoU>=0.5 (plus the max-prob fallback's ECE)."""
import os; os.environ.setdefault("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "1")  # ROCm: faster SDPA (see RESULTS.md); set to 0 to disable
import argparse, json, math, sys
from pathlib import Path
import torch
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from glint.model import load_checkpoint, default_device
from glint.data import load_jsonl, to_record, normalize_teacher, normalize_box_gold, derive_presence_gold
from glint.joint_schema_model import collate_records
from glint.schema import encode_glint_record
from glint.box import box_iou, box_metrics, decode_box, max_prob_confidence

def ece(confs, hits, bins=15):
    tot, n = 0.0, len(confs)
    for b in range(bins):
        lo, hi = b / bins, (b + 1) / bins
        idx = [i for i, c in enumerate(confs) if lo < c <= hi or (b == 0 and c == 0)]
        if idx:
            tot += len(idx) / n * abs(sum(hits[i] for i in idx) / len(idx) - sum(confs[i] for i in idx) / len(idx))
    return tot

@torch.inference_mode()
def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--ckpt", required=True); ap.add_argument("--data", required=True)
    args = ap.parse_args(); dev = default_device()
    model, proc, _ = load_checkpoint(args.ckpt, dev)
    confs, hits, nll, brier, agree, nt = [], [], 0.0, 0.0, 0, 0
    ious, bconfs, bmaxp, any_box = [], [], [], False
    for ex in load_jsonl(args.data):
        enc = encode_glint_record(proc.tokenizer, to_record(ex), processor=proc)
        all_logits, boxes = model.decide(collate_records([enc], proc.tokenizer.pad_token_id, torch.device(dev)))
        logits = all_logits[0]
        teacher = normalize_teacher(ex)
        if boxes[0] is not None:
            any_box = True; gold_boxes = normalize_box_gold(ex); pred = boxes[0]
            edges = decode_box(pred.logits).cpu(); conf = pred.conf_logit.float().sigmoid().cpu()
            maxp = max_prob_confidence(pred.logits).cpu()
            for i, qid in enumerate(pred.question_ids):
                g = gold_boxes.get(qid)
                if g is not None:
                    ious.append(float(box_iou(edges[i], torch.tensor(g)))); bconfs.append(float(conf[i])); bmaxp.append(float(maxp[i]))
            ex = {**ex, "gold": derive_presence_gold(ex)}
        for q, z in zip(enc.questions, logits):
            p = z.float().softmax(-1).cpu(); pred = int(p.argmax())
            g = (ex.get("gold") or {}).get(q.question_id)
            if g is not None:
                y = q.option_ids.index(str(g)); confs.append(float(p[pred])); hits.append(int(pred == y))
                nll -= math.log(max(float(p[y]), 1e-9)); brier += float(((p - torch.nn.functional.one_hot(torch.tensor(y), len(p))) ** 2).sum())
            if q.question_id in teacher:
                t = teacher[q.question_id]; nt += 1; agree += int(q.option_ids[pred] == max(t, key=t.get))
    n = max(len(hits), 1)
    print(json.dumps({"n_questions": len(hits), "accuracy": sum(hits) / n, "nll": nll / n, "brier": brier / n,
                      "ece_15bin": ece(confs, hits), "teacher_agreement": agree / max(nt, 1),
                      **({"box": box_metrics(ious, bconfs, bmaxp)} if any_box else {})}, indent=1))

if __name__ == "__main__":
    main()