"""Load a Glint checkpoint (LoRA + head) and answer decision questions about screenshots.
Output: JSON with per-option probabilities (SystemOne-style answers) + latency.
`box` questions are answered Jev-1.14 style: {"type": "box", "box": {"top","bottom","left","right"}, "confidence"}."""
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, time
from pathlib import Path
import torch
from PIL import Image
from glint.model import load_checkpoint, default_device
from glint.data import load_jsonl, to_record
from glint.joint_schema_model import collate_records, systemone_answer
from glint.schema import encode_glint_record
from glint.box import box_answer, decode_box

@torch.inference_mode()
def decide(model, proc, record, device, max_length=4096):
    sync = torch.cuda.synchronize if device == "cuda" else (lambda: None)
    sync(); t0 = time.perf_counter()
    enc = encode_glint_record(proc.tokenizer, record, max_length=max_length, processor=proc)
    batch = collate_records([enc], proc.tokenizer.pad_token_id, torch.device(device))
    t1 = time.perf_counter()
    all_logits, boxes = model.decide(batch)
    logits, box_pred = all_logits[0], boxes[0]
    sync(); t2 = time.perf_counter()
    answers, probs = {}, {}
    for q, z in zip(enc.questions, logits):
        p = dict(zip(q.option_ids, z.float().softmax(-1).tolist()))
        probs[q.question_id] = {k: round(v, 4) for k, v in p.items()}
        answers[q.question_id] = systemone_answer(record["questions"][q.question_id], p)
    if box_pred is not None:
        edges = decode_box(box_pred.logits).cpu(); conf = box_pred.conf_logit.float().sigmoid().cpu()
        for i, qid in enumerate(box_pred.question_ids):
            answers[qid] = box_answer(edges[i], float(conf[i]))
    answers = {qid: answers[qid] for qid in record["questions"] if qid in answers}   # request order
    return {"answers": answers, "probabilities": probs,
            "usage": {"input_tokens": len(enc.input_ids)},
            "latency_ms": {"preprocess": round((t1 - t0) * 1000, 1), "forward": round((t2 - t1) * 1000, 1),
                           "total": round((t2 - t0) * 1000, 1)}}

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--image", help="screenshot path (with --questions)")
    ap.add_argument("--questions", help="questions JSON file or inline JSON")
    ap.add_argument("--state", default="", help="state string or JSON")
    ap.add_argument("--jsonl", help="alternatively: run on records of a JSONL file")
    ap.add_argument("--limit", type=int, default=5)
    ap.add_argument("--max-pixels", type=int, default=None, help="override the checkpoint's max_pixels")
    ap.add_argument("--no-merge", action="store_true")
    ap.add_argument("--warmup", type=int, default=1)
    args = ap.parse_args()
    dev = default_device()
    model, proc, meta = load_checkpoint(args.ckpt, dev, args.max_pixels, merge=not args.no_merge)
    if args.jsonl:
        rows = load_jsonl(args.jsonl)[: args.limit]
        records = [(r.get("id"), to_record(r), r.get("gold")) for r in rows]
    else:
        qs = args.questions
        qs = json.loads(Path(qs).read_text(encoding="utf-8")) if Path(qs).exists() else json.loads(qs)
        try:
            state = json.loads(args.state)
        except Exception:
            state = args.state
        records = [(args.image, {"state": state, "images": [Image.open(args.image).convert("RGB")], "questions": qs}, None)]
    for _ in range(args.warmup):
        decide(model, proc, records[0][1], dev)
    for rid, rec, gold in records:
        out = {"id": rid, **decide(model, proc, rec, dev)}
        if gold:
            out["gold"] = gold
        print(json.dumps(out, ensure_ascii=False, indent=1), flush=True)

if __name__ == "__main__":
    main()