import os
import h5py
from config import config
import tensorflow as tf
from transformer import Transformer
from translator import Translator
from load_model import load_transformer, load_sp_model, load_emb

import gradio as gr

import warnings
warnings.filterwarnings("ignore")

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 
tf.get_logger().setLevel('ERROR')

# load the embedding matrix
eng_emb_path = 'Embedding_matrix/en_embedding_matrix.h5'
ur_emb_path = 'Embedding_matrix/ur_embedding_matrix.h5'

eng_embedding_matrix = load_emb(eng_emb_path)
urdu_embedding_matrix = load_emb(ur_emb_path)

# load tokenizers
spm_path_en = 'Tokenizer/mix_en_spm.model'
spm_path_ur = 'Tokenizer/mix_ur_spm_rev.model'
sp_model_en, sp_model_ur = load_sp_model(spm_path_en,spm_path_ur)

# load the transformer models
E2U_model = load_transformer(eng_embedding_matrix,urdu_embedding_matrix,'T_mix_E2U_weights/T_mix_E2U.tf', config)
U2E_model = load_transformer(urdu_embedding_matrix,eng_embedding_matrix,'T_mix_U2E_weights/T_mix_U2E.tf', config)

U2E_translator = Translator(sp_model_ur, sp_model_en, U2E_model)
E2U_translator = Translator(sp_model_en, sp_model_ur, E2U_model)

def translate(inp, direction):
    if direction == 'en->ur':
        # Translate from English to Urdu
        translated_text = E2U_translator(inp)
    else:
        # Translate from Urdu to English
        translated_text = U2E_translator(inp)

    return translated_text


description = """
<p>
<center>
This app is built as part of an MS project, the app leverages advanced Transformer models and custom-trained tokenizers to deliver accurate and context-aware translations
</center>
</p>
"""
article = "<p style='text-align: center'><a href='https://www.linkedin.com/in/syedhuzaifanafees/' target='_blank'>by Syed Huzaifa Nafees</a> | <a href='https://www.linkedin.com/in/shaider/' target='_blank'>Supervisor: Dr. Sajjad Haider</a> | Contact: <a href='mailto:huzaifahtu@gmail.com' target='_blank'>Huzaifa</a></p></center></p>"
examples = [
    ["This is a cool application", "en->ur"],
    ["آج اچھا دن ہے", "ur->en"]
]
iface = gr.Interface(
    fn=translate,

    title="English-Urdu Translation",
    description=description,
    article=article,
    examples=examples,
    inputs=[
        gr.inputs.Textbox(lines=5, placeholder="Enter text (maximum 35 words)", label="Input"),
        gr.inputs.Radio(
            choices=[
                'en->ur',
                'ur->en'],
            default='en->ur',
            label='Direction'),
    ],
    outputs="text")

iface.launch(enable_queue=True)