#!/usr/bin/env python3
"""Mide el ajuste ACTUAL contra el propuesto, con las mismas frases.

Las dos vueltas anteriores compararon cosas distintas entre si; esto compara lo
que hay hoy en la cache con lo que se quiere poner, para poder decir una cifra
honesta de cuanto mejora.

De paso deja creada la referencia buena: ref/jarvis_sesion.wav, las tomas de
una misma sesion pegadas. Mezclar sesiones distintas puntuaba peor —el
promediado de XTTS diluye el timbre cuando los clips no suenan igual—, que es
justo lo contrario de lo que se supuso al montarlo.

    python antes_despues.py [cuda|cpu]
"""
import glob
import os
import subprocess
import sys
import tempfile

import torch

DISPOSITIVO = sys.argv[1] if len(sys.argv) > 1 else "cuda"
BASE = os.path.dirname(os.path.abspath(__file__))
REF = os.path.join(BASE, "ref")
SALIDA = os.path.join(BASE, "comparativa")
DESCARGAS = os.path.expanduser("~/Downloads")
SESION = os.path.join(REF, "jarvis_sesion.wav")

os.environ.setdefault("COQUI_TOS_AGREED", "1")

FRASES = [
    "Buenas noches, señor. Soy su asistente personal.",
    "Un momento, lo miro.",
    "He encontrado catorce ficheros, señor. Dígame cuál abro.",
]

# lo que hay hoy en generar_cache_voz.py
ANTES = dict(
    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),
    gen=dict(temperature=0.65, repetition_penalty=2.0, speed=1.0),
)
# lo propuesto
DESPUES = dict(
    refs=[SESION],
    cond=dict(gpt_cond_len=30, gpt_cond_chunk_len=6, max_ref_length=30),
    gen=dict(temperature=0.65, repetition_penalty=2.0, speed=1.0),
)


def parecido(a, b):
    return float(torch.nn.functional.cosine_similarity(
        a.squeeze().float(), b.squeeze().float(), dim=0))


def construye_sesion():
    """Pega las tomas de la misma sesion en un unico wav de referencia."""
    if os.path.exists(SESION):
        return SESION
    mp3s = sorted(glob.glob(os.path.join(DESCARGAS, "Jarvis-*Buenas-noches*.mp3")))
    if not mp3s:
        raise SystemExit("no encuentro los mp3 de la sesion en ~/Downloads")
    with tempfile.TemporaryDirectory() as tmp:
        trozos = []
        for i, m in enumerate(mp3s):
            d = os.path.join(tmp, f"{i}.wav")
            subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-i", m,
                            "-ar", "24000", "-ac", "1", d], check=True)
            trozos.append(d)
        lista = os.path.join(tmp, "lista.txt")
        with open(lista, "w") as f:
            for t in trozos:
                f.write(f"file '{t}'\n")
        subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-f", "concat",
                        "-safe", "0", "-i", lista, "-ar", "24000", "-ac", "1",
                        SESION], check=True)
    print(f"  referencia nueva: {SESION}")
    return SESION


def main():
    os.makedirs(SALIDA, exist_ok=True)
    construye_sesion()

    from TTS.api import TTS
    import soundfile as sf
    print(f"cargando XTTS en {DISPOSITIVO}...", flush=True)
    api = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to(DISPOSITIVO)
    modelo = api.synthesizer.tts_model

    largo = dict(gpt_cond_len=30, gpt_cond_chunk_len=6, max_ref_length=30)

    def emb(rutas):
        _, e = modelo.get_conditioning_latents(audio_path=list(rutas), **largo)
        return e

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

    print()
    for etiqueta, cfg in (("ANTES  (lo que suena hoy)", ANTES),
                          ("DESPUES(lo propuesto)    ", DESPUES)):
        gpt, hablante = modelo.get_conditioning_latents(
            audio_path=cfg["refs"], **cfg["cond"])
        puntos = []
        for j, frase in enumerate(FRASES):
            out = modelo.inference(text=frase, language="es",
                                   gpt_cond_latent=gpt, speaker_embedding=hablante,
                                   enable_text_splitting=True, **cfg["gen"])
            nombre = ("antes" if cfg is ANTES else "despues") + f"_{j}.wav"
            sf.write(os.path.join(SALIDA, nombre), out["wav"], 24000)
            puntos.append(parecido(patron, emb([os.path.join(SALIDA, nombre)])))
        media = sum(puntos) / len(puntos)
        detalle = " ".join(f"{p:.3f}" for p in puntos)
        print(f"  {etiqueta}  {media:.4f}   [{detalle}]", flush=True)

    print(f"\n  wavs para escuchar en {SALIDA} (antes_*.wav / despues_*.wav)")


if __name__ == "__main__":
    main()