import argparse
import os
from pathlib import Path

import logging
import re_matching

logging.getLogger("numba").setLevel(logging.WARNING)
logging.getLogger("markdown_it").setLevel(logging.WARNING)
logging.getLogger("urllib3").setLevel(logging.WARNING)
logging.getLogger("matplotlib").setLevel(logging.WARNING)

logging.basicConfig(
    level=logging.INFO, format="| %(name)s | %(levelname)s | %(message)s"
)

logger = logging.getLogger(__name__)

import librosa
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm
from clap_wrapper import get_clap_audio_feature, get_clap_text_feature


import gradio as gr

import utils
from config import config

import torch
import commons
from text import cleaned_text_to_sequence, get_bert
from text.cleaner import clean_text
import utils

from models import SynthesizerTrn
from text.symbols import symbols
import sys

from tools.translate import translate

net_g = None

device = (
        "cuda:0"
        if torch.cuda.is_available()
        else (
            "mps"
            if sys.platform == "darwin" and torch.backends.mps.is_available()
            else "cpu"
        )
    )

#device = "cpu"
BandList = {
        "PoppinParty":["香澄","有咲","たえ","りみ","沙綾"],
        "Afterglow":["蘭","モカ","ひまり","巴","つぐみ"],
        "HelloHappyWorld":["こころ","美咲","薫","花音","はぐみ"],
        "PastelPalettes":["彩","日菜","千聖","イヴ","麻弥"],
        "Roselia":["友希那","紗夜","リサ","燐子","あこ"],
        "RaiseASuilen":["レイヤ","ロック","ますき","チュチュ","パレオ"],
        "Morfonica":["ましろ","瑠唯","つくし","七深","透子"],
        "MyGo":["燈","愛音","そよ","立希","楽奈"],
        "AveMujica":["祥子","睦","海鈴","にゃむ","初華"],
}

def get_net_g(model_path: str,  device: str, hps):
    net_g = SynthesizerTrn(
        len(symbols),
        hps.data.filter_length // 2 + 1,
        hps.train.segment_size // hps.data.hop_length,
        n_speakers=hps.data.n_speakers,
        **hps.model,
    ).to(device)
    _ = net_g.eval()
    _ = utils.load_checkpoint(model_path, net_g, None, skip_optimizer=True)
    return net_g

def get_text(text, language_str, hps, device, style_text=None, style_weight=0.7):
    style_text = None if style_text == "" else style_text
    norm_text, phone, tone, word2ph = clean_text(text, language_str)
    phone, tone, language = cleaned_text_to_sequence(phone, tone, language_str)

    if hps.data.add_blank:
        phone = commons.intersperse(phone, 0)
        tone = commons.intersperse(tone, 0)
        language = commons.intersperse(language, 0)
        for i in range(len(word2ph)):
            word2ph[i] = word2ph[i] * 2
        word2ph[0] += 1
    bert_ori = get_bert(
        norm_text, word2ph, language_str, device, style_text, style_weight
    )
    del word2ph
    assert bert_ori.shape[-1] == len(phone), phone

    if language_str == "ZH":
        bert = bert_ori
        ja_bert = torch.randn(1024, len(phone))
        en_bert = torch.randn(1024, len(phone))
    elif language_str == "JP":
        bert = torch.randn(1024, len(phone))
        ja_bert = bert_ori
        en_bert = torch.randn(1024, len(phone))
    elif language_str == "EN":
        bert = torch.randn(1024, len(phone))
        ja_bert = torch.randn(1024, len(phone))
        en_bert = bert_ori
    else:
        raise ValueError("language_str should be ZH, JP or EN")

    assert bert.shape[-1] == len(
        phone
    ), f"Bert seq len {bert.shape[-1]} != {len(phone)}"

    phone = torch.LongTensor(phone)
    tone = torch.LongTensor(tone)
    language = torch.LongTensor(language)
    return bert, ja_bert, en_bert, phone, tone, language


