# -*- coding: utf-8 -*-
"""sweep_b2.py — findet den B2-Sweet-Spot auf DIESER Maschine (Ryzen 9700X, 8C/16T).
Variiert WORKERS × THREADS-pro-Worker und misst je Konfig auf einer langen, mehrsätzigen Antwort:
  - wall   : Gesamt-Wall-Clock (Generierung der ganzen Antwort)
  - ttfb   : Time-to-first-audio (Reaktionszeit = erster Satz fertig)
  - audio  : erzeugte Audiolänge (variiert leicht, da temp>0 stochastisch)
  - RTF    : wall/audio (längen-normalisiert -> fairer Durchsatz-Vergleich; <1 = schneller als Echtzeit)
Nutzt die ECHTEN pocket_server-Funktionen (_worker_init/_warmup/_gen_sentences_ordered).
Schreibt sweep_b2_result.json + eine Tabelle und nennt am Ende den Sweet Spot.
"""
import os, time, json, statistics as st
import numpy as np, librosa
from concurrent.futures import ProcessPoolExecutor
from pocket_tts import TTSModel
import pocket_server as ps

BASE = ps.BASE
LONG = ("Guten Morgen, Commander. "
        "Das nächtliche Backup ist sauber durchgelaufen und es gab keine Fehler. "
        "Der Dienst läuft stabil und die Engine antwortet zügig. "
        "Ich habe die Modelle vorgewärmt und die Latenz im Blick behalten. "
        "Die Protokolle zeigen keine Auffälligkeiten in den letzten Stunden. "
        "Wenn du möchtest, fasse ich die offenen Punkte für heute zusammen. "
        "Danach kümmere ich mich um die anstehenden Updates und melde mich wieder.")

# (workers, threads_pro_worker). Ziel: workers*threads ~ 6–8 (8 physische Kerne).
CONFIGS = [("seriell", 0, 0), ("pool", 2, 4), ("pool", 2, 3),
           ("pool", 3, 2), ("pool", 3, 3), ("pool", 4, 2)]
REPS = 3


def _measure(sents):
    """Eine Durchführung: (wall, ttfb, audio_s) über den aktuell in STATE gesetzten Pfad."""
    t0 = time.time(); first = None; total = 0
    for i, a in enumerate(ps._gen_sentences_ordered(sents)):
        if i == 0:
            first = time.time() - t0
        total += a.size
    return time.time() - t0, first, total / 24000.0


def run_serial(reps):
    m = TTSModel.load_model(language=ps.LANG, lsd_decode_steps=ps.LSD, temp=ps.TEMP,
                            noise_clamp=ps.NOISE_CLAMP, quantize=ps.QUANTIZE)
    ref = ps._prep_ref()
    try:
        vs = m.get_state_for_audio_prompt(ps.VOICE_ST)
    except Exception:
        vs = m.get_state_for_audio_prompt(ref)
    ra, _ = librosa.load(ref, sr=m.sample_rate, mono=True)
    ps.STATE.clear()
    ps.STATE.update(m=m, vs=vs, sr=m.sample_rate, ref_fp=ps._fingerprint(ra, m.sample_rate))
    sents = ps._split_sentences(LONG)
    rows = [_measure(sents) for _ in range(reps)]
    ps.STATE.clear(); del m
    return rows, len(sents)


def run_pool(workers, threads, reps):
    os.environ["LUCY_WORKER_THREADS"] = str(threads)   # vom Worker beim Spawn gelesen
    pool = ProcessPoolExecutor(max_workers=workers, initializer=ps._worker_init)
    list(pool.map(ps._warmup, range(workers)))
    ref = ps._prep_ref(); ra, _ = librosa.load(ref, sr=24000, mono=True)
    ps.STATE.clear()
    ps.STATE.update(sr=24000, ref_fp=ps._fingerprint(ra, 24000), pool=pool)
    sents = ps._split_sentences(LONG)
    rows = [_measure(sents) for _ in range(reps)]
    pool.shutdown(wait=True); ps.STATE.clear()
    return rows, len(sents)


if __name__ == "__main__":
    results = []
    print(f"Text: {len(LONG)} Zeichen, REPS={REPS}\n")
    print(f"{'config':>14} | {'wall':>6} {'ttfb':>6} {'audio':>6} {'RTF':>5}")
    print("-" * 48)
    for kind, w, t in CONFIGS:
        try:
            rows, nsent = (run_serial(REPS) if kind == "seriell" else run_pool(w, t, REPS))
        except Exception as e:
            print(f"{kind} w{w} t{t}: FEHLER {e}")
            continue
        wall = st.median([r[0] for r in rows])
        ttfb = st.median([r[1] for r in rows])
        audio = st.median([r[2] for r in rows])
        rtf = wall / max(audio, 0.01)
        label = "seriell(1×all)" if kind == "seriell" else f"{w}w×{t}t"
        print(f"{label:>14} | {wall:6.2f} {ttfb:6.2f} {audio:6.2f} {rtf:5.2f}")
        results.append(dict(label=label, kind=kind, workers=(1 if kind == "seriell" else w),
                            threads=t, wall=wall, ttfb=ttfb, audio=audio, rtf=rtf, sents=nsent))

    json.dump(results, open(os.path.join(BASE, "sweep_b2_result.json"), "w"), indent=2)

    pool_rows = [r for r in results if r["kind"] == "pool"]
    if pool_rows:
        best_tput = min(pool_rows, key=lambda r: r["rtf"])
        best_ttfb = min(pool_rows, key=lambda r: r["ttfb"])
        # Sweet Spot: bester Durchsatz, aber TTFB nicht >20% über dem TTFB-Sieger (Reaktion zählt)
        cand = [r for r in pool_rows if r["ttfb"] <= best_ttfb["ttfb"] * 1.20]
        sweet = min(cand, key=lambda r: r["rtf"]) if cand else best_tput
        print("\n>> bester Durchsatz :", best_tput["label"], f"(RTF {best_tput['rtf']:.2f})")
        print(">> beste Reaktion   :", best_ttfb["label"], f"(TTFB {best_ttfb['ttfb']:.2f}s)")
        print(">> SWEET SPOT       :", sweet["label"],
              f"-> LUCY_WORKERS={sweet['workers']} LUCY_WORKER_THREADS={sweet['threads']}")
    print("SWEEP_DONE")