#!/usr/bin/env python3
"""Frases cortas generadas con carrerilla, y recortadas despues.

Medido en prueba_longitud.py: un acuse de segundo y medio puntua 0,51 de
parecido, y el mismo estilo dicho en catorce segundos puntua 0,70. XTTS
arranca cada generacion en frio y no le da tiempo a asentar la voz antes de
que la frase termine. Como la cache de JARVIS son casi todo acuses de un
segundo, lo que mas se oye es lo que peor suena.

La idea: no generar "Voy, señor" a secas, sino una frase de arrastre delante,
y quedarse solo con el final. La voz llega al trozo util ya asentada.

Se prueban dos formas de cortar:
  - por silencio: XTTS deja pausa entre oraciones, asi que el ultimo tramo
    sonoro es la frase que se queria
  - repitiendo la frase tres veces y quedandose con la ultima

    python prueba_arrastre.py [cuda|cpu] [repeticiones]
"""
import os
import statistics
import sys
import tempfile

import numpy as np
import torch

DISPOSITIVO = sys.argv[1] if len(sys.argv) > 1 else "cuda"
REPES = int(sys.argv[2]) if len(sys.argv) > 2 else 4

BASE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, BASE)
from clonador import CONDICIONADO, GENERACION, FRECUENCIA, referencias

os.environ.setdefault("COQUI_TOS_AGREED", "1")
PATRON = os.path.join(BASE, "ref", "jarvis_lento_b.wav")
MEDIR = dict(gpt_cond_len=30, gpt_cond_chunk_len=6, max_ref_length=30)

# Frase de arrastre: larga, del mismo registro, y que acaba en punto para que
# XTTS deje pausa antes de lo que viene.
ARRASTRE = ("Buenas noches, señor. He revisado el sistema y todo está en orden, "
            "sin ninguna anomalía digna de mención.")

CORTAS = ["Voy, señor.", "Enseguida, señor.", "Hecho, señor.", "A la orden, señor."]


def ultimo_tramo(wav, sr, minimo=0.35):
    """El ultimo trozo sonoro, separado por silencios."""
    import librosa
    tramos = librosa.effects.split(wav, top_db=32, frame_length=1024, hop_length=256)
    tramos = [(a, b) for a, b in tramos if (b - a) / sr >= minimo]
    if not tramos:
        return wav
    a, b = tramos[-1]
    margen = int(0.05 * sr)
    return wav[max(0, a - margen):min(len(wav), b + margen)]


def main():
    from TTS.api import TTS
    import soundfile as sf

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

    patron = emb(PATRON)

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

    gpt, hablante = modelo.get_conditioning_latents(
        audio_path=referencias(), **CONDICIONADO)

    def genera(texto):
        out = modelo.inference(text=texto, language="es", gpt_cond_latent=gpt,
                               speaker_embedding=hablante,
                               enable_text_splitting=False, **GENERACION)
        return np.asarray(out["wav"], dtype=np.float32)

    metodos = {
        "directo (como ahora)": lambda f: genera(f),
        "con arrastre delante": lambda f: ultimo_tramo(
            genera(f"{ARRASTRE} {f}"), FRECUENCIA),
        "repetida x3, la ultima": lambda f: ultimo_tramo(
            genera(" ".join([f] * 3)), FRECUENCIA),
    }

    print(f"  {len(CORTAS)} frases cortas x {REPES} repeticiones\n")
    print(f"  {'metodo':24s} {'dur':>6s} {'parecido':>9s} {'desv':>7s}")
    print("  " + "-" * 50)

    with tempfile.TemporaryDirectory() as tmp:
        for nombre, hacer in metodos.items():
            puntos, duraciones = [], []
            for frase in CORTAS:
                for r in range(REPES):
                    wav = hacer(frase)
                    d = os.path.join(tmp, f"{abs(hash(nombre))}_{len(puntos)}.wav")
                    sf.write(d, wav, FRECUENCIA)
                    duraciones.append(len(wav) / FRECUENCIA)
                    puntos.append(parecido(d))
            print(f"  {nombre:24s} {statistics.mean(duraciones):5.1f}s "
                  f"{statistics.mean(puntos):9.4f} {statistics.pstdev(puntos):7.4f}",
                  flush=True)


if __name__ == "__main__":
    main()