"""
🗣️ Translator - Translate text from one language to another.

Application file made with Streamlit.
"""

import re
import streamlit as st

from datetime import datetime
from transformers import pipeline
from available_models import MODELS


st.set_page_config(page_title="Translate", page_icon="🗣️")
st.title("Translate", "main")
st.subheader("Using cutting-edge AI/ML technology to translate text for you.")
st.markdown("""
Welcome to Translate. Here, we harness the power of cutting-edge Artificial Intelligence and Machine Learning to translate text for _you_.

Translate is a free tool to translate quickly and easily! Start today by typing your text into the form below and selecting a language pair!
""")
hide_streamlit_style = """
            <style>
            #MainMenu {visibility: hidden;}
            footer {visibility: hidden;}
            </style>
            """
st.markdown(hide_streamlit_style, unsafe_allow_html=True) 
lang1, lang2 = st.columns(2)
lang1.selectbox(
    "Source Language", ["🇺🇸 English", "🇫🇷 French", "🇩🇪 German", "🇪🇸 Spanish", "🇨🇳 Chinese"], 
    key="input_lang", index=0,
)
lang2.selectbox(
    "Target Language", ["🇺🇸 English", "🇫🇷 French", "🇩🇪 German", "🇪🇸 Spanish", "🇨🇳 Chinese"], 
    key="output_lang", index=3,
)

selected_model = MODELS[f"{st.session_state['input_lang']}->{st.session_state['output_lang']}"]


if selected_model[0] == None:
    st.write("Sorry, we can't translate this language pair! Please try a different one!")
elif selected_model[0] == 0:
    st.write("Haha! We know what you're trying to do and you can't do it! You're not allowed to translate languages from one language to the same one!")
else:
    input_text = st.text_area("Enter text to translate:", height=200, key="input")
    translate_text = st.button("Translate")

    if translate_text:
        with st.spinner(text="Just a sec, it takes a minute to load up the AI/ML!"):
            task = pipeline(
                "translation",
                model=selected_model[0],
                tokenizer=selected_model[0],
            )

        progress_bar = st.progress(0)
        with st.spinner(text="Just a little bit more time! We wish it was easy, but AI takes some time!"):
            text_to_translate = re.split('(?<=[.!?]) +', input_text)
            total_progress = len(text_to_translate)

            for i, text in enumerate(text_to_translate):
                translation = task(text)
                text_to_translate[i] = translation[0]["translation_text"]
                progress_bar.progress((i + 1) / total_progress)

        st.success("Yay! It's done! Check out your translation below!")
        st.write("**Here's your translation:**")
        st.code(' '.join(text_to_translate))