"""HF-stack decode speed probe (go/no-go for the HF port of our method).
Loads one model (bf16 / compressed-tensors / GPTQ -- auto-detected from the
checkpoint), generates greedily on image prompts, reports decode tok/s with
prefill subtracted via a 1-token call. transformers-native generate, KV cache.
"""
import argparse, io, json, statistics, time
import torch
from PIL import Image
import pyarrow.parquet as pq

from stage0_gonogo import build_inputs


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--tok-model", default="")
    ap.add_argument("--parquet", required=True)
    ap.add_argument("--pids-from", required=True)
    ap.add_argument("--n", type=int, default=3)
    ap.add_argument("--max-new", type=int, default=256)
    ap.add_argument("--img-max-side", type=int, default=1024)
    ap.add_argument("--bf16", action="store_true", help="force bf16 dtype (unquantized)")
    args = ap.parse_args()

    from transformers import AutoProcessor, AutoModelForImageTextToText
    proc = AutoProcessor.from_pretrained(args.tok_model or args.model)
    kw = dict(device_map="cuda:0")
    if args.bf16:
        kw["dtype"] = torch.bfloat16
    try:
        model = AutoModelForImageTextToText.from_pretrained(args.model, **kw)
    except TypeError:
        kw.pop("dtype", None)
        if args.bf16:
            kw["torch_dtype"] = torch.bfloat16
        model = AutoModelForImageTextToText.from_pretrained(args.model, **kw)
    model.eval()
    print(f"[init] {args.model} loaded, mem={torch.cuda.memory_allocated()/2**30:.1f}GiB",
          flush=True)

    pids = []
    for line in open(args.pids_from):
        d = json.loads(line)
        if "pid" in d:
            pids.append(d["pid"])
        if len(pids) >= args.n:
            break
    tbl = pq.read_table(args.parquet, columns=["pid", "query", "decoded_image"])
    rowmap = {r["pid"]: r for r in tbl.to_pylist()}

    speeds = []
    for pid in pids:
        row = rowmap[pid]
        img = Image.open(io.BytesIO(row["decoded_image"]["bytes"])).convert("RGB")
        if max(img.size) > args.img_max_side:
            sc = args.img_max_side / max(img.size)
            img = img.resize((int(img.width * sc), int(img.height * sc)))
        enc = build_inputs(proc, row["query"], img)
        enc = {k: (v.to("cuda:0") if hasattr(v, "to") else v) for k, v in enc.items()}
        with torch.inference_mode():
            model.generate(**enc, max_new_tokens=4, do_sample=False)
            torch.cuda.synchronize()
            t0 = time.time()
            model.generate(**enc, max_new_tokens=1, do_sample=False)
            torch.cuda.synchronize()
            t_pref = time.time() - t0
            t0 = time.time()
            out = model.generate(**enc, max_new_tokens=args.max_new, do_sample=False)
            torch.cuda.synchronize()
            t_full = time.time() - t0
        ntok = out.shape[1] - enc["input_ids"].shape[1]
        dec = (ntok - 1) / max(t_full - t_pref, 1e-6)
        speeds.append(dec)
        print(f"[{pid}] gen={ntok} prefill={t_pref:.2f}s decode={dec:.1f} tok/s", flush=True)
    print(f"[agg] model={args.model} n={len(speeds)} "
          f"decode median={statistics.median(speeds):.1f} mean={statistics.mean(speeds):.1f} tok/s",
          flush=True)


if __name__ == "__main__":
    main()