"""How fast does a model answer the way the live server asks it?

The server scales the 2040 px camera frame by ``infer_scale`` (0.5 in production) and
hands it to the model, which letterboxes it to ``imgsz`` 640. This times exactly that call
in PyTorch on this machine (the beamline runs a TensorRT engine on its own GPU, so the
absolute number differs there; the ratio between models carries over):

    python tools/bench_latency.py runs/<id>/ultralytics/train/weights/best.pt --out runs/<id>/latency.json

Prints and writes one JSON line: median and p90 milliseconds per frame, frames per second.
"""
from __future__ import annotations

import argparse
import json
import statistics
import sys
import time
from pathlib import Path

import cv2

ROOT = Path(__file__).resolve().parent.parent


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("weights")
    ap.add_argument("--images", default=str(ROOT / "data" / "gold" / "images"), help="frames to time on")
    ap.add_argument("--frames", type=int, default=150)
    ap.add_argument("--warmup", type=int, default=20)
    ap.add_argument("--imgsz", type=int, default=640)
    ap.add_argument("--infer-scale", type=float, default=0.5)
    ap.add_argument("--device", default="0")
    ap.add_argument("--half", action="store_true", default=True)
    ap.add_argument("--out", default="")
    args = ap.parse_args()

    from ultralytics import YOLO
    from ultralytics.utils import LOGGER
    LOGGER.setLevel("ERROR")   # 'half' is deprecated in 8.4 but still does the job; no warning per frame

    paths = sorted(Path(args.images).glob("*.jpg"))[: args.frames + args.warmup]
    if len(paths) < args.warmup + 10:
        sys.exit(f"not enough frames under {args.images}")
    frames = []
    for p in paths:
        img = cv2.imread(str(p))
        if img is None:
            continue
        if args.infer_scale < 1.0:
            img = cv2.resize(img, (max(1, int(img.shape[1] * args.infer_scale)), max(1, int(img.shape[0] * args.infer_scale))),
                             interpolation=cv2.INTER_AREA)
        frames.append(img)
    model = YOLO(args.weights)
    kw = dict(imgsz=args.imgsz, device=args.device, half=args.half, conf=0.25, verbose=False)
    for img in frames[: args.warmup]:
        model.predict(img, **kw)
    times = []
    for img in frames[args.warmup:]:
        t = time.perf_counter()
        model.predict(img, **kw)
        times.append((time.perf_counter() - t) * 1000.0)
    times.sort()
    med = statistics.median(times)
    out = {"weights": str(args.weights), "frames": len(times), "imgsz": args.imgsz, "infer_scale": args.infer_scale,
           "half": bool(args.half), "ms_median": round(med, 2), "ms_p90": round(times[int(0.9 * (len(times) - 1))], 2),
           "ms_max": round(times[-1], 2), "fps_median": round(1000.0 / med, 1), "device": args.device,
           "backend": "pytorch"}
    try:
        import torch
        out["gpu"] = torch.cuda.get_device_name(0)
        import ultralytics
        out["ultralytics"] = ultralytics.__version__
        out["torch"] = torch.__version__
    except Exception:
        pass
    line = json.dumps(out)
    print(line)
    if args.out:
        Path(args.out).write_text(line + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()