"""Run a student checkpoint over a JSONL and dump per-record predictions (no metrics here; see student_report.py).

Each output line: id, per-question probs + logits (choice/noul/score), per-box-question decoded box
[top, bottom, left, right] + IoU-confidence + max-prob confidence, plus gold, normalized teacher, scenario and box_pairs
copied from the input so the report is standalone. --sample N evaluates a seeded random subset (e.g. train-gap check).
--resume appends and skips ids already in the output.

usage: python scripts/eval_student.py --ckpt checkpoints/x --data d.jsonl --out runs/preds.jsonl [--sample 300 --seed 0] [--resume]
"""
import os; os.environ.setdefault("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "1")
import argparse, json, random, sys, time
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
from glint.joint_schema_model import collate_records
from glint.schema import encode_glint_record
from glint.box import box_iou, decode_box, max_prob_confidence


@torch.inference_mode()
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True); ap.add_argument("--data", required=True); ap.add_argument("--out", required=True)
    ap.add_argument("--sample", type=int, default=None); ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--resume", action="store_true"); ap.add_argument("--stats", default=None)
    args = ap.parse_args(); dev = default_device()
    rows = load_jsonl(args.data)
    if args.sample is not None and args.sample < len(rows):
        keep = set(random.Random(args.seed).sample(range(len(rows)), args.sample))
        rows = [r for i, r in enumerate(rows) if i in keep]
    done = set()
    if args.resume and Path(args.out).exists():
        for line in Path(args.out).read_text(encoding="utf-8").splitlines():
            try:
                done.add(json.loads(line)["id"])
            except (ValueError, KeyError):
                pass
    todo = [r for r in rows if r["id"] not in done]
    print(f"{len(rows)} records, {len(done)} already done, {len(todo)} to go", flush=True)
    t0 = time.time(); model, proc, meta = load_checkpoint(args.ckpt, dev); load_s = time.time() - t0
    print(f"loaded {args.ckpt} in {load_s:.1f}s", flush=True)
    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
    lat = []
    if dev == "cuda":
        torch.cuda.reset_peak_memory_stats()
    with open(args.out, "a" if args.resume else "w", encoding="utf-8", newline="\n") as f:
        for i, ex in enumerate(todo):
            ts = time.time()
            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)))
            out = {"id": ex["id"], "probs": {}, "logits": {}, "boxes": {}}
            for q, z in zip(enc.questions, all_logits[0]):
                z = z.float().cpu()
                out["logits"][q.question_id] = dict(zip(q.option_ids, [round(v, 5) for v in z.tolist()]))
                out["probs"][q.question_id] = dict(zip(q.option_ids, [round(v, 6) for v in z.softmax(-1).tolist()]))
            if boxes[0] is not None:
                b = boxes[0]; edges = decode_box(b.logits).cpu(); conf = b.conf_logit.float().sigmoid().cpu()
                maxp = max_prob_confidence(b.logits).cpu(); gb = normalize_box_gold(ex)
                for j, qid in enumerate(b.question_ids):
                    g = gb.get(qid)
                    out["boxes"][qid] = {"box": [round(v, 5) for v in edges[j].tolist()], "conf": round(float(conf[j]), 5),
                                         "maxp": round(float(maxp[j]), 5),
                                         "iou": None if g is None else round(float(box_iou(edges[j], torch.tensor(g))), 5)}
            if dev == "cuda":
                torch.cuda.synchronize()
            dt = time.time() - ts; lat.append(dt)
            m = ex.get("meta") or {}
            out.update({"sec": round(dt, 3), "gold": ex.get("gold") or {}, "teacher": normalize_teacher(ex),
                        "meta": {"scenario": m.get("scenario"), "box_pairs": m.get("box_pairs"), "lang": m.get("lang"),
                                 "theme": m.get("theme"), "style_look": m.get("style_look")}})
            f.write(json.dumps(out, ensure_ascii=False) + "\n"); f.flush()
            if i % 25 == 0:
                print(json.dumps({"i": i, "id": ex["id"], "sec": round(dt, 3), "elapsed_min": round((time.time() - t0) / 60, 1)}), flush=True)
    srt = sorted(lat[1:] or lat)
    stats = {"ckpt": args.ckpt, "data": args.data, "n": len(lat), "load_seconds": load_s,
             "median_sec": srt[len(srt) // 2] if srt else None, "total_seconds": time.time() - t0,
             "peak_mem_gb": torch.cuda.max_memory_allocated() / 2**30 if dev == "cuda" else None}
    print(json.dumps(stats), flush=True)
    if args.stats:
        Path(args.stats).write_text(json.dumps(stats, indent=1), encoding="utf-8")


if __name__ == "__main__":
    main()