print("NLTK")
import nltk
nltk.download('punkt')

import gradio as gr
import numpy as np
import whisper
import scipy.io.wavfile

#StyleTTS2 imports
import torch
torch.manual_seed(0)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True

import random
random.seed(0)
np.random.seed(0)

# load packages
import yaml
from munch import Munch
import torch
from torch import nn
import torch.nn.functional as F
import torchaudio
import librosa
from nltk.tokenize import word_tokenize
from models import *
from utils import *
from text_utils import TextCleaner
textclenaer = TextCleaner()
import phonemizer


# Global values
sample_rate_value=24000
original_voice_path = "ref_voice.wav"

to_mel = torchaudio.transforms.MelSpectrogram(
    n_mels=80, n_fft=2048, win_length=1200, hop_length=300)
mean, std = -4, 4

def length_to_mask(lengths):
    mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
    mask = torch.gt(mask+1, lengths.unsqueeze(1))
    return mask

def preprocess(wave):
    wave_tensor = torch.from_numpy(wave).float()
    mel_tensor = to_mel(wave_tensor)
    mel_tensor = (torch.log(1e-5 + mel_tensor.unsqueeze(0)) - mean) / std
    return mel_tensor

def compute_style(path):
    wave, sr = librosa.load(path, sr=24000)
    audio, index = librosa.effects.trim(wave, top_db=30)
    if sr != 24000:
        audio = librosa.resample(audio, sr, 24000)
    mel_tensor = preprocess(audio).to(device)

    with torch.no_grad():
        ref_s = model.style_encoder(mel_tensor.unsqueeze(1))
        ref_p = model.predictor_encoder(mel_tensor.unsqueeze(1))

    return torch.cat([ref_s, ref_p], dim=1)

device = 'cuda' if torch.cuda.is_available() else 'cpu'

# load phonemizer
#phonemizer = Phonemizer.from_checkpoint(str(cached_path('https://public-asai-dl-models.s3.eu-central-1.amazonaws.com/DeepPhonemizer/en_us_cmudict_ipa_forward.pt')))
global_phonemizer = phonemizer.backend.EspeakBackend(language='en-us', preserve_punctuation=True,  with_stress=True)

config = yaml.safe_load(open("Models/LibriTTS/config.yml"))

# load pretrained ASR model
ASR_config = config.get('ASR_config', False)
ASR_path = config.get('ASR_path', False)
text_aligner = load_ASR_models(ASR_path, ASR_config)

# load pretrained F0 model
F0_path = config.get('F0_path', False)
pitch_extractor = load_F0_models(F0_path)

# load BERT model
from Utils.PLBERT.util import load_plbert
BERT_path = config.get('PLBERT_dir', False)
plbert = load_plbert(BERT_path)

model_params = recursive_munch(config['model_params'])
model = build_model(model_params, text_aligner, pitch_extractor, plbert)
_ = [model[key].eval() for key in model]
_ = [model[key].to(device) for key in model]

params_whole = torch.load("Models/LibriTTS/epochs_2nd_00020.pth", map_location='cpu')
params = params_whole['net']


for key in model:
    if key in params:
        print('%s loaded' % key)
        try:
            model[key].load_state_dict(params[key])
        except:
            from collections import OrderedDict
            state_dict = params[key]
            new_state_dict = OrderedDict()
            for k, v in state_dict.items():
                name = k[7:] # remove `module.`
                new_state_dict[name] = v
            # load params
            model[key].load_state_dict(new_state_dict, strict=False)
#             except:
#                 _load(params[key], model[key])
_ = [model[key].eval() for key in model]

from Modules.diffusion.sampler import DiffusionSampler, ADPM2Sampler, KarrasSchedule

sampler = DiffusionSampler(
    model.diffusion.diffusion,
    sampler=ADPM2Sampler(),
    sigma_schedule=KarrasSchedule(sigma_min=0.0001, sigma_max=3.0, rho=9.0), # empirical parameters
    clamp=False
)

def inference(text, ref_s, alpha = 0.3, beta = 0.7, diffusion_steps=5, embedding_scale=1):
    text = text.strip()
    ps = global_phonemizer.phonemize([text])
    #ps = phonemizer([text], lang='en_us')
    ps = word_tokenize(ps[0])
    ps = ' '.join(ps)
    tokens = textclenaer(ps)
    tokens.insert(0, 0)
    tokens = torch.LongTensor(tokens).to(device).unsqueeze(0)
    
    with torch.no_grad():
        input_lengths = torch.LongTensor([tokens.shape[-1]]).to(device)
        text_mask = length_to_mask(input_lengths).to(device)

        t_en = model.text_encoder(tokens, input_lengths, text_mask)
        bert_dur = model.bert(tokens, attention_mask=(~text_mask).int())
        d_en = model.bert_encoder(bert_dur).transpose(-1, -2) 

        s_pred = sampler(noise = torch.randn((1, 256)).unsqueeze(1).to(device), 
                                          embedding=bert_dur,
                                          embedding_scale=embedding_scale,
                                            features=ref_s, # reference from the same speaker as the embedding
                                             num_steps=diffusion_steps).squeeze(1)


        s = s_pred[:, 128:]
        ref = s_pred[:, :128]

        ref = alpha * ref + (1 - alpha)  * ref_s[:, :128]
        s = beta * s + (1 - beta)  * ref_s[:, 128:]

        d = model.predictor.text_encoder(d_en, 
                                         s, input_lengths, text_mask)

        x, _ = model.predictor.lstm(d)
        duration = model.predictor.duration_proj(x)

        duration = torch.sigmoid(duration).sum(axis=-1)
        pred_dur = torch.round(duration.squeeze()).clamp(min=1)


        pred_aln_trg = torch.zeros(input_lengths, int(pred_dur.sum().data))
        c_frame = 0
        for i in range(pred_aln_trg.size(0)):
            pred_aln_trg[i, c_frame:c_frame + int(pred_dur[i].data)] = 1
            c_frame += int(pred_dur[i].data)

        # encode prosody
        en = (d.transpose(-1, -2) @ pred_aln_trg.unsqueeze(0).to(device))
        if model_params.decoder.type == "hifigan":
            asr_new = torch.zeros_like(en)
            asr_new[:, :, 0] = en[:, :, 0]
            asr_new[:, :, 1:] = en[:, :, 0:-1]
            en = asr_new

        F0_pred, N_pred = model.predictor.F0Ntrain(en, s)

        asr = (t_en @ pred_aln_trg.unsqueeze(0).to(device))
        if model_params.decoder.type == "hifigan":
            asr_new = torch.zeros_like(asr)
            asr_new[:, :, 0] = asr[:, :, 0]
            asr_new[:, :, 1:] = asr[:, :, 0:-1]
            asr = asr_new

        out = model.decoder(asr, 
                                F0_pred, N_pred, ref.squeeze().unsqueeze(0))
    
        
    return out.squeeze().cpu().numpy()[..., :-50] # weird pulse at the end of the model, need to be fixed later


