#!/usr/bin/env python
"""benchmark.py — reproduce the Audio8-TTS single-stream fast-path RTF numbers.

Measures real-time factor (RTF = autoregressive_seconds / audio_seconds; lower is
faster, and RTF < 1 is faster than real time) for a batch-1 single stream on one GPU:

    stock    stock model.generate (eager reference implementation)
    fast-ar  stock model.generate with only the fast-AR codebook step torch.compiled
    eager    FastArktts static-shape rewrite, eager
    compile  FastArktts whole-frame torch.compile(fullgraph=True, dynamic=False)
    maxat    FastArktts whole-frame torch.compile(mode="max-autotune")   [best]

Codec decode time is reported separately and NOT counted in RTF (the AR loop is the
throughput wall; decode is a one-shot ~40 ms).

    python benchmark.py                       # all modes, default text/voice
    python benchmark.py --modes eager compile maxat
    python benchmark.py --model scrappylabsai/warble --runs 5
"""
from __future__ import annotations

import argparse
import time

import torch

from fast_arktts import DEFAULT_MODEL_ID, FastArktts, build_prompt, load_model

FPS = 44100 / 2048  # codec frame rate ~= 21.53 frames / second

# long enough to generate ~150 frames so RTF is stable and comparable
DEFAULT_TEXT = (
    "<|speaker:2|>Thanks for calling. I can help you with that. Let me pull up your "
    "account real quick and get you sorted out today."
)


def rtf_of(seconds, frames):
    return seconds / (frames / FPS)


def time_runs(call, warmup, runs, label):
    for _ in range(warmup):
        call()
    out = []
    for _ in range(runs):
        torch.cuda.synchronize()
        t = time.perf_counter()
        frames = call()
        dt = time.perf_counter() - t
        out.append((dt, frames, rtf_of(dt, frames)))
    best = min(out, key=lambda r: r[2])
    med = sorted(out, key=lambda r: r[2])[len(out) // 2]
    print(f"  {label:<9} best AR={best[0]:.3f}s frames={best[1]} "
          f"RTF={best[2]:.3f} ({1/best[2]:.1f}x realtime) | median RTF={med[2]:.3f}",
          flush=True)
    return best, med


# ---- mode: stock model.generate (optionally with the fast-AR step compiled) ----

def stock_caller(model, processor, text, max_new, compile_fast_ar=False):
    inp = processor(text=[text], return_tensors="pt")
    inp = {k: v.to(model.device) for k, v in inp.items()}
    if compile_fast_ar:
        # compile ONLY the fast-AR codebook generator; the slow AR + sampling stay eager
        model._generate_codebooks = torch.compile(model._generate_codebooks)

    def call():
        codes = model.generate(**inp, max_new_tokens=max_new, temperature=0.7,
                               top_p=0.9, top_k=50, do_sample=True)
        return codes.shape[-1]

    return call


# ---- mode: FastArktts ----

def fast_caller(model, prompt, mask, mode, max_new):
    fast = FastArktts(model, max_new_tokens=max_new)
    if mode == "compile":
        fast.compile("default")
    elif mode == "maxat":
        fast.compile("max-autotune")
    elif mode == "graph":
        fast.capture_graph(prompt, mask)
    # "eager" -> leave as-is

    def call():
        return fast.generate(prompt, mask).shape[-1]

    return call


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default=DEFAULT_MODEL_ID)
    ap.add_argument("--text", default=DEFAULT_TEXT)
    ap.add_argument("--modes", nargs="+",
                    default=["stock", "fast-ar", "eager", "compile", "maxat"])
    ap.add_argument("--warmup", type=int, default=3)
    ap.add_argument("--runs", type=int, default=3)
    ap.add_argument("--max-new", type=int, default=400)
    args = ap.parse_args()

    if not torch.cuda.is_available():
        raise SystemExit("benchmark.py requires a CUDA GPU")
    torch._dynamo.config.cache_size_limit = 64

    t0 = time.time()
    model, processor = load_model(args.model)
    print(f"loaded {args.model} in {time.time()-t0:.1f}s "
          f"| VRAM {torch.cuda.memory_allocated()/1e9:.2f} GB "
          f"| {torch.cuda.get_device_name(0)}", flush=True)

    prompt, mask = build_prompt(model, processor, args.text)
    print(f"prompt_width={prompt.shape[-1]}  fps={FPS:.2f}\n", flush=True)
    print("Results (RTF = AR seconds / audio seconds; lower is faster):", flush=True)

    for mode in args.modes:
        if mode in ("stock", "fast-ar"):
            call = stock_caller(model, processor, args.text, args.max_new,
                                compile_fast_ar=(mode == "fast-ar"))
        else:
            call = fast_caller(model, prompt, mask, mode, args.max_new)
        time_runs(call, args.warmup, args.runs, mode)

    # codec decode cost (reported, not counted in RTF) — eager generate is fine here
    gen = FastArktts(model, max_new_tokens=args.max_new)
    c = gen.generate(prompt, mask)
    torch.cuda.synchronize()
    t = time.perf_counter()
    model.decode_audio(c)
    torch.cuda.synchronize()
    print(f"\ncodec decode = {(time.perf_counter()-t)*1000:.0f} ms (one-shot, not in RTF)")
    print(f"VRAM peak = {torch.cuda.max_memory_allocated()/1e9:.2f} GB")


if __name__ == "__main__":
    main()