import spaces
import gradio as gr
import librosa
import numpy as np
import torch

from transformers import pipeline

synthesiser = pipeline("text-to-speech", "techiaith/microsoft_speecht5_finetuned_bu_tts_cy_en")
synthesiser.to('cuda')

speaker_embeddings = {
    "GGP": "spkemb/speaker0.npy",
    "BGP": "spkemb/speaker1.npy",
    "BDP": "spkemb/speaker2.npy",
}

spk_model_name = "speechbrain/spkrec-xvect-voxceleb"
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f">>>>> DEVICE {device}")
speaker_model = EncoderClassifier.from_hparams(
    source=spk_model_name,
    run_opts={"device": device},
    savedir=os.path.join("/tmp", spk_model_name),
)


def create_speaker_embedding(waveform):
    with torch.no_grad():
        se = speaker_model.encode_batch(torch.tensor(waveform))
        se = torch.nn.functional.normalize(se, dim=2)
        se = se.squeeze().cpu().numpy()
    return se

@spaces.GPU
def predict(text, speaker):
    if len(text.strip()) == 0:
        return (16000, np.zeros(0).astype(np.int16))
    speaker_embedding = np.load(speaker_embeddings[speaker[:3]])
    speaker_embedding = prepare_dataset(speaker_embedding)
    speaker_embedding = torch.tensor(speaker_embedding).unsqueeze(0)
    speech = synthesiser(text, forward_params={"speaker_embeddings": speaker_embedding})
    speech = (speech.numpy() * 32767).astype(np.int16)
    return (16000, speech)


title = "Techiaith Finetune Microsoft/SpeechT5: Speech Synthesis"

description = """
Lleisiau TTS microsoft_speech_T5_finetune_bu_tts_cy_en
"""

examples = [
    ["Rhyfeddod neu ffenomenon optegol a meteorolegol yw enfys, pan fydd sbectrwm o olau yn ymddangos yn yr awyr pan fo'r haul yn disgleirio ar ddiferion o leithder yn atmosffer y ddaear.", "GGP (gwryw-gogledd-pro)"],
    ["Rhyfeddod neu ffenomenon optegol a meteorolegol yw enfys, pan fydd sbectrwm o olau yn ymddangos yn yr awyr pan fo'r haul yn disgleirio ar ddiferion o leithder yn atmosffer y ddaear.", "BGP (benyw-gogledd-pro)"],
    ["Rhyfeddod neu ffenomenon optegol a meteorolegol yw enfys, pan fydd sbectrwm o olau yn ymddangos yn yr awyr pan fo'r haul yn disgleirio ar ddiferion o leithder yn atmosffer y ddaear.", "BDP (benyw-de-pro)"],
]

gr.Interface(
    fn=predict,
    inputs=[
        gr.Text(label="Input Text"),
        gr.Radio(label="Speaker", choices=[
            "GGP (gwryw-gogledd-pro)",
            "BGP (benyw-gogledd-pro)",
            "BDP (benyw-de-pro)",
        ],
        value="GGP (gwryw-gogledd-pro)"),
    ],
    outputs=[
        gr.Audio(label="Generated Speech", type="numpy"),
    ],
    title=title,
    description=description,
    examples=examples,
).launch()