# -*- coding: utf-8 -*-
"""F5-TTS Qualitäts-/Tempo-Test (CPU): dt. Finetune + Lucy-Klon. Gleiche Sätze wie pocket -> A/B."""
import os, time, soundfile as sf, numpy as np, torch
# torchcodec/ffmpeg fehlt -> torchaudio.load/save auf soundfile umbiegen (wie bei OuteTTS-Patch)
import torchaudio
def _ta_load(path, *a, **k):
    data, sr = sf.read(str(path), dtype="float32", always_2d=True)
    return torch.from_numpy(data.T.copy()), sr
def _ta_save(path, tensor, sr, *a, **k):
    arr = np.asarray(tensor.detach().cpu().numpy())
    sf.write(str(path), arr.T if arr.ndim == 2 else arr, sr)
torchaudio.load = _ta_load
torchaudio.save = _ta_save
from f5_tts.api import F5TTS

BASE = r"F:\Coding Stuff\mission-control-2\client\lucy-f5"
OUT = r"C:\Users\TobisPC\Desktop\lucy_f5_test"; os.makedirs(OUT, exist_ok=True)
ref = os.path.join(BASE, "lucy_ref.wav")
ref_text = open(os.path.join(BASE, "lucy_ref.txt"), encoding="utf-8").read().strip()
print("REF_TEXT:", ref_text[:80], flush=True)

t0 = time.time()
f5 = F5TTS(model="F5TTS_Base",
           ckpt_file=os.path.join(BASE, "model_f5tts_german.safetensors"),
           vocab_file=os.path.join(BASE, "vocab.txt"), device="cpu")
print(f"Modell geladen in {time.time()-t0:.1f}s (Vocoder evtl. erst geladen)", flush=True)

SENT = {
 "kurz":  "Hallo Commander, ich höre dich.",
 "mittel":"Guten Morgen, Commander. Das Backup ist sauber durchgelaufen und es gab keine Fehler.",
 "lang":  "Natürlich kümmere ich mich darum, Commander. Ich starte den Dienst neu, prüfe die Protokolle und melde mich, sobald alles wieder läuft.",
}
for name, text in SENT.items():
    t0 = time.time()
    wav, sr, _ = f5.infer(ref_file=ref, ref_text=ref_text, gen_text=text,
                          nfe_step=32, target_rms=0.1, remove_silence=True)
    dt = time.time() - t0
    wav = np.asarray(wav, dtype=np.float32).reshape(-1)
    peak = float(np.max(np.abs(wav)))           # Peak-Limiter gegen Clipping (F5 traf 1.0)
    if peak > 0.95: wav = wav * (0.95 / peak)
    sf.write(os.path.join(OUT, f"{name}.wav"), wav, sr)
    secs = len(wav) / sr
    print(f"[{name:6}] gen={dt:6.1f}s audio={secs:5.1f}s RTF={dt/max(secs,0.01):5.2f}", flush=True)
print("F5_TEST_DONE", flush=True)