# -*- coding: utf-8 -*-
"""Chatterbox TTS kiyaslama araci (sunucudan bagimsiz, kendi modelini yukler).

Kullanim (venv ile):
  tools\\.venv-tts\\Scripts\\python.exe tools\\tts_bench.py --mode base
  tools\\.venv-tts\\Scripts\\python.exe tools\\tts_bench.py --mode profile
  tools\\.venv-tts\\Scripts\\python.exe tools\\tts_bench.py --mode tf32

Cikti: T3 (otoregresif) ve S3Gen (vokoder) sureleri, token/sn, RTF.
"""
import argparse
import os
import sys
import time
from collections import defaultdict

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
HF_HOME = os.path.join(ROOT, "models", "hf-cache")
os.environ["HF_HOME"] = HF_HOME
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"

import torch  # noqa: E402
from chatterbox.mtl_tts import ChatterboxMultilingualTTS  # noqa: E402

TEXT = ("Merhaba, ben FIKO. Bug\u00fcn sana d\u00f6rt konuda yard\u0131mc\u0131 olabilirim. "
        "\u00d6nce takvimini kontrol edebilirim.")
REF = os.path.join(ROOT, "models", "voices", "referans.wav")


def sync():
    torch.cuda.synchronize()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--mode", default="base", choices=["base", "profile", "tf32"])
    ap.add_argument("--runs", type=int, default=3)
    ap.add_argument("--cfg", type=float, default=0.3)
    ap.add_argument("--exaggeration", type=float, default=0.6)
    args = ap.parse_args()

    if args.mode == "tf32":
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
        torch.set_float32_matmul_precision("high")

    print("[bench] mod=%s torch=%s gpu=%s" % (args.mode, torch.__version__, torch.cuda.get_device_name(0)), flush=True)
    model = ChatterboxMultilingualTTS.from_pretrained("cuda")
    if os.path.isfile(REF):
        model.prepare_conditionals(REF, exaggeration=args.exaggeration)

    acc = defaultdict(float)

    def wrap(obj, name, key, count_tokens=False):
        orig = getattr(obj, name)

        def f(*a, **k):
            sync()
            t = time.perf_counter()
            r = orig(*a, **k)
            sync()
            acc[key] += time.perf_counter() - t
            if count_tokens:
                acc["tokens"] += r.shape[-1]
            return r
        setattr(obj, name, f)

    wrap(model.t3, "inference", "t3", count_tokens=True)
    wrap(model.s3gen, "inference", "s3gen")

    def one():
        torch.manual_seed(1234)
        with torch.inference_mode():
            return model.generate(TEXT, language_id="tr", exaggeration=args.exaggeration,
                                  cfg_weight=args.cfg, temperature=0.8)

    print("[bench] isinma...", flush=True)
    one()
    acc.clear()

    if args.mode == "profile":
        from torch.profiler import profile, ProfilerActivity
        sync()
        t0 = time.perf_counter()
        with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
            wav = one()
            sync()
        wall = time.perf_counter() - t0
        gpu_busy = 0.0
        for e in prof.key_averages():
            gpu_busy += getattr(e, "self_device_time_total", 0.0) or getattr(e, "self_cuda_time_total", 0.0)
        gpu_busy /= 1e6
        n_launch = sum(e.count for e in prof.key_averages() if (getattr(e, "self_device_time_total", 0) or 0) > 0)
        ses = wav.shape[-1] / model.sr
        print("[profile] wall=%.2fs  gpu_kernel_toplam=%.2fs  (gpu_dolu=%%%.0f)  kernel_sayisi=%d  tokens=%d  ses=%.2fs"
              % (wall, gpu_busy, 100.0 * gpu_busy / wall, n_launch, acc["tokens"], ses), flush=True)
        print(prof.key_averages().table(sort_by="self_cpu_time_total", row_limit=12), flush=True)
        return

    rows = []
    for i in range(args.runs):
        acc.clear()
        sync()
        t0 = time.perf_counter()
        wav = one()
        sync()
        wall = time.perf_counter() - t0
        ses = wav.shape[-1] / model.sr
        tok = acc["tokens"]
        rows.append((wall, acc["t3"], acc["s3gen"], tok, ses))
        print("[run %d] wall=%.2fs  t3=%.2fs (%.1f tok/sn, %d token)  s3gen=%.2fs  ses=%.2fs  RTF=%.2f"
              % (i + 1, wall, acc["t3"], tok / max(acc["t3"], 1e-9), tok, acc["s3gen"], ses, wall / ses), flush=True)
    n = len(rows)
    avg = [sum(r[j] for r in rows) / n for j in range(5)]
    print("[ORTALAMA] wall=%.2fs t3=%.2fs s3gen=%.2fs tok/sn=%.1f RTF=%.2f vram_tepe=%.1fGB"
          % (avg[0], avg[1], avg[2], avg[3] / max(avg[1], 1e-9), avg[0] / avg[4],
             torch.cuda.max_memory_allocated() / 1e9), flush=True)


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