# -*- coding: utf-8 -*-
"""bench_lucy.py — A3/B1: misst Lucys Stimme (pocket-tts) über temp/lsd.

Metriken je Konfig: RTF, gen-Zeit und ROHE Kollaps-Rate (EINE Generierung, OHNE best-of-N-Gate)
— also genau das, was der Stimm-Wächter aktuell wegbügeln muss (= Latenz-Treiber). Niedrigere
temp sollte die rohe Kollaps-Rate senken -> weniger Regenerationen -> niedrigere effektive Latenz.

Nutzung (im ptts-venv der lucy-tts-Maschine):
  python bench_lucy.py                          # Default-Sweep
  python bench_lucy.py --reps 5
  python bench_lucy.py --temps 0.7,0.8,0.9 --lsds 8,10,12
  python bench_lucy.py --whisper                # + Verständlichkeits-Check (faster-whisper)

Reuse: Kollaps-Logik (F0 + MFCC-Fingerabdruck) und cleanup() kommen 1:1 aus pocket_server.py,
damit der Benchmark exakt die Produktions-Kriterien misst.
"""
import os, time, argparse, statistics as st
import numpy as np, soundfile as sf, librosa
from pocket_tts import TTSModel
import pocket_server as ps

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.",
}


def _voice_state(m):
    """Voice-State laden: bevorzugt das gecachte safetensors (A1), sonst aus der Referenz klonen."""
    if os.path.exists(ps.VOICE_ST):
        try:
            return m.get_state_for_audio_prompt(ps.VOICE_ST)
        except Exception as e:
            print(f"  (safetensors-Load fehlgeschlagen, klone aus Referenz: {e})")
    return m.get_state_for_audio_prompt(ps._prep_ref())


def _raw_gen(m, vs, ref_fp, text, sr):
    """EINE rohe Generierung (kein Gate). -> (gen_s, audio_s, male, sim, audio)."""
    t0 = time.time()
    audio = m.generate_audio(vs, ps.LEAD + text, frames_after_eos=ps.FEOS)
    gen = time.time() - t0
    a = audio.numpy() if hasattr(audio, "numpy") else np.asarray(audio)
    a = np.asarray(a, dtype=np.float32).reshape(-1)
    f0 = ps._voiced_f0(a, sr)
    male = (f0 == f0) and f0 < ps.F0_FLOOR              # NaN-sicher
    fp = ps._fingerprint(ps._crop_lead(a, sr), sr)
    sim = float(np.dot(ref_fp, fp)) if (ref_fp is not None and fp is not None) else 1.0
    return gen, a.size / sr, bool(male), sim, a


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--temps", default="0.7,0.8,0.9")
    ap.add_argument("--lsds", default="8,10")
    ap.add_argument("--reps", type=int, default=4)
    ap.add_argument("--whisper", action="store_true")
    args = ap.parse_args()
    temps = [float(x) for x in args.temps.split(",")]
    lsds = [int(x) for x in args.lsds.split(",")]

    OUT = os.path.join(ps.BASE, "out_bench"); os.makedirs(OUT, exist_ok=True)
    ref = ps._prep_ref()
    print(f"LANG={ps.LANG} quantize={ps.QUANTIZE} noise_clamp={ps.NOISE_CLAMP} "
          f"F0_FLOOR={ps.F0_FLOOR} FP_FLOOR={ps.FP_FLOOR} reps={args.reps}")
    print(f"{'temp':>5} {'lsd':>4} | {'RTF':>5} {'gen/s':>6} | "
          f"{'collapse%':>9} {'sim_min':>7} {'sim_avg':>7}")
    rows = []
    for lsd in lsds:
        for temp in temps:
            m = TTSModel.load_model(language=ps.LANG, lsd_decode_steps=lsd, temp=temp,
                                    noise_clamp=ps.NOISE_CLAMP, quantize=ps.QUANTIZE)
            sr = m.sample_rate
            ref_audio, _ = librosa.load(ref, sr=sr, mono=True)
            ref_fp = ps._fingerprint(ref_audio, sr)
            vs = _voice_state(m)
            rtfs, gens, sims, collapses, n = [], [], [], 0, 0
            for name, text in SENT.items():
                for r in range(args.reps):
                    gen, asec, male, sim, a = _raw_gen(m, vs, ref_fp, text, sr)
                    rtfs.append(gen / max(asec, 0.01)); gens.append(gen); sims.append(sim)
                    collapses += int(male or sim < ps.FP_FLOOR); n += 1
                    if r == 0:  # ein Sample je Satz zum Reinhören
                        sf.write(os.path.join(OUT, f"t{temp}_l{lsd}_{name}.wav"), ps.cleanup(a, sr), sr)
            cr = 100.0 * collapses / max(n, 1)
            print(f"{temp:>5} {lsd:>4} | {st.mean(rtfs):>5.2f} {st.mean(gens):>6.2f} | "
                  f"{cr:>8.1f}% {min(sims):>7.3f} {st.mean(sims):>7.3f}")
            rows.append((temp, lsd, st.mean(rtfs), cr, min(sims)))
            del m

    best = sorted(rows, key=lambda x: (x[3], x[2]))[0]  # min Kollaps, dann beste RTF
    print(f"\n>> Vorschlag: temp={best[0]} lsd={best[1]} "
          f"(collapse={best[3]:.1f}%, RTF={best[2]:.2f}) -> in pocket_server via "
          f"LUCY_TEMP / LUCY_LSD setzen.")

    if args.whisper:
        print("\n=== Whisper-Verständlichkeit (faster-whisper small/int8) ===")
        from faster_whisper import WhisperModel
        import glob
        wm = WhisperModel("small", device="cpu", compute_type="int8")
        for p in sorted(glob.glob(os.path.join(OUT, "*.wav"))):
            segs, _ = wm.transcribe(p, language="de", beam_size=5)
            print(os.path.basename(p), "::", " ".join(s.text.strip() for s in segs))


if __name__ == "__main__":
    main()