#!/usr/bin/env python3
"""Duelo honesto entre dos configuraciones de voz.

Por que existe esto: la primera comparacion uso UNA frase, y la segunda tres.
Con esos tamaños se concluyo que el ajuste nuevo mejoraba... y al regenerar la
cache entera salio peor. XTTS muestrea en cada generacion (do_sample, temperatura
0,65), asi que dos pasadas de la MISMA frase con la MISMA configuracion ya dan
numeros distintos. Comparar pocas muestras es leer ruido.

Aqui cada configuracion genera las mismas N frases R veces, en el mismo proceso
y el mismo dispositivo, y se comparan las medias. Ademas se mide la dispersion
DENTRO de cada configuracion: si la diferencia entre las dos es menor que lo que
varia una consigo misma, no hay diferencia que defender.

    python duelo_ajustes.py [cuda|cpu] [repeticiones]
"""
import json
import os
import statistics
import sys
import tempfile

import torch

DISPOSITIVO = sys.argv[1] if len(sys.argv) > 1 else "cuda"
REPES = int(sys.argv[2]) if len(sys.argv) > 2 else 3

BASE = os.path.dirname(os.path.abspath(__file__))
REF = os.path.join(BASE, "ref")
CATALOGO = os.path.join(BASE, "..", "newelle", "data", "acciones", "acciones.json")

os.environ.setdefault("COQUI_TOS_AGREED", "1")
LARGO = dict(gpt_cond_len=30, gpt_cond_chunk_len=6, max_ref_length=30)

VIEJA = dict(
    nombre="vieja (3 refs, cond 6s)",
    refs=[os.path.join(REF, n) for n in
          ("jarvis_lento_b.wav", "jarvis_lento_a.wav", "jarvis_ref.wav")],
    cond=dict(gpt_cond_len=6, gpt_cond_chunk_len=6, max_ref_length=30),
)
NUEVA = dict(
    nombre="nueva (1 ref sesion, cond 30s)",
    refs=[os.path.join(REF, "jarvis_sesion.wav")],
    cond=LARGO,
)
GEN = dict(temperature=0.65, repetition_penalty=2.0, speed=1.0)


def frases(n=8):
    with open(CATALOGO, encoding="utf-8") as f:
        d = json.load(f)
    acc = d.get("acciones", d if isinstance(d, list) else [])
    return sorted({a["acuse"] for a in acc if a.get("acuse")})[:n]


def main():
    from TTS.api import TTS
    import soundfile as sf

    api = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to(DISPOSITIVO)
    modelo = api.synthesizer.tts_model

    def emb(ruta):
        _, e = modelo.get_conditioning_latents(audio_path=[ruta], **LARGO)
        return e

    patron = emb(os.path.join(REF, "jarvis_lento_b.wav"))

    def parecido(ruta):
        return float(torch.nn.functional.cosine_similarity(
            patron.squeeze().float(), emb(ruta).squeeze().float(), dim=0))

    textos = frases()
    print(f"  {len(textos)} frases x {REPES} repeticiones por configuracion\n")

    resumen = {}
    with tempfile.TemporaryDirectory() as tmp:
        for cfg in (VIEJA, NUEVA):
            refs = [r for r in cfg["refs"] if os.path.exists(r)]
            gpt, hablante = modelo.get_conditioning_latents(
                audio_path=refs, **cfg["cond"])
            puntos = []
            for i, texto in enumerate(textos):
                for r in range(REPES):
                    out = modelo.inference(
                        text=texto, language="es", gpt_cond_latent=gpt,
                        speaker_embedding=hablante, enable_text_splitting=True,
                        **GEN)
                    d = os.path.join(tmp, f"{id(cfg)}_{i}_{r}.wav")
                    sf.write(d, out["wav"], 24000)
                    puntos.append(parecido(d))
            resumen[cfg["nombre"]] = puntos
            print(f"    {cfg['nombre']:32s} media {statistics.mean(puntos):.4f}"
                  f"  desv {statistics.pstdev(puntos):.4f}"
                  f"  ({len(puntos)} muestras)", flush=True)

    a, b = list(resumen.values())
    ma, mb = statistics.mean(a), statistics.mean(b)
    ruido = (statistics.pstdev(a) + statistics.pstdev(b)) / 2
    print(f"\n    diferencia entre configuraciones: {abs(mb-ma):.4f}")
    print(f"    variacion dentro de cada una:      {ruido:.4f}")
    if abs(mb - ma) < ruido:
        print("\n    => la diferencia es MENOR que el ruido: no hay mejora que defender")
    else:
        gana = list(resumen)[1] if mb > ma else list(resumen)[0]
        print(f"\n    => gana: {gana}")


if __name__ == "__main__":
    main()