#!/usr/bin/env python3
"""Benchmark a Coral Edge TPU model against the CPU TFLite model."""

from pathlib import Path
import sys
import time

import cv2
import numpy as np
from tflite_runtime.interpreter import Interpreter, load_delegate


ROOT = Path(__file__).resolve().parent
EDGE_MODEL = ROOT / "model.tflite"
CPU_MODEL = ROOT / "model_cpu.tflite"
IMAGE = ROOT / "parrot.jpg"
WARMUP_RUNS = 10
TIMED_RUNS = 100


def benchmark(model_path: Path, *, delegate=None, threads=None):
    kwargs = {"model_path": str(model_path)}
    if delegate is not None:
        kwargs["experimental_delegates"] = [delegate]
    if threads is not None:
        kwargs["num_threads"] = threads

    interpreter = Interpreter(**kwargs)
    interpreter.allocate_tensors()
    input_info = interpreter.get_input_details()[0]

    image = cv2.imread(str(IMAGE))
    if image is None:
        raise RuntimeError(f"Cannot read test image: {IMAGE}")
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    height, width = input_info["shape"][1:3]
    tensor = cv2.resize(image, (int(width), int(height))).astype(input_info["dtype"])
    interpreter.set_tensor(input_info["index"], tensor[None, ...])

    for _ in range(WARMUP_RUNS):
        interpreter.invoke()

    times_ms = []
    for _ in range(TIMED_RUNS):
        start = time.perf_counter_ns()
        interpreter.invoke()
        times_ms.append((time.perf_counter_ns() - start) / 1_000_000)

    return {
        "mean": float(np.mean(times_ms)),
        "median": float(np.median(times_ms)),
        "p95": float(np.percentile(times_ms, 95)),
    }


def print_result(name, result):
    print(
        f"{name}: mean={result['mean']:.3f} ms, "
        f"median={result['median']:.3f} ms, "
        f"p95={result['p95']:.3f} ms, "
        f"fps={1000 / result['mean']:.2f}"
    )


def main():
    missing = [path.name for path in (EDGE_MODEL, CPU_MODEL, IMAGE) if not path.is_file()]
    if missing:
        print(f"Missing benchmark assets in {ROOT}: {', '.join(missing)}", file=sys.stderr)
        return 1

    edge = benchmark(EDGE_MODEL, delegate=load_delegate("libedgetpu.so.1.0"))
    cpu = benchmark(CPU_MODEL, threads=8)

    print(f"=== Coral benchmark: {WARMUP_RUNS} warm-up + {TIMED_RUNS} timed inferences ===")
    print_result("Edge TPU", edge)
    print_result("CPU (8 threads)", cpu)
    print(f"Speed-up: {cpu['mean'] / edge['mean']:.2f}x")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())