import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load both translation models from Hugging Face
tokenizer_eng_to_darija = AutoTokenizer.from_pretrained("lachkarsalim/Helsinki-translation-English_Moroccan-Arabic")
model_eng_to_darija = AutoModelForSeq2SeqLM.from_pretrained("lachkarsalim/Helsinki-translation-English_Moroccan-Arabic")

tokenizer_darija_to_msa = AutoTokenizer.from_pretrained("Saidtaoussi/AraT5_Darija_to_MSA")
model_darija_to_msa = AutoModelForSeq2SeqLM.from_pretrained("Saidtaoussi/AraT5_Darija_to_MSA")

# Translation functions
def translate_darija_to_msa(darija_text):
    inputs = tokenizer_darija_to_msa(darija_text, return_tensors="pt", padding=True)
    translated = model_darija_to_msa.generate(**inputs)
    translated_text = tokenizer_darija_to_msa.decode(translated[0], skip_special_tokens=True)
    return translated_text

def translate_eng_to_darija(eng_text, direction="eng_to_darija"):
    if direction == "eng_to_darija":
        inputs = tokenizer_eng_to_darija(eng_text, return_tensors="pt", padding=True)
        translated = model_eng_to_darija.generate(**inputs)
        translated_text = tokenizer_eng_to_darija.decode(translated[0], skip_special_tokens=True)
    else:
        # Reverse translation (Darija to English)
        inputs = tokenizer_eng_to_darija(eng_text, return_tensors="pt", padding=True)
        translated = model_eng_to_darija.generate(**inputs)
        translated_text = tokenizer_eng_to_darija.decode(translated[0], skip_special_tokens=True)
    return translated_text

# Respond function
def respond(message, translation_choice):
    if translation_choice == "Moroccan Arabic to MSA":
        return translate_darija_to_msa(message)
    elif translation_choice == "English to Moroccan Arabic":
        return translate_eng_to_darija(message, direction="eng_to_darija")

# Gradio Interface Layout with organized components
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            # Header with emojis and logo
            gr.Markdown("""
                <h1 style="text-align: center;">
                    🇲🇦 🚀 Moroccan Arabic Translation Demo 😍
                </h1>
                <p style="text-align: center;">
                    🌟 Dima Meghrib 🌟 <br>
                    Select the translation direction and type your text. <br>
                    Get quick translations between **English** and **Moroccan Arabic (Darija)** or **Darija to Modern Standard Arabic (MSA)**! 🔥
                </p>
                <p style="text-align: center;">
                    This demo uses two advanced models:<br>
                    - English to Moroccan Arabic (Darija)<br>
                    - Moroccan Arabic (Darija) to Modern Standard Arabic (MSA)<br>
                    Choose your desired translation direction and get started!<br>
                </p>
                <p style="text-align: center;">
                    <img src="https://moroccan-culture-image.s3.eu-north-1.amazonaws.com/2159558.png" 
                    style="width: 150px; display: block; margin: 20px auto;" alt="Moroccan Flag" />
                </p>
            """)
        # Translation Inputs and Outputs
        user_input = gr.Textbox(label="Enter Your Text", placeholder="Type your sentence here...")
        translation_choice = gr.Dropdown(
            label="Choose Translation Direction",
            choices=["English to Moroccan Arabic", "Moroccan Arabic to MSA"],
            value="English to Moroccan Arabic"
        )
        submit_button = gr.Button("Submit", elem_id="submit_button")
        
    with gr.Row():
        # Output area for translated text
        output = gr.Textbox(label="Translated Text", placeholder="Translation will appear here...")

    # Footer with your name at the bottom
    gr.Markdown("<p style='text-align: center; font-size: 14px;'>Created by Eng Amal 🌟</p>")

    # Define the action for submit
    submit_button.click(fn=respond, inputs=[user_input, translation_choice], outputs=output)

# Launch the interface
demo.launch()