# load the libraries for the application
# -------------------------------------------
import os
import re
import nltk
import torch
import librosa
import tempfile
import subprocess

import gradio as gr

from scipy.io import wavfile
from nnet import utils, commons
from transformers import pipeline
from scipy.io.wavfile import write
from faster_whisper import WhisperModel
from nnet.models import SynthesizerTrn as vitsTRN
from nnet.models_vc import SynthesizerTrn as freeTRN
from nnet.mel_processing import mel_spectrogram_torch
from configurations.get_constants import constantConfig

from speaker_encoder.voice_encoder import SpeakerEncoder

from df_local.enhance import enhance, init_df, load_audio, save_audio
from configurations.get_hyperparameters import hyperparameterConfig
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline

nltk.download('punkt')
from nltk.tokenize import sent_tokenize

# making the FreeVC function
# ---------------------------------
class FreeVCModel:
    def __init__(self, config, ptfile, speaker_model, wavLM_model, device='cpu'):
        self.hps = utils.get_hparams_from_file(config)

        self.net_g = freeTRN(
            self.hps.data.filter_length // 2 + 1,
            self.hps.train.segment_size // self.hps.data.hop_length,
            **self.hps.model
        ).to(hyperparameters.device)
        _ = self.net_g.eval()
        _ = utils.load_checkpoint(ptfile, self.net_g, None, True)
        
        self.cmodel = utils.get_cmodel(device, wavLM_model)

        if self.hps.model.use_spk:
            self.smodel = SpeakerEncoder(speaker_model)

    def convert(self, src, tgt):
        fs_src, src_audio = src
        fs_tgt, tgt_audio = tgt

        src = f"{constants.temp_audio_folder}/src.wav"
        tgt = f"{constants.temp_audio_folder}/tgt.wav"
        out = f"{constants.temp_audio_folder}/cnvr.wav"
        with torch.no_grad():
            wavfile.write(tgt, fs_tgt, tgt_audio)
            wav_tgt, _ = librosa.load(tgt, sr=self.hps.data.sampling_rate)
            wav_tgt, _ = librosa.effects.trim(wav_tgt, top_db=20)
            if self.hps.model.use_spk:
                g_tgt = self.smodel.embed_utterance(wav_tgt)
                g_tgt = torch.from_numpy(g_tgt).unsqueeze(0).to(hyperparameters.device.type)
            else:
                wav_tgt = torch.from_numpy(wav_tgt).unsqueeze(0).to(hyperparameters.device.type)
                mel_tgt = mel_spectrogram_torch(
                    wav_tgt,
                    self.hps.data.filter_length,
                    self.hps.data.n_mel_channels,
                    self.hps.data.sampling_rate,
                    self.hps.data.hop_length,
                    self.hps.data.win_length,
                    self.hps.data.mel_fmin,
                    self.hps.data.mel_fmax,
                )
            wavfile.write(src, fs_src, src_audio)
            wav_src, _ = librosa.load(src, sr=self.hps.data.sampling_rate)
            wav_src = torch.from_numpy(wav_src).unsqueeze(0).to(hyperparameters.device.type)
            c = utils.get_content(self.cmodel, wav_src)

            if self.hps.model.use_spk:
                audio = self.net_g.infer(c, g=g_tgt)
            else:
                audio = self.net_g.infer(c, mel=mel_tgt)
            audio = audio[0][0].data.cpu().float().numpy()
            write(out, 24000, audio)

            return out

# load the system configurations
constants       = constantConfig()
hyperparameters = hyperparameterConfig()

# load the models
model, df_state, _  = init_df(hyperparameters.voice_enhacing_model, config_allow_defaults=True) # voice enhancing model
stt_model           = WhisperModel(hyperparameters.stt_model, device=hyperparameters.device.type, compute_type="float32") #speech to text model

trans_model     = AutoModelForSeq2SeqLM.from_pretrained(constants.model_name_dict[hyperparameters.nllb_model], torch_dtype=torch.bfloat16).to(hyperparameters.device)
trans_tokenizer = AutoTokenizer.from_pretrained(constants.model_name_dict[hyperparameters.nllb_model])

modelConvertSpeech = FreeVCModel(config=hyperparameters.text2speech_config, ptfile=hyperparameters.text2speech_model,
                                 speaker_model=hyperparameters.text2speech_encoder, wavLM_model=hyperparameters.wavlm_model,
                                 device=hyperparameters.device.type)

