import os
import warnings
import wave

import paddle
import pyaudio
import soundfile as sf

warnings.filterwarnings("ignore")
from paddlespeech.t2s.frontend.zh_frontend import Frontend


# 聲學模型路徑
am_model_path = 'models/fastspeech2/model'
# 模型的發聲字典
phones_dict_path = 'models/fastspeech2/phone_id_map.txt'

# 聲碼器模型路徑
voc_model_path = 'models/wavegan/model'

# 要合成的文本
text = '我是夜雨飄零，我愛深度學習！'
# 輸出文件的保持路徑
output_path = 'output/1.wav'

# 獲取文本前端
frontend = Frontend(g2p_model='g2pM', phone_vocab_path=phones_dict_path)

# 聲學模型
am_inference = paddle.jit.load(am_model_path)
# 聲碼器模型
voc_inference = paddle.jit.load(voc_model_path)

# 文本轉模型輸入
input_ids = frontend.get_input_ids(text, merge_sentences=False)
phone_ids = input_ids['phone_ids']
print(phone_ids)

# 模型輸出結果拼接
wav_all = None
for i in range(len(phone_ids)):
    part_phone_ids = phone_ids[i]
    # 獲取聲學模型的輸出
    mel = am_inference(part_phone_ids)
    # 獲取聲碼器模型輸出
    wav = voc_inference(mel)
    if wav_all is None:
        wav_all = wav
    else:
        wav_all = paddle.concat([wav_all, wav])

# Tensor轉numpy
wav = wav_all.numpy()
os.makedirs(os.path.dirname(output_path), exist_ok=True)
sf.write(output_path, wav, samplerate=24000)
print(f'音頻已保存：{output_path}')

chunk = 1024
wf  = wave.open(output_path, 'rb')
p = pyaudio.PyAudio()
stream = p.open(format=p.get_format_from_width(width=wf.getsampwidth()),
                channels=wf.getnchannels(),
                rate=wf.getframerate(),
                output=True)

# 讀取音頻
data = wf.readframes(chunk)

# 循環讀取音頻
while len(data) > 0:
    stream.write(data)
    data = wf.readframes(chunk)

stream.stop_stream()
stream.close()
p.terminate()