import torch

from transformers import pipeline

import numpy as np
import gradio as gr

def _grab_best_device(use_gpu=True):
    if torch.cuda.device_count() > 0 and use_gpu:
        device = "cuda"
    else:
        device = "cpu"
    return device

device = _grab_best_device()

default_model_per_language = {
    "marathi": "facebook/mms-tts-mar"
}

models_per_language = {
    "marathi": ["ylacombe/mms-mar-finetuned-monospeaker"]
}

HUB_PATH = "ylacombe/vits_ljs_midlands_male_monospeaker"


pipe_dict = {
    "current_model": "ylacombe/vits_ljs_midlands_male_monospeaker",
    "pipe":  pipeline("text-to-speech", model=HUB_PATH, device=0),
    "original_pipe": pipeline("text-to-speech", model=default_model_per_language["marathi"], device=0),
    "language": "english",
}

title =      """
Marathi Parkinson Enabler: Speaking is a big challenge during Parakinsons. Patients show slurred speech and cannot communicate effectively. 
This is marathi text to speech model for parkinson users who want to communicate in Marathi. 
            """

max_speakers = 1


# Inference
def generate_audio(text, model_id, language):

    if pipe_dict["language"] != language:
        gr.Warning(f"Language has changed - loading new default model: {default_model_per_language[language]}")
        pipe_dict["language"] = language
        pipe_dict["original_pipe"] = pipeline("text-to-speech", model=default_model_per_language[language], device=0)
    
    if pipe_dict["current_model"] != model_id:
        gr.Warning("Model has changed - loading new model")
        pipe_dict["pipe"] = pipeline("text-to-speech", model=model_id, device=0)
        pipe_dict["current_model"] = model_id

    num_speakers = pipe_dict["pipe"].model.config.num_speakers

    out = []
    # first generate original model result
    output = pipe_dict["original_pipe"](text)
    output =  gr.Audio(value = (output["sampling_rate"], output["audio"].squeeze()), type="numpy", autoplay=False, label=f"Non finetuned model prediction {default_model_per_language[language]}", show_label=True,
                               visible=True)
    out.append(output)
    
    
    if num_speakers>1:
        for i in range(min(num_speakers, max_speakers - 1)):
            forward_params = {"speaker_id": i}
            output = pipe_dict["pipe"](text, forward_params=forward_params)
            
            output =  gr.Audio(value = (output["sampling_rate"], output["audio"].squeeze()), type="numpy", autoplay=False, label=f"Generated Audio - speaker {i}", show_label=True,
                               visible=True)
            out.append(output)
        out.extend([gr.Audio(visible=False)]*(max_speakers-num_speakers))
    else:
        output = pipe_dict["pipe"](text)
        output =  gr.Audio(value = (output["sampling_rate"], output["audio"].squeeze()), type="numpy", autoplay=False, label="Generated Audio - Mono speaker", show_label=True,
                               visible=True)
        out.append(output)
        out.extend([gr.Audio(visible=False)]*(max_speakers-2))
    return out


css = """
#container{
    margin: 0 auto;
    max-width: 80rem;
}
#intro{
    max-width: 100%;
    text-align: center;
    margin: 0 auto;
}
"""
# Gradio blocks demo    
with gr.Blocks(css=css) as demo_blocks:
    gr.Markdown(title, elem_id="intro")

    with gr.Row():
        with gr.Column():
            inp_text = gr.Textbox(label="Input Text", info="What sentence would you like to synthesise?")
            btn = gr.Button("Generate Audio!")
            language = gr.Dropdown(
                default_model_per_language.keys(),
                value = "marathi",
                label = "language",
                info = "Language that you want to test"
            )
            
            model_id = gr.Dropdown(
                    models_per_language["marathi"],
                    value="ylacombe/mms-mar-finetuned-monospeaker", 
                    label="Model", 
                    info="Model you want to test",
                    )
                
        with gr.Column():
            outputs = []
            for i in range(max_speakers):
                out_audio = gr.Audio(type="numpy", autoplay=False, label=f"Generated Audio - speaker {i}", show_label=True, visible=False)
                outputs.append(out_audio)

    with gr.Accordion("Datasets and models details", open=False):
        gr.Markdown("""
        
### Marathi
* **Model**: [Marathi MMS TTS](https://huggingface.co/facebook/mms-tts-mar).
* **Datasets**:
    - [Marathi TTS dataset](https://huggingface.co/datasets/ylacombe/google-chilean-marathi). 
                    """) 


    language.change(lambda language: gr.Dropdown(
                    models_per_language[language],
                    value=models_per_language[language][0], 
                    label="Model", 
                    info="Model you want to test",
                    ),
                    language,
                    model_id
                   )
    
    btn.click(generate_audio, [inp_text, model_id, language], outputs)
    

demo_blocks.queue().launch()