def infer(
    text,
    sdp_ratio,
    noise_scale,
    noise_scale_w,
    length_scale,
    sid,
    style_text=None,
    style_weight=0.7,
):

    language= 'JP' if is_japanese(text) else 'ZH'
    bert, ja_bert, en_bert, phones, tones, lang_ids = get_text(
        text,
        language,
        hps,
        device,
        style_text=style_text,
        style_weight=style_weight,
    )
    with torch.no_grad():
        x_tst = phones.to(device).unsqueeze(0)
        tones = tones.to(device).unsqueeze(0)
        lang_ids = lang_ids.to(device).unsqueeze(0)
        bert = bert.to(device).unsqueeze(0)
        ja_bert = ja_bert.to(device).unsqueeze(0)
        en_bert = en_bert.to(device).unsqueeze(0)
        x_tst_lengths = torch.LongTensor([phones.size(0)]).to(device)
        # emo = emo.to(device).unsqueeze(0)
        del phones
        speakers = torch.LongTensor([hps.data.spk2id[sid]]).to(device)
        audio = (
            net_g.infer(
                x_tst,
                x_tst_lengths,
                speakers,
                tones,
                lang_ids,
                bert,
                ja_bert,
                en_bert,
                sdp_ratio=sdp_ratio,
                noise_scale=noise_scale,
                noise_scale_w=noise_scale_w,
                length_scale=length_scale,
            )[0][0, 0]
            .data.cpu()
            .float()
            .numpy()
        )
        del (
            x_tst,
            tones,
            lang_ids,
            bert,
            x_tst_lengths,
            speakers,
            ja_bert,
            en_bert,
        )  # , emo
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        return (hps.data.sampling_rate,gr.processing_utils.convert_to_16_bit_wav(audio))

def is_japanese(string):
        for ch in string:
            if ord(ch) > 0x3040 and ord(ch) < 0x30FF:
                return True
        return False

def loadmodel(model):
    _ = net_g.eval()
    _ = utils.load_checkpoint(model, net_g, None, skip_optimizer=True)
    return "success"

