"""CPU inference latency benchmark for the two CV backbones.

Motivation: the original model comparison claimed MobileNetV3 is
"deployment-friendly / faster". Training throughput did not confirm that, so
this script measures the metric that actually matters for the CPU-only
HuggingFace Space: single-image forward-pass latency. It also reports the
parameter count and on-disk model size.

Usage:
    python -m src.cv.benchmark [--runs 100] [--warmup 10]
"""

from __future__ import annotations

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

import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from src.config import MODELS_DIR  # noqa: E402
from src.cv.train import build_model  # noqa: E402


LATENCY_REPORT = MODELS_DIR / "cv_latency_report.json"

BACKBONES = ["resnet18", "mobilenet_v3_small"]
DEFAULT_NUM_CLASSES = 10


def _param_count_millions(model: torch.nn.Module) -> float:
    return sum(p.numel() for p in model.parameters()) / 1e6


def _model_size_mb(model: torch.nn.Module) -> float:
    """Serialized state_dict size in MB (proxy for the deployed artifact)."""
    import io

    buffer = io.BytesIO()
    torch.save(model.state_dict(), buffer)
    return buffer.getbuffer().nbytes / (1024 * 1024)


@torch.no_grad()
def benchmark_model(
    name: str,
    num_classes: int,
    runs: int,
    warmup: int,
) -> dict:
    torch.manual_seed(42)
    device = torch.device("cpu")
    model = build_model(name, num_classes).to(device).eval()
    dummy = torch.randn(1, 3, 224, 224, device=device)

    for _ in range(warmup):
        model(dummy)

    timings_ms: list[float] = []
    for _ in range(runs):
        start = time.perf_counter()
        model(dummy)
        timings_ms.append((time.perf_counter() - start) * 1000.0)

    arr = np.array(timings_ms)
    return {
        "model": name,
        "params_m": round(_param_count_millions(model), 3),
        "size_mb": round(_model_size_mb(model), 2),
        "latency_ms_mean": round(float(arr.mean()), 2),
        "latency_ms_p50": round(float(np.percentile(arr, 50)), 2),
        "latency_ms_p95": round(float(np.percentile(arr, 95)), 2),
        "runs": runs,
    }


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--runs", type=int, default=100)
    p.add_argument("--warmup", type=int, default=10)
    p.add_argument("--num-classes", type=int, default=DEFAULT_NUM_CLASSES)
    args = p.parse_args()

    torch.set_num_threads(max(1, torch.get_num_threads()))
    print(f"[cv.benchmark] CPU threads: {torch.get_num_threads()}")

    results = [
        benchmark_model(name, args.num_classes, args.runs, args.warmup)
        for name in BACKBONES
    ]

    for r in results:
        print(
            f"[cv.benchmark] {r['model']:>20}  params={r['params_m']:.2f}M  "
            f"size={r['size_mb']:.1f}MB  "
            f"latency_mean={r['latency_ms_mean']:.1f}ms  "
            f"p95={r['latency_ms_p95']:.1f}ms"
        )

    faster = min(results, key=lambda r: r["latency_ms_p95"])
    payload = {
        "device": "cpu",
        "threads": torch.get_num_threads(),
        "results": results,
        "fastest_p95": faster["model"],
    }
    LATENCY_REPORT.write_text(json.dumps(payload, indent=2))
    print(f"[cv.benchmark] wrote {LATENCY_REPORT}")


if __name__ == "__main__":
    main()