#!/usr/bin/env python3
"""NAM inference benchmark — measure latency at 48kHz with various model sizes.

Usage:
    python scripts/benchmark_nam.py              # all models, all block sizes
    python scripts/benchmark_nam.py --quick      # default model only, 256-block
    python scripts/benchmark_nam.py --model convnet --block 256

Results are extrapolated for RPi 4B (Cortex-A72) using a 4-8x slowdown factor
vs x86_64 PyTorch CPU.
"""

import argparse
import time
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

SAMPLE_RATE = 48000
RPI_SLOWDOWN_LOWER = 4.0  # Conservative RPi slowdown estimate
RPI_SLOWDOWN_UPPER = 8.0  # Pessimistic RPi slowdown estimate


def build_models():
    """Return list of (name, factory, params_estimate)."""
    from nam.models.conv_net import ConvNet

    models = [
        ("nano (4ch, 3 dil)", lambda: ConvNet(
            channels=4, dilations=[1, 2, 4],
            batchnorm=False, activation="Tanh",
        )),
        ("feather (8ch, 5 dil)", lambda: ConvNet(
            channels=8, dilations=[1, 2, 4, 8, 16],
            batchnorm=False, activation="Tanh",
        )),
        ("lite (12ch, 6 dil)", lambda: ConvNet(
            channels=12, dilations=[1, 2, 4, 8, 16, 32],
            batchnorm=False, activation="Tanh",
        )),
        ("standard (16ch, 8 dil)", lambda: ConvNet(
            channels=16, dilations=[1, 2, 4, 8, 16, 32, 64, 128],
            batchnorm=False, activation="Tanh",
        )),
        ("default (32ch, 4 dil)", lambda: ConvNet(
            channels=32, dilations=[1, 2, 4, 8],
            batchnorm=False, activation="Tanh",
        )),
    ]
    return models


def benchmark_model(model_fn, block_size, n_warmup=10, n_runs=500):
    """Run model on block_size samples, return average ms."""
    import torch

    model = model_fn()
    model.eval()
    with torch.no_grad():
        x = torch.randn(1, block_size)
        for _ in range(n_warmup):
            _ = model(x)

        start = time.perf_counter()
        for _ in range(n_runs):
            _ = model(x)
        elapsed = time.perf_counter() - start

    avg_ms = elapsed / n_runs * 1000.0
    params = sum(p.numel() for p in model.parameters())
    return avg_ms, params


def main():
    parser = argparse.ArgumentParser(description="NAM inference benchmark")
    parser.add_argument("--quick", action="store_true", help="Only default model, 256-block")
    parser.add_argument("--model", type=str, default=None, help="Model name substring to filter")
    parser.add_argument("--block", type=int, default=None, help="Block size (override)")
    args = parser.parse_args()

    models = build_models()
    block_sizes = [128, 256, 512]

    if args.quick:
        block_sizes = [256]
        models = [m for m in models if "default" in m[0]]

    if args.block:
        block_sizes = [args.block]

    if args.model:
        models = [m for m in models if args.model.lower() in m[0].lower()]

    print(f"{'Model':<30} {'Params':>8} {'Block':>6} {'x86 ms':>8} {'RPi ms (4x)':>12} {'RPi ms (8x)':>12} {'Budget%':>8}")
    print("-" * 90)

    for name, factory in models:
        for bs in block_sizes:
            block_ms = bs / SAMPLE_RATE * 1000.0
            avg_ms, params = benchmark_model(factory, bs)

            rpi_low = avg_ms * RPI_SLOWDOWN_LOWER
            rpi_high = avg_ms * RPI_SLOWDOWN_UPPER
            budget_pct = (avg_ms / block_ms) * 100

            ok = "✓" if rpi_low < block_ms else "⚠" if rpi_high < block_ms else "✗"

            print(f"{ok} {name:<27} {params:>8,} {bs:>6} {avg_ms:>7.3f}ms {rpi_low:>8.2f}ms {rpi_high:>8.2f}ms {budget_pct:>6.1f}%")

    print()
    print("Legend: ✓ = safe on RPi (4x est < block budget), ⚠ = marginal, ✗ = too slow")
    print(f"Block budget: {block_sizes[0] / SAMPLE_RATE * 1000:.2f}ms @ {SAMPLE_RATE}Hz")
    print(f"RPi estimate: {RPI_SLOWDOWN_LOWER:.0f}x-{RPI_SLOWDOWN_UPPER:.0f}x × x86 PyTorch CPU time")


if __name__ == "__main__":
    main()