#!/usr/bin/env python3
"""RTF bench for hop/window speechbrain enhance.

Stop livekit-cameras.service first or the numbers fight the fleet.

  .venv/bin/python scripts/bench_enhance.py
  .venv/bin/python scripts/bench_enhance.py --stereo --hop 0.1
"""
import argparse
import sys
import time
from pathlib import Path

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

import numpy as np

from audio_cleanup import AudioCleaner


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--stereo", action="store_true")
    ap.add_argument("--hop", type=float, default=0.2)
    ap.add_argument("--chunk", type=float, default=1.0)
    ap.add_argument("--n", type=int, default=50)
    args = ap.parse_args()

    c = AudioCleaner(
        mode="force",
        beamform="force" if args.stereo else "off",
        chunk_s=args.chunk,
        hop_s=args.hop,
        torch_num_threads=1,
    )
    hop_n = int(args.hop * 16000)
    if args.stereo:
        t = np.arange(hop_n)
        ch0 = (np.sin(2 * np.pi * 440 * t / 16000) * 800).astype(np.int16)
        pcm = np.stack([ch0, np.roll(ch0, 2)], axis=1)
    else:
        pcm = (np.random.randn(hop_n) * 800).astype(np.int16)
    warmup = max(8, int(args.chunk / args.hop) + 2)
    for _ in range(warmup):
        c.process(pcm)
    n = args.n
    t0 = time.perf_counter()
    for _ in range(n):
        c.process(pcm)
    dt = time.perf_counter() - t0
    audio_s = n * args.hop
    rtf = dt / audio_s
    print(
        f"backend={c.backend} stereo={args.stereo} hops={n} "
        f"wall={dt:.3f}s audio={audio_s:.1f}s rtf={rtf:.3f} "
        f"last_ms={c.last_dt_ms:.1f}"
    )
    limit = 0.5 if args.stereo else 0.8
    if rtf > limit:
        raise SystemExit(f"RTF too high for 20-way realtime (limit {limit})")
    return 0


if __name__ == "__main__":
    sys.exit(main())