#!/usr/bin/env python3
"""Segunda vuelta: que referencia condiciona mejor, medido sobre varias frases.

La primera vuelta (comparar_ajustes.py) comparo parametros con UNA frase. Una
sola medida engaña: XTTS muestrea, y dos generaciones de lo mismo no puntuan
igual. Aqui cada candidata se mide sobre varias frases y se promedia.

Lo que se compara es de donde sale el timbre:
  - una sola referencia larga (la mejor de las normalizadas)
  - todas las normalizadas juntas, promediadas por XTTS
  - un unico wav con todo el material pegado, para que el condicionado vea
    mas fonemas del mismo hablante de una vez
  - solo las tomas de la MISMA sesion de grabacion (las dos "Buenas noches"),
    por si mezclar sesiones con tono distinto es lo que diluye

    python afinar_referencia.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")
DESCARGAS = os.path.expanduser("~/Downloads")

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

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


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


def pega(entradas, destino):
    """Concatena varios audios en un wav mono 24k."""
    lista = destino + ".txt"
    with open(lista, "w") as f:
        for e in entradas:
            f.write(f"file '{e}'\n")
    subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-f", "concat",
                    "-safe", "0", "-i", lista, "-ar", "24000", "-ac", "1",
                    destino], check=True)
    return destino


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

    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")])

    with tempfile.TemporaryDirectory() as tmp:
        normalizadas = sorted(glob.glob(os.path.join(REF, "*.wav")))
        pegadas = pega(normalizadas, os.path.join(tmp, "todo.wav"))

        misma_sesion = sorted(glob.glob(
            os.path.join(DESCARGAS, "Jarvis-*Buenas-noches*.mp3")))
        sesion_wavs = []
        for i, m in enumerate(misma_sesion):
            d = os.path.join(tmp, f"s{i}.wav")
            subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-i", m,
                            "-ar", "24000", "-ac", "1", d], check=True)
            sesion_wavs.append(d)
        sesion_pegada = pega(sesion_wavs, os.path.join(tmp, "sesion.wav")) \
            if sesion_wavs else None

        candidatas = [
            ("una_larga", [os.path.join(REF, "jarvis_lento_b.wav")]),
            ("tres_normalizadas", [os.path.join(REF, n) for n in
                                   ("jarvis_lento_b.wav", "jarvis_lento_a.wav",
                                    "jarvis_ref.wav")]),
            ("todo_pegado", [pegadas]),
        ]
        if sesion_pegada:
            candidatas.append(("misma_sesion_pegada", [sesion_pegada]))

        print(f"\n  cada candidata sobre {len(FRASES)} frases:\n")
        tabla = []
        for nombre, refs in candidatas:
            gpt, hablante = modelo.get_conditioning_latents(
                audio_path=refs, **LARGO)
            puntos = []
            for j, frase in enumerate(FRASES):
                salida = modelo.inference(
                    text=frase, language="es", gpt_cond_latent=gpt,
                    speaker_embedding=hablante, temperature=0.65,
                    repetition_penalty=2.0, speed=1.0,
                    enable_text_splitting=True)
                d = os.path.join(tmp, f"{nombre}_{j}.wav")
                sf.write(d, salida["wav"], 24000)
                puntos.append(parecido(patron, emb([d])))
            media = sum(puntos) / len(puntos)
            tabla.append((media, nombre, puntos))
            detalle = " ".join(f"{p:.3f}" for p in puntos)
            print(f"    {media:.4f}  {nombre:22s} [{detalle}]", flush=True)

        print("\n  mejor:", max(tabla)[1])


if __name__ == "__main__":
    main()