#!/usr/bin/env python3
"""CPU baseline: thread-capped MetricGAN enhance_batch (project .venv)."""
from __future__ import annotations

import argparse
import os
import statistics
import time

os.environ.setdefault("OMP_NUM_THREADS", "1")
os.environ.setdefault("MKL_NUM_THREADS", "1")
os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")

import torch
from speechbrain.inference.enhancement import SpectralMaskEnhancement

SAVEDIR = "/root/.cache/speechbrain-enhancement"


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 main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--batch", type=int, default=14)
    ap.add_argument("--samples", type=int, default=16000)
    ap.add_argument("--iters", type=int, default=30)
    ap.add_argument("--warmup", type=int, default=5)
    args = ap.parse_args()

    torch.set_num_threads(1)
    torch.set_num_interop_threads(1)
    m = SpectralMaskEnhancement.from_hparams(
        source="speechbrain/metricgan-plus-voicebank",
        savedir=SAVEDIR,
        run_opts={"device": "cpu"},
    )
    wav = torch.randn(args.batch, args.samples)
    lengths = torch.ones(args.batch)
    with torch.no_grad():
        for _ in range(args.warmup):
            m.enhance_batch(wav, lengths=lengths)
        times = []
        for _ in range(args.iters):
            t0 = time.perf_counter()
            m.enhance_batch(wav, lengths=lengths)
            times.append((time.perf_counter() - t0) * 1000.0)
        feats = m.compute_features(wav)
        net = m.mods.enhance_model
        for _ in range(args.warmup):
            net(feats, lengths=lengths)
        mask_times = []
        for _ in range(args.iters):
            t0 = time.perf_counter()
            net(feats, lengths=lengths)
            mask_times.append((time.perf_counter() - t0) * 1000.0)

    def report(name, xs):
        per = [x / args.batch for x in xs]
        print(
            f"{name} batch={args.batch} n={len(xs)} "
            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"
        )

    report("cpu_enhance_batch", times)
    report("cpu_mask_only", mask_times)


if __name__ == "__main__":
    main()