"""Measure train step time, inference latency and peak memory for several max_pixels settings."""
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, statistics, time
import torch
from glint.model import build_student, load_processor, default_device
from glint.data import load_jsonl, to_record, normalize_teacher, normalize_box_gold, derive_presence_gold
from glint.losses import decision_loss, record_box_loss
from glint.box import is_box_question
from glint.joint_schema_model import collate_records
from glint.schema import encode_glint_record as encode_record   # == Clef encode_record unless box questions exist

def set_pixels(proc, max_pixels):
    ip = proc.image_processor
    ip.size = {"shortest_edge": 65536, "longest_edge": max_pixels}
    if hasattr(ip, "max_pixels"):
        ip.max_pixels = max_pixels

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", required=True)
    ap.add_argument("--model", default="Qwen/Qwen3.5-2B")
    ap.add_argument("--pixels", default="430080,1048576")
    ap.add_argument("--train-steps", type=int, default=8)
    ap.add_argument("--infer-runs", type=int, default=8)
    ap.add_argument("--dtype", default="bf16")
    ap.add_argument("--grad-ckpt", action="store_true")
    args = ap.parse_args(); dev = default_device(); sync = torch.cuda.synchronize if dev == "cuda" else (lambda: None)
    data = load_jsonl(args.data); results = []
    has_box = any(is_box_question(q) for ex in data for q in ex["questions"].values())   # box head only when the data asks for boxes
    model, proc, meta = build_student(args.model, 430080, dev, args.dtype, grad_ckpt=args.grad_ckpt,
                                      box_head_cfg={} if has_box else None)
    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=1e-5)
    for mp in [int(x) for x in args.pixels.split(",")]:
        set_pixels(proc, mp)
        encs = [(ex, encode_record(proc.tokenizer, to_record(ex), processor=proc)) for ex in data[: max(args.train_steps, args.infer_runs) + 2]]
        grid = encs[0][1].media["image_grid_thw"][0].tolist()
        img_tokens = grid[0] * grid[1] * grid[2] // 4
        # --- train
        model.train(); times = []
        if dev == "cuda":
            torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats()
        for i in range(args.train_steps + 2):
            ex, enc = encs[i % len(encs)]
            sync(); t0 = time.perf_counter()
            batch = collate_records([enc], proc.tokenizer.pad_token_id, torch.device(dev))
            if has_box:
                logits, boxes = model.decide(batch)
                loss, _ = decision_loss(logits[0], enc, normalize_teacher(ex), derive_presence_gold(ex))
                bl, _, _ = record_box_loss(boxes[0], normalize_box_gold(ex))
                loss = loss + (bl if bl is not None else 0.0)
            else:
                loss, _ = decision_loss(model(batch)[0], enc, normalize_teacher(ex), ex.get("gold") or {})
            loss.backward(); opt.step(); opt.zero_grad(set_to_none=True)
            sync(); times.append(time.perf_counter() - t0)
        train_peak = torch.cuda.max_memory_allocated() / 2**30 if dev == "cuda" else None
        # --- inference (LoRA unmerged, same as training graph)
        model.eval(); lat = []
        if dev == "cuda":
            torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats()
        with torch.inference_mode():
            for i in range(args.infer_runs + 2):
                ex = data[i % len(data)]
                sync(); t0 = time.perf_counter()
                enc = encode_record(proc.tokenizer, to_record(ex), processor=proc)
                model.decide(collate_records([enc], proc.tokenizer.pad_token_id, torch.device(dev)))
                sync(); lat.append(time.perf_counter() - t0)
        infer_peak = torch.cuda.max_memory_allocated() / 2**30 if dev == "cuda" else None
        r = {"max_pixels": mp, "image_grid_thw": grid, "image_tokens": img_tokens, "seq_tokens": len(encs[0][1].input_ids),
             "train_step_s_median": round(statistics.median(times[2:]), 3), "train_peak_mem_gb": train_peak and round(train_peak, 2),
             "infer_ms_median": round(statistics.median(lat[2:]) * 1000, 1), "infer_ms_min": round(min(lat[2:]) * 1000, 1),
             "infer_peak_mem_gb": infer_peak and round(infer_peak, 2)}
        results.append(r); print(json.dumps(r), flush=True)
    print(json.dumps({"model": args.model, "dtype": args.dtype, "device": torch.cuda.get_device_name(0) if dev == "cuda" else "cpu",
                      "results": results}, indent=1))

if __name__ == "__main__":
    main()