#!/usr/bin/env python3
"""Audition official Silero v5_5_ru locally on CPU, without changing the reader."""
import argparse
import array
import hashlib
import json
import os
from pathlib import Path
import sys
import time
import wave

ROOT = Path(__file__).resolve().parent.parent
os.environ["CUDA_VISIBLE_DEVICES"] = ""
os.environ.setdefault("TORCH_HOME", str(ROOT / ".local/torch"))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--voices", nargs="+", default=["xenia", "baya", "aidar"])
    args = parser.parse_args()
    import torch

    torch.set_num_threads(4)
    torch.set_num_interop_threads(1)
    torch.manual_seed(42)
    model_path = ROOT / ".local/tts/silero-v5_5-ru/model.pt"
    output = ROOT / "composeApp/build/tts-samples"
    output.mkdir(parents=True, exist_ok=True)
    start = time.perf_counter()
    model = torch.package.PackageImporter(str(model_path)).load_pickle("tts_models", "model")
    model.to(torch.device("cpu"))
    loaded = time.perf_counter() - start
    print(f"Silero v5_5_ru, CPU, 4 threads; load {loaded:.3f}s; voices {model.speakers}", flush=True)
    samples = {
        "ru": "В библиотеке было тихо. Она открыла книгу и начала читать. За окном шёл осенний дождь. На двери висел старый замок, а вдали виднелся замок на холме. Завтра, подумала она, мы наконец доберёмся до моря.",
        "ru-yo": "Елка росла у дома. Ежик шел по лесу. Ребенок пьет теплое молоко. Все уже пришли. Он прочитал все. Она все еще ждет. Он узнает об этом завтра. Он узнает ее по голосу.",
    }
    report = {"model": "v5_5_ru", "sha256": hashlib.sha256(model_path.read_bytes()).hexdigest(),
              "torch": torch.__version__, "device": "cpu", "threads": 4, "load_seconds": loaded,
              "put_accent": True, "put_yo": True, "put_stress_homo": True, "put_yo_homo": True, "samples": []}
    for voice in args.voices:
        for label, text in samples.items():
            start = time.perf_counter()
            with torch.inference_mode():
                audio = model.apply_tts(text=text, speaker=voice, sample_rate=48000,
                                        put_accent=True, put_yo=True,
                                        put_stress_homo=True, put_yo_homo=True).detach().cpu().flatten()
            elapsed = time.perf_counter() - start
            assert audio.numel() and torch.isfinite(audio).all() and audio.abs().max() > 0.001
            duration = audio.numel() / 48000
            pcm = array.array("h", (audio.clamp(-1, 1) * 32767).round().to(torch.int16).tolist())
            if sys.byteorder != "little":
                pcm.byteswap()
            path = output / f"silero-{label}-{voice}.wav"
            with wave.open(str(path), "wb") as wav:
                wav.setparams((1, 2, 48000, 0, "NONE", "not compressed"))
                wav.writeframes(pcm.tobytes())
            path.with_suffix(".txt").write_text(text + "\n")
            report["samples"].append({"file": path.name, "synthesis_seconds": elapsed,
                                      "audio_seconds": duration, "rtf": elapsed / duration})
            print(f"{path.name}: {elapsed:.3f}s synthesis, {duration:.3f}s audio, RTF {elapsed / duration:.3f}", flush=True)
    (output / "silero-benchmark.json").write_text(json.dumps(report, indent=2) + "\n")


if __name__ == "__main__":
    main()