import streamlit as st
from huggingface_hub import hf_hub_download
from utils import *

# APP
st.set_page_config(
    page_title="Hengam",
    page_icon="🕒",
)

# Layout adjustments
st.markdown(
    """
<style>
.reportview-container .main .block-container {
    max-width: 1400px;
}
</style>
""",
    unsafe_allow_html=True,
)

st.markdown("## **🕒 Hengam**")
st.write("")
with st.expander("ℹ️ - About this app", expanded=True):
    st.write(
        """     
- Online Demo for Hengam: An Adversarially Trained Transformer for Persian Temporal Tagging [Code](https://github.com/kargaranamir/hengam)!
- This paper introduces Hengam, an adversarially trained transformer for Persian temporal tagging outperforming state-of-the-art approaches on a diverse and manually created dataset.
    """
    )
    st.markdown("")

    
# Load models lazily
@st.cache(allow_output_mutation=True)
def load_ner_model(model_path):
    return NER(model_path=model_path, tags=['B-TIM', 'I-TIM', 'B-DAT', 'I-DAT', 'O'])

with st.spinner(text="Please wait while the model is loading...."):

    # Download the models
    HengamTransW = hf_hub_download(repo_id="kargaranamir/Hengam", filename="HengamTransW.pth")
    HengamTransA = hf_hub_download(repo_id="kargaranamir/Hengam", filename="HengamTransA.pth")

    # cache
    load_ner_model(HengamTransW)
    load_ner_model(HengamTransA)


st.markdown("")
st.markdown("## **📌 Paste any Persian (Farsi) text you want to extract its temporal markers.**")
with st.form(key="my_form"):
    c1, c2 = st.columns([2, 3])
    
    with c1:
        model_paths = {
            "HengamTransW.pth": HengamTransW,
            "HengamTransA.pth": HengamTransA,
        }
        
        default_model = "HengamTransA.pth"  # Set the default model
        ModelType = st.selectbox(
            "Choose your model",
            list(model_paths.keys()),
            index=list(model_paths.keys()).index(default_model),  # Select the default model
            help="At present, you can choose between 2 models (HengamTransW or HengamTransA) to extract temporal markers. More to come!",
        )
        
        ner = load_ner_model(model_paths[ModelType])

    with c2:
        doc = st.text_area(
            "Paste your text below",
            "Example: ساعت ۸ صبح من و علی قرار گذاشتیم که به دوشنبه بازار بریم ...",
            height=80,
        )

        submit_button = st.form_submit_button(label="✨ Extract Temporal Markers!")

if submit_button:
    result = ner("# " + doc)
    st.write("")  # Add vertical spacing
    st.markdown("**Result:**")
    st.code(result, language="python")