def transcribe(audio):
    transcribed_text = ""
    try:
        whisper_model = whisper.load_model("base")
        result = whisper_model.transcribe(audio)
        transcribed_text = result["text"]
    except Exception as exc: 
        print(exc) 
        transcribed_text = "An error occured. Please try again."  
        
    print(transcribed_text)
    # ref_s = compute_style(original_voice_path) # run locally
    ref_s = compute_style(audio) # run on HF
    wav = inference(transcribed_text, ref_s, alpha=0.1, beta=0.5, diffusion_steps=10, embedding_scale=1)
    scaled = np.int16(wav / np.max(np.abs(wav)) * 32767)
    return (sample_rate_value, scaled)



def record_speaker(audio):
    sr, voice = audio
    scaled = np.int16(voice / np.max(np.abs(voice)) * 32767)
    scipy.io.wavfile.write(original_voice_path, sr, scaled)

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(""" # AccentCoach: Transform Any Accent into American Accent.
        **This is an educational app designed to transform the speech of a non-native English speaker into a native American accent.**
        
        **The tool aims to coach learners in <ins>accent reduction</ins> and pronunciation improvement. It performs much better on <ins>longer speech</ins>.**
        **The code is based on style diffusion and adversarial training with LSLMs outlined in StyleTTS2 paper.**

        **It is strongly advised to duplicate this space and run it on a powerful GPU. Inference time can be reduced to less than a second when utilizing an Nvidia 3090.**
    """)
    # with gr.Accordion("First-Time Users (Click Here):", open=False):
    #     gr.Markdown("""
    #         **Record the reference voice:** Kindly capture your voice as you read the provided
    #         text. Please ensure that you have granted microphone access in your browser settings.

    #         > I must not fear. Fear is the mind-killer. Fear is the little-death that brings total obliteration. 
    #         I will face my fear. I will permit it to pass over me and through me. And when it has gone past I 
    #         will turn the inner eye to see its path. Where the fear has gone there will be nothing. Only I will remain.

    #         Ensure clarity in your pronunciation, as the quality of this recording
    #         significantly influences the future results. Once done please click "Save".
    #         If the quality of the native voice is not satisfactory, you can come back and
    #         re-record your voice here again.

    #         You can also upload your voice or someone else's voice.
    #     """)
    #     speaker_voice = gr.Audio(sources=["microphone", "upload"], format="wav", label="Record reference voice:",show_download_button="True")
    #     ref_btn = gr.Button("Save")
    #     ref_btn.click(record_speaker, inputs= speaker_voice, outputs=None)


    with gr.Column():
        gr.Markdown("""
            *Initiate the recording process by selecting the **Record** button. Speak Clearly and ensure a noise-free environment.*  
            """)
        inp = gr.Audio(sources=["microphone", "upload"], format="wav", type="filepath",
            label="Original accent:",show_download_button="True")
        gr.Markdown("""
            *Press the **Run** button to listen to your native accent:*  
            """)
        out = gr.Audio(label="Native accent:", autoplay="True", show_download_button="True")
        btn = gr.Button("Run")
        btn.click(transcribe, inputs=inp, outputs=out)
        gr.Examples(
            examples=[
                ["https://dl.sndup.net/9y9x/Albert-Einstein.wav",],
                ["https://dl.sndup.net/p6gz/Arnold-Schwarzenegger.wav" ,], 
            ],
            inputs=inp,
            outputs=out,
            fn=transcribe,
            cache_examples=True, 
        )

        gr.Markdown(
            """
            ## Remarks: 
            - **The optimal performance of the model is achieved when running on a GPU with a 
            minimum of 8GB of VRAM. However, due to budget constraints, the author is currently 
            limited to utilizing the free CPU on HF, resulting in slower inference speeds.**
            - **Longer sentences yield a more naturally flowing result.
            Brief expressions like "Hi" or "How are you" may yield suboptimal outcomes.**
            - **The model might occasionally produce noise or generate random speech. 
            Consider re-recording or re-running for enhanced clarity and accuracy.**
            - **By utilizing this application, you provide consent for your voice to
            be synthesized by pre-trained models.**
            - **If encountering an error, please try re-running or reloading the page.**
            - **This app primarily functions as an educational tool for English learners. 
            The author does not endorse or support any malicious or misuse of this application.**
            - **The user acknowledges and agrees that the use of the software is at the user's sole risk.**
            """)


if __name__ == "__main__":
    demo.launch()