#!/usr/bin/env python3
"""Compara ajustes de clonado y mide cuanto se parece cada uno al original.

El problema: "se parece pero no acaba de ser el". Juzgarlo solo de oido es
lento y poco fiable —dos escuchas seguidas ya no opinan igual—, asi que aqui se
genera la misma frase con varias configuraciones y se puntua cada resultado
contra la referencia con el propio codificador de hablante de XTTS.

La puntuacion es coseno entre embeddings de hablante: 1.0 seria identico. No es
la verdad absoluta —el oido manda— pero ordena las opciones y evita perder la
tarde comparando cosas que objetivamente van peor.

El hallazgo que motiva esto: la API de alto nivel usa `gpt_cond_len=6`, o sea
que de una referencia de 25 segundos solo condiciona con los 6 primeros. Lo
demas se tira. Subirlo es lo primero que hay que probar.

    python comparar_ajustes.py [cuda|cpu]
"""
import os
import sys
import time

import torch

DISPOSITIVO = sys.argv[1] if len(sys.argv) > 1 else "cuda"
BASE = os.path.dirname(os.path.abspath(__file__))
REF_DIR = os.path.join(BASE, "ref")
SALIDA = os.path.join(BASE, "comparativa")

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

# La misma frase que dice uno de los audios descargados, para poder comparar
# el clon con el original diciendo exactamente lo mismo.
FRASE = "Buenas noches, señor. Soy su asistente personal."

TODAS = ["jarvis_lento_b.wav", "jarvis_lento_a.wav", "jarvis_ref.wav"]
LARGA = ["jarvis_lento_b.wav"]          # la mas larga y limpia (25,8 s)

# cada configuracion: (nombre, referencias, kwargs de latents, kwargs de sintesis)
CONFIGS = [
    ("A_actual", TODAS,
     dict(gpt_cond_len=6, gpt_cond_chunk_len=6, max_ref_length=30),
     dict(temperature=0.65, repetition_penalty=2.0, speed=1.0)),

    ("B_condicion_larga", TODAS,
     dict(gpt_cond_len=30, gpt_cond_chunk_len=6, max_ref_length=30),
     dict(temperature=0.65, repetition_penalty=2.0, speed=1.0)),

    ("C_una_ref_larga", LARGA,
     dict(gpt_cond_len=30, gpt_cond_chunk_len=6, max_ref_length=30),
     dict(temperature=0.65, repetition_penalty=2.0, speed=1.0)),

    ("D_una_ref_fria", LARGA,
     dict(gpt_cond_len=30, gpt_cond_chunk_len=6, max_ref_length=30),
     dict(temperature=0.55, repetition_penalty=5.0, speed=1.0)),

    ("E_una_ref_normalizada", LARGA,
     dict(gpt_cond_len=30, gpt_cond_chunk_len=6, max_ref_length=30,
          sound_norm_refs=True),
     dict(temperature=0.6, repetition_penalty=5.0, speed=1.0)),
]


def rutas(nombres):
    return [os.path.join(REF_DIR, n) for n in nombres
            if os.path.exists(os.path.join(REF_DIR, n))]


def parecido(a, b):
    """Coseno entre dos embeddings de hablante."""
    a = a.squeeze().float()
    b = b.squeeze().float()
    return float(torch.nn.functional.cosine_similarity(a, b, dim=0))


def main():
    os.makedirs(SALIDA, exist_ok=True)
    from TTS.api import TTS

    print(f"cargando XTTS en {DISPOSITIVO}...", flush=True)
    t0 = time.time()
    api = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to(DISPOSITIVO)
    modelo = api.synthesizer.tts_model
    print(f"cargado en {time.time()-t0:.1f} s\n", flush=True)

    # patron a batir: el embedding de la referencia mas larga
    _, patron = modelo.get_conditioning_latents(
        audio_path=rutas(LARGA), gpt_cond_len=30, gpt_cond_chunk_len=6,
        max_ref_length=30)

    resultados = []
    for nombre, refs, cond_kw, gen_kw in CONFIGS:
        paths = rutas(refs)
        if not paths:
            print(f"  {nombre}: sin referencias, saltada")
            continue
        t0 = time.time()
        gpt_latent, hablante = modelo.get_conditioning_latents(
            audio_path=paths, **cond_kw)
        salida = modelo.inference(
            text=FRASE, language="es",
            gpt_cond_latent=gpt_latent, speaker_embedding=hablante,
            enable_text_splitting=True, **gen_kw)
        tardo = time.time() - t0

        destino = os.path.join(SALIDA, f"{nombre}.wav")
        import soundfile as sf
        sf.write(destino, salida["wav"], 24000)

        # puntuar el RESULTADO contra el patron, con el mismo codificador
        _, emb_salida = modelo.get_conditioning_latents(
            audio_path=[destino], gpt_cond_len=30, gpt_cond_chunk_len=6,
            max_ref_length=30)
        p = parecido(patron, emb_salida)
        resultados.append((p, nombre, tardo))
        print(f"  {nombre:24s} parecido {p:.4f}   ({tardo:.1f}s)", flush=True)

    print("\n  ranking (mas alto = mas se parece al original):")
    for p, nombre, tardo in sorted(resultados, reverse=True):
        print(f"    {p:.4f}  {nombre}")
    print(f"\n  wavs en {SALIDA}")


if __name__ == "__main__":
    main()