#!/usr/bin/env python3
"""OpenVINO GPU/CPU bench of MetricGAN mask ONNX (.venv-ov)."""
from __future__ import annotations

import argparse
import statistics
import time
from pathlib import Path

import numpy as np
import openvino as ov


def pct(xs, p):
    xs = sorted(xs)
    i = int(round((p / 100.0) * (len(xs) - 1)))
    return xs[max(0, min(i, len(xs) - 1))]


def bench(compiled, feats, warmup, iters):
    infer = compiled.create_infer_request()
    inp = compiled.inputs[0]
    for _ in range(warmup):
        infer.infer({inp: feats})
    xs = []
    for _ in range(iters):
        t0 = time.perf_counter()
        infer.infer({inp: feats})
        xs.append((time.perf_counter() - t0) * 1000.0)
    return xs


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--onnx", default="scripts/metricgan_enhance.onnx")
    ap.add_argument("--batch", type=int, default=14)
    ap.add_argument("--time", type=int, default=63)
    ap.add_argument("--feat", type=int, default=257)
    ap.add_argument("--iters", type=int, default=40)
    ap.add_argument("--warmup", type=int, default=10)
    args = ap.parse_args()

    core = ov.Core()
    print("openvino", ov.__version__)
    print("devices", core.available_devices)
    for d in core.available_devices:
        try:
            print(" ", d, core.get_property(d, "FULL_DEVICE_NAME"))
        except Exception as e:
            print(" ", d, e)

    onnx = Path(args.onnx)
    model = ov.convert_model(str(onnx))
    feats = np.random.randn(args.batch, args.time, args.feat).astype(np.float32)

    for device in core.available_devices:
        try:
            compiled = core.compile_model(model, device)
        except Exception as e:
            print(f"compile {device} FAILED: {type(e).__name__}: {e}")
            continue
        try:
            xs = bench(compiled, feats, args.warmup, args.iters)
        except Exception as e:
            print(f"infer {device} FAILED: {type(e).__name__}: {e}")
            continue
        per = [x / args.batch for x in xs]
        print(
            f"ov_{device} batch={args.batch} "
            f"p50={statistics.median(xs):.2f}ms p95={pct(xs,95):.2f}ms "
            f"p99={pct(xs,99):.2f}ms "
            f"per_cam_p50={statistics.median(per):.2f}ms "
            f"per_cam_p95={pct(per,95):.2f}ms"
        )


if __name__ == "__main__":
    main()