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

# 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 = """
<div class="title">English-Urdu Translation</div>
<div class="description">This app leverages advanced Transformer models and custom-trained tokenizers to deliver accurate and context-aware translations.</div>
"""
article = "<p style='text-align: center'>Created by: <a href='https://www.linkedin.com/in/syedhuzaifanafees/' target='_blank'>Syed Huzaifa Nafees</a> | Supervisor: <a href='https://www.linkedin.com/in/shaider/' target='_blank'>Dr. Sajjad Haider</a></p></center></p>"
examples = [
    ["the weather is good today", "en->ur"],
    ["آج اچھا دن ہے", "ur->en"]
]
iface = gr.Interface(
    fn=translate,
    #title="English-Urdu Translation",
    description=description,
    article=article,
    examples=examples,
    inputs=[gr.Textbox(lines=5, placeholder="Enter text (maximum 35 words)", label="Input"),
        gr.Radio(choices=['en->ur','ur->en'], value='en->ur', label='Select the direction of translation'),],
    outputs="text",
    css= 'style.css')

iface.launch()