"""Throughput of the local build against the padded-bank conv1d formulation.

    python dev/bench.py
"""
import os
import sys
import time

import torch
import torch.nn.functional as Fn

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
import load_local  # noqa: E402

rp = load_local.load()

PAIRS = [("44.1k -> 16k  (speech in)", 44100, 16000),
         ("48k   -> 16k  (speech in)", 48000, 16000),
         ("16k   -> 24k  (tts out)",   16000, 24000),
         ("44.1k -> 48k  (device)",    44100, 48000)]
SECONDS = 30.0
REPS = 25


def timeit(fn, reps=REPS):
    for _ in range(5):
        fn()
    best = 1e30
    for _ in range(reps):
        t0 = time.perf_counter()
        fn()
        best = min(best, time.perf_counter() - t0)
    return best


def ref(x, plan):
    L, M, width, kern = plan.L, plan.M, plan.width, plan.kernel
    B, T = x.shape
    pad = Fn.pad(x, (width, width + M))
    r = Fn.conv1d(pad[:, None], kern[:, None, :], stride=M)
    return r.transpose(1, 2).reshape(B, -1)[:, :plan.out_len(T)]


print(f"threads={torch.get_num_threads()}  audio={SECONDS:.0f}s mono\n")
print(f"{'conversion':28} {'resample-poly':>13} {'conv1d bank':>12} "
      f"{'speedup':>8} {'taps/out':>9} {'path':>6}")

for name, orig, new in PAIRS:
    T = int(orig * SECONDS)
    g = torch.Generator().manual_seed(0)
    x = torch.randn(1, T, generator=g, dtype=torch.float32)
    plan = rp.Resampler(orig, new)

    t_mine = timeit(lambda: plan(x))
    t_ref = timeit(lambda: ref(x, plan))
    a, b = plan(x), ref(x, plan)
    rel = ((a - b).abs().max() / b.abs().max()).item()
    assert rel < 2e-5, rel
    print(f"{name:28} {t_mine*1e3:12.2f}ms {t_ref*1e3:11.2f}ms "
          f"{t_ref/t_mine:7.2f}x {plan.taps_per_output:9.1f} "
          f"{'fused' if plan.fused else 'bank':>6}")

plan = rp.Resampler(44100, 16000)
x = torch.randn(1, int(44100 * SECONDS), dtype=torch.float32)
t = timeit(lambda: plan(x))
print(f"\n44.1k->16k realtime factor: {SECONDS / t:,.0f}x  "
      f"({t*1e3:.2f} ms for {SECONDS:.0f} s)")
print(f"taps stored {plan.taps.numel()} of {plan.kernel.numel()} "
      f"padded-bank entries ({plan.density*100:.1f}%)")