# download the language model if doesn't existing
# ----------------------------------------------------
def download(lang, lang_directory):
    
    if not os.path.exists(f"{lang_directory}/{lang}"):
        cmd = ";".join([
                f"wget {constants.language_download_web}/{lang}.tar.gz -O {lang_directory}/{lang}.tar.gz",
                f"tar zxvf {lang_directory}/{lang}.tar.gz -C {lang_directory}"
        ])
        subprocess.check_output(cmd, shell=True)
    try:
        os.remove(f"{lang_directory}/{lang}.tar.gz")
    except:
        pass
    return f"{lang_directory}/{lang}"

def preprocess_char(text, lang=None):
    """
    Special treatement of characters in certain languages
    """
    if lang == 'ron':
        text = text.replace("ț", "ţ")
    return text

def preprocess_text(txt, text_mapper, hps, uroman_dir=None, lang=None):
    txt = preprocess_char(txt, lang=lang)
    is_uroman = hps.data.training_files.split('.')[-1] == 'uroman'
    if is_uroman:
        txt  = text_mapper.uromanize(txt, f'{uroman_dir}/bin/uroman.pl')
            
    txt = txt.lower()
    txt = text_mapper.filter_oov(txt)
    return txt

def detect_language(text,LID):
    predictions = LID.predict(text)
    detected_lang_code = predictions[0][0].replace("__label__", "")
    return detected_lang_code

# text to speech
class TextMapper(object):
    def __init__(self, vocab_file):
        self.symbols = [x.replace("\n", "") for x in open(vocab_file, encoding="utf-8").readlines()]
        self.SPACE_ID = self.symbols.index(" ")
        self._symbol_to_id = {s: i for i, s in enumerate(self.symbols)}
        self._id_to_symbol = {i: s for i, s in enumerate(self.symbols)}

    def text_to_sequence(self, text, cleaner_names):
        '''Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
        Args:
        text: string to convert to a sequence
        cleaner_names: names of the cleaner functions to run the text through
        Returns:
        List of integers corresponding to the symbols in the text
        '''
        sequence = []
        clean_text = text.strip()
        for symbol in clean_text:
            symbol_id = self._symbol_to_id[symbol]
            sequence += [symbol_id]
        return sequence
    
    def uromanize(self, text, uroman_pl):
        with tempfile.NamedTemporaryFile() as tf, \
             tempfile.NamedTemporaryFile() as tf2:
            with open(tf.name, "w") as f:
                f.write("\n".join([text]))
            cmd = f"perl " + uroman_pl
            cmd += f" -l xxx "
            cmd +=  f" < {tf.name} > {tf2.name}"
            os.system(cmd)
            outtexts = []
            with open(tf2.name) as f:
                for line in f:
                    line =  re.sub(r"\s+", " ", line).strip()
                    outtexts.append(line)
            outtext = outtexts[0]
        return outtext

    def get_text(self, text, hps):
        text_norm = self.text_to_sequence(text, hps.data.text_cleaners)
        if hps.data.add_blank:
            text_norm = commons.intersperse(text_norm, 0)
        text_norm = torch.LongTensor(text_norm)
        return text_norm
    
    def filter_oov(self, text):
        val_chars = self._symbol_to_id
        txt_filt = "".join(list(filter(lambda x: x in val_chars, text)))
        return txt_filt

def speech_to_text(audio_file):
    try:
        fs, audio = audio_file
        wavfile.write(constants.input_speech_file, fs, audio)
        audio0, _ = load_audio(constants.input_speech_file, sr=df_state.sr())

        # Enhance the SNR of the audio
        enhanced = enhance(model, df_state, audio0)
        save_audio(constants.enhanced_speech_file, enhanced, df_state.sr())

        segments, info = stt_model.transcribe(constants.enhanced_speech_file)

        speech_text = ''
        for segment in segments:
            speech_text = f'{speech_text}{segment.text}'
        try:
            source_lang_nllb = [k for k, v in constants.flores_codes_to_tts_codes.items() if v[:2] == info.language][0]
        except:
            source_lang_nllb = 'language cant be determined, select manually'
        
        # text translation
        return speech_text, gr.Dropdown.update(value=source_lang_nllb)
    except:
        return '', gr.Dropdown.update(value='English')

