#!/usr/bin/env python3
"""¿El clon no se parece porque esta una octava por debajo?

Medido: los originales de fish.audio rondan los 210 Hz de tono medio con un
recorrido de 105 a 349; los clones salen a 105 Hz y con diez hercios de
recorrido. O sea una octava mas grave y ademas planos.

Esto lo comprueba de la unica forma que vale: subiendo el clon de tono y
midiendo si el parecido sube. Si sube mucho, el problema es el tono y se
arregla; si no se mueve, el tono es un sintoma y el problema esta en otro sitio.

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

import numpy as np
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", "jarvis_lento_b.wav")
os.environ.setdefault("COQUI_TOS_AGREED", "1")
LARGO = dict(gpt_cond_len=30, gpt_cond_chunk_len=6, max_ref_length=30)

SEMITONOS = [0, 3, 5, 7, 9, 12]


def desplaza(origen, destino, semitonos):
    if semitonos == 0:
        subprocess.run(["cp", origen, destino], check=True)
        return destino
    subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-i", origen,
                    "-af", f"rubberband=pitch={semitonos}", destino], check=True)
    return destino


def tono(ruta):
    import librosa
    y, sr = librosa.load(ruta, sr=22050, mono=True)
    f, _, _ = librosa.pyin(y, sr=sr, fmin=60, fmax=400, frame_length=2048)
    v = f[~np.isnan(f)]
    return float(np.median(v)) if len(v) else 0.0


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], **LARGO)
        return e

    patron = emb(REF)

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

    # los clones mas largos de la cache, que son los que mejor se miden
    clones = sorted(glob.glob(os.path.join(BASE, "cache_voz", "*.wav")),
                    key=os.path.getsize, reverse=True)[:4]

    print(f"  referencia: {tono(REF):.0f} Hz\n")
    print(f"  {'semitonos':>10s} {'tono clon':>10s} {'parecido':>9s}")
    print("  " + "-" * 33)

    with tempfile.TemporaryDirectory() as tmp:
        for st in SEMITONOS:
            puntos, tonos = [], []
            for i, c in enumerate(clones):
                d = os.path.join(tmp, f"c{i}_{st}.wav")
                desplaza(c, d, st)
                puntos.append(parecido(d))
                tonos.append(tono(d))
            print(f"  {st:>+10d} {np.mean(tonos):9.0f}Hz {np.mean(puntos):9.4f}",
                  flush=True)


if __name__ == "__main__":
    main()