"""Model-family benchmark: YOLOv8n / v8s / 11n / 11s on identical config (plan 3.1).

Trains each candidate on the same data, epochs, batch, image size, and seed as
the original yolov8s run, evaluates on the val split, and writes one
accuracy-vs-latency-vs-size table. Meant for a Colab GPU session:

    !uv run python scripts/benchmark_models.py                      # full run (hours)
    !uv run python scripts/benchmark_models.py --models yolov8n,yolo11n --epochs 10   # quick pass
    !uv run python scripts/benchmark_models.py --skip-train         # re-tabulate existing runs

Each model trains into runs/detect/bench_<model>/; --skip-train reuses those
weights, so an interrupted session can resume without retraining finished
models. Output table: docs/benchmark_models.md.
"""

from __future__ import annotations

import argparse
from pathlib import Path

from ultralytics import YOLO

DATA = "work/data.yaml"
# YOLO12 needs ultralytics >= 8.3.78; this repo pins 8.3.40, so the ladder tops
# out at YOLO11. The v8s row doubles as a sanity check against the original run.
MODELS = ["yolov8n", "yolov8s", "yolo11n", "yolo11s"]


def run_name(model: str) -> str:
    return f"bench_{model}"


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--models", default=",".join(MODELS), help="comma-separated Ultralytics model names")
    p.add_argument("--data", default=DATA)
    # Defaults mirror the original training run (see runs/detect/yolov8n_v1_train/args.yaml).
    p.add_argument("--epochs", type=int, default=35)
    p.add_argument("--batch", type=int, default=32)
    p.add_argument("--imgsz", type=int, default=640)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--device", default=None, help="e.g. '0' on Colab GPU; omit to auto-pick")
    p.add_argument("--skip-train", action="store_true", help="only evaluate existing runs/detect/bench_* weights")
    args = p.parse_args()

    models = [m.strip() for m in args.models.split(",")]
    rows: list[tuple[str, float, float, float, float, float, float]] = []

    for name in models:
        best = Path("runs/detect") / run_name(name) / "weights" / "best.pt"
        if not args.skip_train and not best.exists():
            print(f"\n=== Training {name} ===")
            YOLO(f"{name}.pt").train(
                data=args.data,
                epochs=args.epochs,
                batch=args.batch,
                imgsz=args.imgsz,
                seed=args.seed,
                device=args.device,
                name=run_name(name),
                exist_ok=True,
            )
        if not best.exists():
            print(f"skip {name}: {best} not found")
            continue

        print(f"\n=== Evaluating {name} ===")
        model = YOLO(best)
        metrics = model.val(data=args.data, split="val", imgsz=args.imgsz, device=args.device)
        n_params = sum(x.numel() for x in model.model.parameters()) / 1e6
        size_mb = best.stat().st_size / 1e6
        latency = metrics.speed["inference"]  # ms per image on this device
        rows.append((name, n_params, size_mb, metrics.box.map50, metrics.box.map, latency, 1000 / latency))

    lines = [
        "# Model-family benchmark (plan 3.1)",
        "",
        f"_Identical config per model: {args.epochs} epochs, batch {args.batch}, imgsz {args.imgsz},"
        f" seed {args.seed}, val split of `{args.data}`. Latency = Ultralytics val-time forward pass"
        " on the training device — comparable within this table only._",
        "",
        "| Model | Params (M) | Size (MB) | mAP@0.5 | mAP@0.5:0.95 | Latency (ms) | FPS |",
        "|---|---|---|---|---|---|---|",
    ]
    for name, params, size, map50, map5095, ms, fps in rows:
        lines.append(f"| {name} | {params:.1f} | {size:.1f} | {map50:.3f} | {map5095:.3f} | {ms:.1f} | {fps:.0f} |")
    lines += ["", "TODO(human): one paragraph — which model would you deploy, and why?", ""]

    out = Path("docs/benchmark_models.md")
    out.write_text("\n".join(lines))
    print(f"\nWrote {out}\n")
    print("\n".join(lines))


if __name__ == "__main__":
    main()