import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import numpy as np
import torch
import arxiv


def main():
    id_provided = False

    st.set_page_config(
        layout="wide",
        initial_sidebar_state="auto",
        page_title="Political Science Title Generator!",
        page_icon=None,
    )

    st.title("Generate Title from Abstract of a Political Science Paper")
    st.text("")
    st.text("")

    # Take the message which needs to be processed
    message = st.text_area("Paste a paper's abstract to generate a title", height=12)

    st.text("")
    models_to_choose = [
        "ey211/mt5-base-finetuned-dimensions-polisci",
    ]

    BASE_MODEL = st.selectbox("Choose a model to generate the title", models_to_choose)

    def preprocess(text):
        if (BASE_MODEL == "ey211/mt5-base-finetuned-dimensions-polisci"):
            return [text]
        else:
            st.error("Please select a model first")

    @st.cache(allow_output_mutation=True, suppress_st_warning=True, show_spinner=False)
    def load_model():
        tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
        model = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL)
        return model, tokenizer

    def get_summary(text):
        with st.spinner(text="Processing your request"):
            model, tokenizer = load_model()
            preprocessed = preprocess(text)
            inputs = tokenizer(
                preprocessed, truncation=True, padding="longest", return_tensors="pt"
            )
            output = model.generate(
                **inputs,
                max_length=256,
                num_beams=10,
                num_return_sequences=1,
                temperature=1.5,
            )
            target_text = tokenizer.batch_decode(output, skip_special_tokens=True)
            return target_text[0]

    # Define function to run when submit is clicked
    def submit(message):
        if len(message) > 0:
            summary = get_summary(message)
            if id_provided:
                html_str = f"""
                <style>
                p.a {{
                font: 30px Courier;
                }}
                </style>
                <p class="a"><b>Title Generated:></b> {summary} </p> 
                <p class="a"><b>Original Title:></b> {title} </p>
                """
            else:
                html_str = f"""
                <style>
                p.a {{
                font: 30px Courier;
                }}
                </style>
                <p class="a">Title Generated:> <b>{summary} </b></p>
                """

            st.markdown(html_str, unsafe_allow_html=True)
            # st.markdown(emoji)
        else:
            st.error("The text can't be empty")

    # Run algo when submit button is clicked
    if st.button("Submit"):
        submit(message)

    with st.expander("Additional Information"):
        st.markdown("""
        The model used was fine-tuned on title and abstract data from political science papers from [Dimensions](https://dimensions.ai).
        The task of the models is to suggest an appropraite title from the abstract of a scientific paper.
        """,unsafe_allow_html=True,)

    st.text('\n')
    st.text('\n')
    st.markdown(
    '''<span style="color:blue; font-size:10px">App created by [@ey211](https://huggingface.co/ey211)
      </span>''', 
     unsafe_allow_html=True,
    )

if __name__ == "__main__":
    main()