import os
import sys
#os.system('pip install phonemizer')
#os.system('apt-get install espeak-ng')
sys.path.insert(0, "./hifi-gan")
sys.path.insert(0, "./weights")

#for comptability attrdict with python 3.10 and above
import collections
import collections.abc
for type_name in collections.abc.__all__:
    setattr(collections, type_name, getattr(collections.abc, type_name))
#===================================================

import nltk
from nltk.tokenize import word_tokenize
import re
import librosa
import yaml
import gradio as gr
import torch
import glob
from attrdict import AttrDict
import phonemizer
import json
from vocoder import Generator
from models import *
from utils import *
from ipa_uk.ipa_uk import ipa
import unicodedata
from ukrainian_word_stress import Stressifier, StressSymbol

nltk.download('punkt')
nltk.download('punkt_tab')
#global_phonemizer = phonemizer.backend.EspeakBackend(language='uk', preserve_punctuation=True,  with_stress=True)
device = 'cuda' if torch.cuda.is_available() else 'cpu'


# Text Cleaner ===================================================
_pad = "$"
_punctuation = ';:,.!?¡¿—…"«»“” '
_letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
_letters_ipa = "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→'̩'ᵻ"


# Export all symbols:
symbols = [_pad] + list(_punctuation) + list(_letters) + list(_letters_ipa)

dicts = {}
for i in range(len((symbols))):
    dicts[symbols[i]] = i

class TextCleaner:
    def __init__(self, dummy=None):
        self.word_index_dictionary = dicts
    def __call__(self, text):
        indexes = []
        for char in text:
            try:
                indexes.append(self.word_index_dictionary[char])
            except KeyError:
                print(char)
        return indexes

textclenaer = TextCleaner()
#===================================================

def load_checkpoint(filepath, device):
    assert os.path.isfile(filepath)
    print("Loading '{}'".format(filepath))
    checkpoint_dict = torch.load(filepath, map_location=device)
    print("Complete.")
    return checkpoint_dict

def scan_checkpoint(cp_dir, prefix):
    pattern = os.path.join(cp_dir, prefix + '*')
    cp_list = glob.glob(pattern)
    if len(cp_list) == 0:
        return ''
    return sorted(cp_list)[-1]
    

def greet(name):
    return "Hello " + name + "!!"

# Load HiFi-GAN
cp_g = scan_checkpoint("Vocoder/", 'g_')
config_file = os.path.join(os.path.split(cp_g)[0], 'config.json')
with open(config_file) as f:
    data = f.read()
json_config = json.loads(data)
h = AttrDict(json_config)

device = torch.device(device)
generator = Generator(h).to(device)

state_dict_g = load_checkpoint(cp_g, device)
generator.load_state_dict(state_dict_g['generator'])
generator.eval()
generator.remove_weight_norm()
#===================================================

# load TTS Model 
model_path = "./Models/UK/second_stage.pth" 
model_config_path = "./Models/UK/config_uk.yml"
config = yaml.safe_load(open(model_config_path))

# 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)

model = build_model(Munch(config['model_params']), text_aligner, pitch_extractor)
params = torch.load(model_path, map_location=device)
params = params['net']
for key in model:
    if key in params:
        if not "discriminator" in key:
            print('%s loaded' % key)
            model[key].load_state_dict(params[key])
_ = [model[key].eval() for key in model]
_ = [model[key].to(device) for key in model]
#===================================================

# load embbeding refrence
reference_embeddings = np.load("reference_embeddings.npy", allow_pickle=True).item()
#===================================================

def remove_combining_diacritics(input_str):
    return ''.join(char for char in unicodedata.normalize('NFD', input_str) if unicodedata.category(char) != 'Mn')

