#!/usr/bin/env python3
#
# Copyright      2022-2023  Xiaomi Corp.        (authors: Fangjun Kuang)
#
# See LICENSE for clarification regarding multiple authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# References:
# https://gradio.app/docs/#dropdown

import logging
import os,torch
import time
from mtts.models.vocoder import *
import gradio as gr
import yaml
from model import get_pretrained_model, language_to_models
from mtts.text import TextProcessor
from mtts.models.fs2_model import FastSpeech2
title = "# Text-to-speech (TTS)"

description = """
文字转语音
"""

# css style is copied from
# https://huggingface.co/spaces/alphacep/asr/blob/main/app.py#L113
css = """
.result {display:flex;flex-direction:column}
.result_item {padding:15px;margin-bottom:8px;border-radius:15px;width:100%}
.result_item_success {background-color:mediumaquamarine;color:white;align-self:start}
.result_item_error {background-color:#ff7070;color:white;align-self:start}
"""

examples = [
    ["Chinese (Mandarin, 普通话)", "csukuangfj/vits-piper-zh_CN-huayan-medium", "你去做饭吧", 0, 1.0],
    ["Chinese (Mandarin, 普通话)", "csukuangfj/vits-piper-zh_CN-huayan-medium", "吃葡萄不吐葡萄皮", 0, 1.0],
]


def update_model_dropdown(language: str):
    if language in language_to_models:
        choices = language_to_models[language]
        return gr.Dropdown(
            choices=choices,
            value=choices[0],
            interactive=True,
        )

    raise ValueError(f"Unsupported language: {language}")


def build_html_output(s: str, style: str = "result_item_success"):
    return f"""
    <div class='result'>
        <div class='result_item {style}'>
          {s}
        </div>
    </div>
    """


def process(language: str, repo_id: str, text: str, sid: str, speed: float):
    logging.info(f"Input text: {text}. sid: {sid}, speed: {speed}")
    sid = int(sid)
    
    
    start = time.time()
    dst_file, duration = get_pretrained_model(model,text,config,text_processor,vocoder)
    end = time.time()

    

    elapsed_seconds = end - start
    rtf = elapsed_seconds / duration

    info = f"""
    Wave duration  : {duration:.3f} s <br/>
    Processing time: {elapsed_seconds:.3f} s <br/>
    RTF: {elapsed_seconds:.3f}/{duration:.3f} = {rtf:.3f} <br/>
    """

    logging.info(info)
    logging.info(f"\nrepo_id: {repo_id}\ntext: {text}\nsid: {sid}\nspeed: {speed}")

    return dst_file, build_html_output(info)

def __build_vocoder(config):
    vocoder_name = config['vocoder']['type']
    VocoderClass = eval(vocoder_name)
    model = VocoderClass(config=config['vocoder'][vocoder_name])
    return model
demo = gr.Blocks(css=css)
config = "examples/biaobei/config.yaml"
checkpoint = "checkpoints/checkpoint_140000.pth.tar"
if os.path.exists(config):
    print("file cunzai ")
else:
    print("12")
with open(config) as f:
    config = yaml.safe_load(f)
vocoder = __build_vocoder(config)
text_processor = TextProcessor(config)
model = FastSpeech2(config)
if checkpoint != '':
    print("loading model")
    sd = torch.load(checkpoint, map_location="cpu")
    if 'model' in sd.keys():
        sd = sd['model']
model.load_state_dict(sd)
model = model.to("cpu")
with demo:
    gr.Markdown(title)
    language_choices = list(language_to_models.keys())

    language_radio = gr.Radio(
        label="Language",
        choices=language_choices,
        value=language_choices[0],
    )

    model_dropdown = gr.Dropdown(
        choices=language_to_models[language_choices[0]],
        label="Select a model",
        value=language_to_models[language_choices[0]][0],
    )

    language_radio.change(
        update_model_dropdown,
        inputs=language_radio,
        outputs=model_dropdown,
    )

    with gr.Tabs():
        with gr.TabItem("Please input your text"):
            input_text = gr.Textbox(
                label="Input text",
                info="Your text",
                lines=3,
                placeholder="Please input your text here",
            )

            input_sid = gr.Textbox(
                label="Speaker ID",
                info="Speaker ID",
                lines=1,
                max_lines=1,
                value="0",
                placeholder="Speaker ID. Valid only for mult-speaker model",
            )

            input_speed = gr.Slider(
                minimum=0.1,
                maximum=10,
                value=1,
                step=0.1,
                label="Speed (larger->faster; smaller->slower)",
            )

            input_button = gr.Button("Submit")

            output_audio = gr.Audio(label="Output")

            output_info = gr.HTML(label="Info")

            gr.Examples(
                examples=examples,
                fn=process,
                inputs=[
                    language_radio,
                    model_dropdown,
                    input_text,
                    input_sid,
                    input_speed,
                ],
                outputs=[
                    output_audio,
                    output_info,
                ],
            )

        input_button.click(
            process,
            inputs=[
                language_radio,
                model_dropdown,
                input_text,
                input_sid,
                input_speed,
            ],
            outputs=[
                output_audio,
                output_info,
            ],
        )

    gr.Markdown(description)



if __name__ == "__main__":

    formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"

    logging.basicConfig(format=formatter, level=logging.INFO)

    demo.launch()