#!/usr/bin/env python3
"""Run several synthesis cases on the board and save wavs (for ASR / listening)."""
import os, sys, time
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from cutetts_ax import CuteTTSAX
from cutetts_ax import audio as A

MODEL = os.environ.get("CUTETTS_MODEL_DIR", "compile")
ASSETS = os.environ.get("CUTETTS_ASSETS", "sdk/assets")
OUT = os.environ.get("CUTETTS_OUT", "scratch/board_samples")
os.makedirs(OUT, exist_ok=True)

CASES = [
    ("en_default", "The voice generated by this model sounds amazing!",
     "sdk/samples/reference.wav", 42),
    ("en_long", "Hello, this is a test of the CuteTTS text to speech system running on the AX650 neural processing unit.",
     "sdk/samples/reference.wav", 7),
    ("zh_default", "今天天气很好，我们一起去公园散步吧。",
     "sdk/samples/reference.wav", 123),
    ("en_other_voice", "Quick brown fox jumps over the lazy dog.",
     "/ai_data/AX615/hy/workspace/Speech-Translation-ZipVoice/samples/en_16k.wav", 2024),
]

tts = CuteTTSAX(MODEL, ASSETS, engine=os.environ.get("CUTETTS_ENGINE", "ax"))
for tag, text, ref, seed in CASES:
    if not os.path.exists(ref):
        print("skip (no ref)", tag); continue
    t0 = time.time()
    wav = tts.synthesize(text, ref, seed=seed)
    A.write_wav(os.path.join(OUT, tag + ".wav"), wav, tts.sr)
    tm = tts.last_timing
    print(f"{tag}: {tm['audio_seconds']:.2f}s audio, {tm['steps']} patches, "
          f"total {tm['total']:.2f}s, RTF {tm['total']/max(tm['audio_seconds'],1e-6):.2f} | "
          f"vae_enc {tm['vae_enc']:.2f} spk {tm['speaker_enc']:.2f} prefill {tm['prefill']:.2f} "
          f"ar {tm['ar_loop']:.2f} (dit {tm['dit']:.2f} dec {tm['lm_decode']:.2f}) vae_dec {tm['vae_dec']:.2f}")