# Text tp speech
def text_to_speech(text, target_lang):
    txt = text

    # LANG = get_target_tts_lang(target_lang)
    LANG = constants.flores_codes_to_tts_codes[target_lang]
    ckpt_dir = download(LANG, lang_directory=constants.language_directory)

    vocab_file  = f"{ckpt_dir}/{constants.language_vocab_text}"
    config_file = f"{ckpt_dir}/{constants.language_vocab_configuration}"
    hps = utils.get_hparams_from_file(config_file)
    text_mapper = TextMapper(vocab_file)
    net_g = vitsTRN(
        len(text_mapper.symbols),
        hps.data.filter_length // 2 + 1,
        hps.train.segment_size // hps.data.hop_length,
        **hps.model)
    net_g.to(hyperparameters.device)
    _ = net_g.eval()

    g_pth = f"{ckpt_dir}/{constants.language_vocab_model}"

    _ = utils.load_checkpoint(g_pth, net_g, None)

    txt = preprocess_text(txt, text_mapper, hps, lang=LANG, uroman_dir=constants.uroman_directory)
    stn_tst = text_mapper.get_text(txt, hps)
    with torch.no_grad():
        x_tst = stn_tst.unsqueeze(0).to(hyperparameters.device)
        x_tst_lengths = torch.LongTensor([stn_tst.size(0)]).to(hyperparameters.device)
        hyp = net_g.infer(
            x_tst, x_tst_lengths, noise_scale=.667,
            noise_scale_w=0.8, length_scale=1.0
        )[0][0,0].cpu().float().numpy()

    return hps.data.sampling_rate, hyp

def translation(audio, text, source_lang_nllb, target_code_nllb, output_type, sentence_mode):
    target_code    = constants.flores_codes[target_code_nllb]
    translator     = pipeline('translation', model=trans_model, tokenizer=trans_tokenizer, src_lang=source_lang_nllb, tgt_lang=target_code, device=hyperparameters.device)

    # output = translator(text, max_length=400)[0]['translation_text']
    if sentence_mode == "Sentence-wise":
        sentences = sent_tokenize(text)
        translated_sentences = []
        for sentence in sentences:
            translated_sentence = translator(sentence, max_length=400)[0]['translation_text']
            translated_sentences.append(translated_sentence)
        output = ' '.join(translated_sentences)
    else:
        output = translator(text, max_length=1024)[0]['translation_text']

    # get the text to speech 
    fs_out, audio_out = text_to_speech(output, target_code_nllb)

    if output_type == 'own voice':
        out_file = modelConvertSpeech.convert((fs_out, audio_out), audio)
        return output, out_file

    wavfile.write(constants.text2speech_wavfile, fs_out, audio_out)
    return output, constants.text2speech_wavfile

with gr.Blocks(title = "Octopus Translation App") as octopus_translator:
    with gr.Row():
        audio_file = gr.Audio(source="microphone")
    
    with gr.Row():
        input_text  = gr.Textbox(label="Input text")
        source_language     = gr.Dropdown(list(constants.flores_codes.keys()), value='English', label='Source (Autoselected)', interactive=True)
    
    with gr.Row():
        output_text = gr.Textbox(label='Translated text')
        target_language  = gr.Dropdown(list(constants.flores_codes.keys()), value='German', label='Target', interactive=True)
    
    
    with gr.Row():
        output_speech = gr.Audio(label='Translated speech')
        translate_button = gr.Button('Translate')

    
    with gr.Row():
        enhance_audio       = gr.Radio(['yes', 'no'], value='yes', label='Enhance input voice', interactive=True)
        input_type          = gr.Radio(['Whole text', 'Sentence-wise'],value='Sentence-wise', label="Translation Mode", interactive=True)
        output_audio_type   = gr.Radio(['standard speaker', 'voice transfer'], value='voice transfer', label='Enhance output voice', interactive=True)

    audio_file.change(speech_to_text,
                      inputs=[audio_file],
                      outputs=[input_text, source_language])

    translate_button.click(translation,
                           inputs=[audio_file, input_text, 
                                   source_language, target_language, 
                                   output_audio_type, input_type],
                           outputs=[output_text, output_speech])

octopus_translator.launch(share=False)