import argparse
from models import SynthesizerTrn
from text.symbols import symbols
import utils
import gradio as gr
import tempfile
from scipy.io import wavfile
import numpy as np
import re
from synthesize_fastapi import *
config = "./configs/baker_base.json"

import subprocess
def install(name):
    subprocess.call(['pip', 'install', 'packages/ttsfrd-0.0.4-cp38-cp38-linux_x86_64.whl'])
install('ttsfrd')

import ttsfrd

print("---------- Loading VITS Model ----------")
hps = utils.get_hparams_from_file(config)
net_g = SynthesizerTrn(
    len(symbols),
    hps.data.filter_length // 2 + 1,
    hps.train.segment_size // hps.data.hop_length,
    **hps.model)
_ = net_g.eval()

_ = utils.load_checkpoint("./ckpt/pretrained_baker.pth", net_g, None)

lexicon_mandarin = read_lexicon("./lexicon/pinyin-lexicon-r.txt")

# 解压文件
import zipfile
f = zipfile.ZipFile("./resource.zip", 'r')  # 压缩文件位置
for file in f.namelist():
    f.extract(file, "./")  # 解压位置
f.close()

frontend = ttsfrd.TtsFrontendEngine()
model_dir = './resource'
frontend.initialize(model_dir)
frontend.set_lang_type('zhcn')


def tts(text):
    audio_all = np.zeros(1, dtype=np.int16)  # 设置初始音频
    # 删除英文
    text = re.sub('[a-zA-Z]', '', text)
    text = text.replace(',[object Object]', '')
    pinyin_list, prosody_list = g2p_mandarin(frontend, text)
    for texts, phone_prosody in zip(pinyin_list, prosody_list):
        print(texts)
        print(phone_prosody)
        stn_tst = get_text(texts, hps)
        prosody = get_prosody(phone_prosody, hps)
        with torch.no_grad():
            x_tst = stn_tst.unsqueeze(0)
            prosody = prosody.unsqueeze(0)
            x_tst_lengths = torch.LongTensor([stn_tst.size(0)])
            audio = net_g.infer(x_tst, x_tst_lengths, prosody, noise_scale=.667, noise_scale_w=0.0, length_scale=1)[0][
                        0, 0].data.cpu().float().numpy() * 32767.0
        i = np.random.uniform(0.12, 0.35, 1)[0]
        space_time = np.zeros(int(i * 22050), dtype=np.int16)
        audio_all = np.concatenate((audio_all, audio, space_time))
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as fp:
        wavfile.write(
            fp.name,
            22050,
            audio_all.astype(np.int16),
        )
    return fp.name


inputs = [gr.inputs.Textbox(label="Input Text", default='国务院银行业监督管理机构会按照国家法律法规规定的程序处理', lines=6)]
outputs = gr.Audio(label="Output")
interface = gr.Interface(fn=tts, inputs=inputs, outputs=outputs,
                         title='Empathy-TTS',
                         description='Note: This space is running on CPU, inference times will be higher.')
interface.launch(server_name='0.0.0.0')