# sweep.py
from __future__ import annotations
import time
from statistics import median
import numpy as np
import torch
import ssl_mamba

_CUT = 64600

def time_batch_sizes():
    print("device", torch.cuda.get_device_name())
    m = ssl_mamba.SSLMamba()
    m.load()
    cuda = getattr(m, "device", "cpu") == "cuda"
    results = {}
    print("=== batch-size sweep ===")
    for bs in [1, 2, 4, 8, 16, 24, 32]:
        audios = [np.random.randn(_CUT).astype(np.float32) for _ in range(bs)]
        srs = [16000] * bs
        try:
            m.score_batch(audios, srs)            # warm-up
            if cuda: torch.cuda.synchronize()
            reps = []
            for _ in range(5):
                if cuda: torch.cuda.synchronize()
                t0 = time.perf_counter()
                m.score_batch(audios, srs)
                if cuda: torch.cuda.synchronize()
                reps.append(time.perf_counter() - t0)
            med = median(reps); ups = bs / med
            results[bs] = ups
            print(f"bs={bs:3d}  median={med*1000:8.3f} ms  utt/s={ups:8.2f}")
        except torch.cuda.OutOfMemoryError:
            print(f"bs={bs:3d}  OOM (skipped)")
            if cuda: torch.cuda.empty_cache()
    m.unload()
    best = max(results, key=results.get)
    print(f"best batch_size = {best} (utt/s={results[best]:.2f})")
    return best, results

if __name__ == "__main__":
    time_batch_sizes()