#!/usr/bin/env python
"""Slovenian Gepard TTS inference: text -> NanoCodec tokens (GepardRunner) -> waveform
(UnfoldedCodecModel). Run with the gepard-train venv_infer (has gepard + NeMo codec).

  cd third_party/gepard-train
  CUDA_VISIBLE_DEVICES=0 ./venv_infer/bin/python /home/tex/repos/trainer_slo/scripts/gepard_tts.py
"""
import os
import sys

import soundfile as sf
import torch

from gepard.inference.runner import GepardRunner
from gepard.inference.codec_wrapper import UnfoldedCodecModel

CKPT = "/home/tex/repos/trainer_slo/outputs/gepard_sl_lora/final"
CODEC = "nvidia/nemo-nano-codec-22khz-1.89kbps-21.5fps"
OUT = "/tmp/gepard_sl"
TEXTS = [
    "Danes je lep sončen dan in ptice pojejo na drevesih.",
    "Trtna uš je vinograde povsem uničila, ostalo je le nekaj brajdovca.",
    "Umetna inteligenca spreminja svet okoli nas.",
    "Sestanek se začne ob devetih zjutraj v veliki dvorani.",
]


def main():
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    os.makedirs(OUT, exist_ok=True)
    print("loading gepard checkpoint + nano-codec ...", file=sys.stderr)
    runner = GepardRunner.from_checkpoint(CKPT, device=dev)
    codec = UnfoldedCodecModel.from_pretrained(CODEC).eval().to(dev)
    for i, t in enumerate(TEXTS):
        tokens = runner.generate("sl: " + t, temperature=0.4, top_k=0,
                                 cfg_scale=4.5, cfg_frames=25, max_frames=800,
                                 repetition_penalty=1.45, repetition_window=32)  # cfg4.0=čšž; rep_pen=1.45 fixes non-termination
        codes = tokens.unsqueeze(0).to(dev)                      # (num_heads, T) -> (1, D, T)
        clen = torch.tensor([codes.shape[-1]], device=dev)
        with torch.inference_mode():
            audio, alen = codec.decode_from_codes(codes, clen)
        wav = audio[0, : int(alen[0])].detach().cpu().numpy()
        sf.write(f"{OUT}/{i}.wav", wav, 22050)
        print(f"[{i}] {len(wav)/22050:.1f}s  '{t[:50]}'")
    print("wavs ->", OUT)


if __name__ == "__main__":
    main()