if __name__ == "__main__":
    languages = [ "Auto", "ZH", "JP"]
    modelPaths = []
    for dirpath, dirnames, filenames in os.walk('Data/BangDream/models/'):
        for filename in filenames:
            modelPaths.append(os.path.join(dirpath, filename))
    hps = utils.get_hparams_from_file('Data/BangDream/configs/config.json')
    net_g = get_net_g(
        model_path=modelPaths[-1], device=device, hps=hps
    )
    speaker_ids = hps.data.spk2id
    speakers = list(speaker_ids.keys())
    with gr.Blocks() as app:
        gr.Markdown(value="""
            ([Bert-Vits2](https://github.com/Stardust-minus/Bert-VITS2) V2.3)少歌邦邦全员在线语音合成\n
            镜像 [V2.2](https://huggingface.co/spaces/Mahiruoshi/MyGO_VIts-bert)\n
            [好玩的](http://love.soyorin.top/)\n
            该界面的真实链接(国内可用): https://mahiruoshi-bangdream-bert-vits2.hf.space/\n
            API: https://mahiruoshi-bert-vits2-api.hf.space/ \n
            调用方式: https://mahiruoshi-bert-vits2-api.hf.space/?text=%E4%B8%BA%E4%BB%80%E4%B9%88%E8%A6%81%E6%BC%94%E5%A5%8F%E6%98%A5%E6%97%A5%E5%BD%B1&speaker=%E9%A6%99%E6%BE%84\n
            推荐搭配[Legado开源阅读](https://github.com/gedoor/legado)或[聊天bot](https://github.com/Paraworks/BangDreamAi)使用\n
            二创请标注作者：B站@Mahiroshi: https://space.bilibili.com/19874615\n
            训练数据集归属：BangDream及少歌手游,提取自BestDori,[数据集获取流程](https://nijigaku.top/2023/09/29/Bestbushiroad%E8%AE%A1%E5%88%92-vits-%E9%9F%B3%E9%A2%91%E6%8A%93%E5%8F%96%E5%8F%8A%E6%95%B0%E6%8D%AE%E9%9B%86%E5%AF%B9%E9%BD%90/)\n
            BangDream数据集下载[链接](https://huggingface.co/spaces/Mahiruoshi/BangDream-Bert-VITS2/blob/main/%E7%88%AC%E8%99%AB/SortPathUrl.txt)\n
            ！！！注意：huggingface容器仅用作展示，建议在右上角更多选项中克隆本项目或Docker运行app.py/server.py,环境参考requirements.txt\n""")
        for band in BandList:
            with gr.TabItem(band):
                for name in BandList[band]:
                    with gr.TabItem(name):
                        with gr.Row():
                            with gr.Column():
                                with gr.Row():
                                    gr.Markdown(
                                        '<div align="center">'
                                        f'<img style="width:auto;height:400px;" src="https://mahiruoshi-bangdream-bert-vits2.hf.space/file/image/{name}.png">' 
                                        '</div>'
                                    )
                                length_scale = gr.Slider(
                                        minimum=0.1, maximum=2, value=1, step=0.01, label="语速调节"
                                    )
                                with gr.Accordion(label="参数设定", open=True):
                                    sdp_ratio = gr.Slider(
                                    minimum=0, maximum=1, value=0.5, step=0.01, label="SDP/DP混合比"
                                    )
                                    noise_scale = gr.Slider(
                                        minimum=0.1, maximum=2, value=0.6, step=0.01, label="感情调节"
                                    )
                                    noise_scale_w = gr.Slider(
                                        minimum=0.1, maximum=2, value=0.667, step=0.01, label="音素长度"
                                    )
                                    speaker = gr.Dropdown(
                                        choices=speakers, value=name, label="说话人"
                                    ) 
                                with gr.Accordion(label="切换模型", open=False):
                                    modelstrs = gr.Dropdown(label = "模型", choices = modelPaths, value = modelPaths[0], type = "value")
                                    btnMod = gr.Button("载入模型")
                                    statusa = gr.TextArea(label = "模型加载状态")
                                    btnMod.click(loadmodel, inputs=[modelstrs], outputs = [statusa])
                            with gr.Column():
                                text = gr.TextArea(
                                    label="文本输入",
                                    info="输入纯日语或者中文",
                                    value="我是来结束这个乐队的。",
                                )
                                style_text = gr.Textbox(label="辅助文本",info="语言保持跟主文本一致",placeholder="为什么要演奏春日影!")
                                style_weight = gr.Slider(
                                        minimum=0,
                                        maximum=1,
                                        value=0.7,
                                        step=0.1,
                                        label="Weight",
                                        info="主文本和辅助文本的bert混合比率，0表示仅主文本，1表示仅辅助文本",
                                    )
                                btn = gr.Button("点击生成", variant="primary")
                                audio_output = gr.Audio(label="Output Audio")
                                btntran = gr.Button("快速中翻日")
                                translateResult = gr.TextArea(label="百度翻译",value="从这里翻译后的文本")
                                btntran.click(translate, inputs=[text], outputs = [translateResult])
                                
                    btn.click(
                        infer,
                        inputs=[
                            text,
                            sdp_ratio,
                            noise_scale,
                            noise_scale_w,
                            length_scale,
                            speaker,
                            style_text,
                            style_weight,
                        ],
                        outputs=[audio_output],
                    )
        with gr.TabItem('少歌在2.2版本'):
            gr.Markdown(
                '<div align="center">'
                f'<iframe style="width:100%;height:400px;" src="https://mahiruoshi-mygo-vits-bert.hf.space/" frameborder="0"></iframe>'
                '</div>'
            )
    print("推理页面已开启!")
    app.launch(share=True)