def tts(text: str):
    stressify = Stressifier(stress_symbol = '\u02C8')

    #======================================================
    text = text.strip()
    text = text.replace('"', '')
    text = text.replace('+', 'ˈ')
    
    text = re.sub(r'[{}()_*]', '', text) # remove special symbols
    text = re.sub(r'([!?.])[\./]+', r'\1', text) # Replace test.? with test?
    text = text.replace('-', '')
    
    text = re.sub(r'[᠆‐‑‒–—―⁻₋−⸺⸻]', '—', text)
    text = re.sub(r' - ', ': ', text)
    #======================================================
    
    # tokenize

    ps = [ipa(stressify(text), check_accent=False)]
    ps = [remove_combining_diacritics(''.join(ps))]
    
    #ps = global_phonemizer.phonemize([text])
    
    ipa2uk_phonemes = ps[0]
    print(ps)
    
    ps = word_tokenize(ps[0])
    ps = ' '.join(ps)
    tokens = textclenaer(ps)
    tokens.insert(0, 0)
    tokens.append(0)
    tokens = torch.LongTensor(tokens).to(device).unsqueeze(0)
    #===================================================

    # Synthesize speech
    converted_samples = {}
    
    with torch.no_grad():
        input_lengths = torch.LongTensor([tokens.shape[-1]]).to(device)
        m = length_to_mask(input_lengths).to(device)
        t_en = model.text_encoder(tokens, input_lengths, m)
            
        for key, (ref, _) in reference_embeddings.items():
            
            s = ref.squeeze(1)
            style = s
            
            d = model.predictor.text_encoder(t_en, style, input_lengths, m)
            
            x, _ = model.predictor.lstm(d)
            duration = model.predictor.duration_proj(x) / 0.98
            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))
            style = s.expand(en.shape[0], en.shape[1], -1)
    
            F0_pred, N_pred = model.predictor.F0Ntrain(en, s)
    
            out = model.decoder((t_en @ pred_aln_trg.unsqueeze(0).to(device)), 
                                    F0_pred, N_pred, ref.squeeze().unsqueeze(0))
    
            c = out.squeeze()
            y_g_hat = generator(c.unsqueeze(0))
            y_out = y_g_hat.squeeze()
            
            converted_samples[key] = y_out.cpu().numpy()
    #===================================================    
    #return "Hello " + text + "!!"
    # Only one key in converted_samples
    key = list(converted_samples.keys())[0]
    audio_output = converted_samples[key]
    audio_output = (audio_output * 32767).astype(np.int16)
    # Return the audio waveform and sample rate (assuming 24000 Hz)
    return (24000, audio_output), ipa2uk_phonemes


with open("README.md") as file:
    article = file.read()
    article = article[article.find("---\n", 4) + 5 : :]

iface = gr.Interface(
    fn=tts,
    inputs=[
        #gr.components.Textbox(
        gr.Text(
            label="Input",
            value='''ми кр+апельки ртуті на рівному полі, на сірій безмірній пустій площині. 
            – рухливі, прудкі, досконалі і голі, котитись навчились, а жити, ще ні,
            – щенячі забави щоночі, щоднини, тваринна захланність звірячий запал,
            – хоч крапелька кожна це майже людина, і світло тремтливе відлите в метал.
            – але проступають пророцтва забуті, і стеляться світом зневіра і страх.
            – а ми розтік+аємось краплями ртуті, по мінних, - по мінних, по мінних полях. ''',
        ),
    ],
    outputs=[
        gr.components.Audio(label="Output"),
        gr.components.Textbox(label="IPA Phonems"),
    ],
    title="Ukrainian StyleTTS Alexis",
    description= f'''Ukrainian StyleTTS. If the stress is incorrect, use the + symbol before the stressed letter.''',
    cache_examples=False,
    examples=[["""Але щоб ви зрозуміли, звідки виник+ає це хибне уявлення людей, цуратись насолоди і вихваляти страждання, я розкрию перед вами всю картину і роз’ясн+ю, що саме говорив цей чоловік, який відкрив істину, якого я б назвав зодчим щасливого життя."""],
             ["""Ми кр+апельки ртуті на рівному полі, на сірій безмірній пустій площині. 
            – рухливі, прудкі, досконалі і голі, котитись навчились, а жити, ще ні,
            – щенячі забави щоночі, щоднини, тваринна захланність звірячий запал,
            – хоч крапелька кожна це майже людина, і світло тремтливе відлите в метал.
            – але проступають пророцтва забуті, і стеляться світом зневіра і страх.
            – а ми розтік+аємось краплями ртуті, по мінних, - по мінних, по мінних полях. """],
             ["""спини мене, отямся і отям,  така любов буває раз в нік+оли.
            – вона ж промчить над зламаним життям  за нею ж будуть бігти видноколи.
            – вона ж порве нам спокій до струни, вона ж слова поспалює вустами.
            – спини мене, спини і схамени,  ще поки можу думати востаннє.
            – ще поки можу, –  але вже не можу. – настала черга й на мою зорю.
            – чи біля тебе душу відморожу чи біля тебе полум'ям згорю..."""] 
             ],
    article=article,
)
iface.launch()