"""Benchmark single-image inference latency for the trained model.

No dataset needed — runs over the bundled sample images in `images/`. From repo root:

    uv run python scripts/bench_speed.py                 # all available devices
    uv run python scripts/bench_speed.py --device cpu     # one device
    uv run python scripts/bench_speed.py --imgsz 640 --runs 30

Reports mean inference latency (ms) and FPS per device, warmup excluded. The
"inference" figure is the model forward pass only (Ultralytics' `speed['inference']`),
matching the methodology of the model-card inference table.
"""

from __future__ import annotations

import argparse
from pathlib import Path

import torch
from ultralytics import YOLO

WEIGHTS = "runs/detect/yolov8n_v1_train/weights/best.pt"
IMAGES_DIR = "images"


def available_devices() -> list[str]:
    devices = ["cpu"]
    if torch.cuda.is_available():
        devices.append("0")
    if torch.backends.mps.is_available():
        devices.append("mps")
    return devices


def bench(weights: str, device: str, images: list[str], imgsz: int, runs: int) -> tuple[float, float]:
    model = YOLO(weights)
    # Warmup — first calls pay lazy init / kernel-compile costs we don't want to time.
    model.predict(images[0], device=device, imgsz=imgsz, verbose=False)

    latencies: list[float] = []
    for _ in range(runs):
        for img in images:
            r = model.predict(img, device=device, imgsz=imgsz, verbose=False)
            latencies.append(r[0].speed["inference"])  # ms, forward pass only

    mean_ms = sum(latencies) / len(latencies)
    return mean_ms, 1000.0 / mean_ms


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--weights", default=WEIGHTS)
    p.add_argument("--images", default=IMAGES_DIR)
    p.add_argument("--device", default=None, help="cpu / mps / 0 (GPU); omit to run all available")
    p.add_argument("--imgsz", type=int, default=640)
    p.add_argument("--runs", type=int, default=20, help="passes over the image set per device")
    args = p.parse_args()

    images = sorted(str(p) for p in Path(args.images).glob("*") if p.suffix.lower() in {".jpg", ".jpeg", ".png"})
    if not images:
        raise SystemExit(f"No images found in {args.images!r}")

    devices = [args.device] if args.device else available_devices()
    print(f"Weights: {args.weights}")
    print(f"Images:  {len(images)} from {args.images}/  (imgsz={args.imgsz}, {args.runs} passes each)\n")
    print(f"{'Device':<10} {'Latency (ms)':>14} {'FPS':>8}")
    print("-" * 34)
    for dev in devices:
        mean_ms, fps = bench(args.weights, dev, images, args.imgsz, args.runs)
        label = {"cpu": "CPU", "mps": "MPS (GPU)", "0": "CUDA GPU"}.get(dev, dev)
        print(f"{label:<10} {mean_ms:>14.1f} {fps:>8.1f}")


if __name__ == "__main__":
    main()