#!/usr/bin/env python3
"""Comprueba que la cache que SUENA es la mejorada, no solo las pruebas.

Compara, frase por frase, el wav de la cache anterior con el de la nueva, y los
puntua contra la referencia. Es la unica forma de saber que lo medido en el
banco de pruebas llego de verdad a lo que JARVIS dice.

    python verificar_cache.py <cache_vieja> [cuda|cpu] [cuantas]
"""
import glob
import os
import sys

import torch

BASE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, BASE)
from clonador import CONDICIONADO, REFERENCIA, RESPALDO

VIEJA = sys.argv[1]
DISPOSITIVO = sys.argv[2] if len(sys.argv) > 2 else "cuda"
CUANTAS = int(sys.argv[3]) if len(sys.argv) > 3 else 6
NUEVA = os.path.join(BASE, "cache_voz")

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


def main():
    from TTS.api import TTS
    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], **CONDICIONADO)
        return e

    # el patron es la referencia original, la misma contra la que se midio todo
    patron = emb(os.path.join(BASE, "ref", "jarvis_lento_b.wav"))

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

    comunes = sorted(set(os.path.basename(p) for p in glob.glob(os.path.join(NUEVA, "*.wav")))
                     & set(os.path.basename(p) for p in glob.glob(os.path.join(VIEJA, "*.wav"))))
    comunes = comunes[:CUANTAS]
    print(f"  {len(comunes)} frases comparadas (mismas claves en las dos caches)\n")

    peor = mejor = 0
    va, vn = [], []
    for k in comunes:
        a = parecido(os.path.join(VIEJA, k))
        n = parecido(os.path.join(NUEVA, k))
        va.append(a)
        vn.append(n)
        flecha = "sube" if n > a else "baja"
        if n > a:
            mejor += 1
        else:
            peor += 1
        print(f"    {a:.4f} -> {n:.4f}  {flecha}   {k[:12]}")

    print(f"\n    media  {sum(va)/len(va):.4f} -> {sum(vn)/len(vn):.4f}")
    print(f"    mejoran {mejor}, empeoran {peor}")


if __name__ == "__main__":
    main()