import torch
import gradio as gr
from torchaudio.sox_effects import apply_effects_file
from transformers import AutoFeatureExtractor, AutoModelForAudioXVector

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

OUTPUT = """
    <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/bootstrap@5.1.3/dist/css/bootstrap.min.css" integrity="sha256-YvdLHPgkqJ8DVUxjjnGVlMMJtNimJ6dYkowFFvp4kKs=" crossorigin="anonymous">
    <div class="container">
        <div class="row"><h1 style="text-align: center">The speakers are</h1></div>
        <div class="row"><h1 class="display-1" style="text-align: center">{:.1f}%</h1></div>
        <div class="row"><h1 style="text-align: center">similar</h1></div>
    </div>
"""

EFFECTS = [
    ["channels", "1"],
    ["rate", "16000"],
    ["gain", "-3.0"],
    ["silence", "1", "0.1", "0.1%", "-1", "0.1", "0.1%"],
]

model_name = "anton-l/unispeech-sat-base-plus-sv"
feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
model = AutoModelForAudioXVector.from_pretrained(model_name).to(device)
cosine_sim = torch.nn.CosineSimilarity(dim=-1)


def similarity_fn(mic_path1, file_path1, mic_path2, file_path2):
    if not ((mic_path1 or file_path1) and (mic_path2 or file_path2)):
        return '<b style="color:red">ERROR: Please record or upload audio for *both* speakers!</b>'

    wav1, _ = apply_effects_file(mic_path1 if mic_path1 else file_path1, EFFECTS)
    wav2, _ = apply_effects_file(mic_path2 if mic_path2 else file_path2, EFFECTS)

    input1 = feature_extractor(wav1.squeeze(0), return_tensors="pt").input_values.to(device)
    input2 = feature_extractor(wav2.squeeze(0), return_tensors="pt").input_values.to(device)

    with torch.no_grad():
        emb1 = model(input1).embeddings
        emb2 = model(input2).embeddings
    emb1 = torch.nn.functional.normalize(emb1, dim=-1).cpu()
    emb2 = torch.nn.functional.normalize(emb2, dim=-1).cpu()
    similarity = cosine_sim(emb1, emb2).numpy()[0]

    return OUTPUT.format(similarity * 100)


inputs = [
    gr.inputs.Audio(source="microphone", type="filepath", optional=True, label="Speaker #1"),
    gr.inputs.Audio(source="upload", type="filepath", optional=True, label="or"),
    gr.inputs.Audio(source="microphone", type="filepath", optional=True, label="Speaker #2"),
    gr.inputs.Audio(source="upload", type="filepath", optional=True, label="or"),
]
output = gr.outputs.HTML(label="")


description = (
    "Speaker Verification demo based on "
    "UniSpeech-SAT: Universal Speech Representation Learning with Speaker Aware Pre-Training"
)
article = (
    "<p style='text-align: center'>"
    "<a href='https://huggingface.co/microsoft/unispeech-sat-large' target='_blank'>🎙️ Learn more about UniSpeech-SAT</a> | "
    "<a href='https://arxiv.org/abs/2110.05752' target='_blank'>📚 Article on ArXiv</a>"
    "</p>"
)

interface = gr.Interface(
    fn=similarity_fn,
    inputs=inputs,
    outputs=output,
    title="Speaker Verification with UniSpeech-SAT",
    description=description,
    article=article,
    layout="horizontal",
    theme="huggingface",
    allow_flagging=False,
    live=False,
)
interface.launch(enable_queue=True)