import time
import transformers
import pandas as pd
import streamlit as st

paper_contents: str = ""
generated_title: str = ""
maximum_tokens: int = 10
preferred_model: str = ""

AVAILABLE_MODELS = [
    "TohidaRehman/pegasus-large-Abstract-Title-CSPubSum",
    "TohidaRehman/Llama-3-8b-Abstract-Title-CSPubSum",
    "TohidaRehman/t5-base-Abstract-Title",
    "czearing/article-title-generator"
]

def generate_title(input_text: str, model_name: str, max_length: int = 20) -> str:
    prefix: str = "summarize: "
    text_with_prefix = prefix + input_text

    try:
        tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
        model = transformers.AutoModelForSeq2SeqLM.from_pretrained(model_name)

        inputs = tokenizer(
            text_with_prefix, 
            return_tensors='pt',
            max_length=512, 
            truncation=True, 
            padding=True
        )

        predictions = model.generate(
            input_ids=inputs['input_ids'],
            attention_mask=inputs['attention_mask'],
            max_length=max_length,
            num_beams=4,
            do_sample=True,
            min_length=3,
        )

        summary = tokenizer.decode(predictions[0], skip_special_tokens=True)
        return summary
    except Exception as e:
        return str(e)

st.set_page_config(
    layout="wide",
    page_title="Research Paper Title Generator", 
    page_icon=":sun_behind_cloud:",
)

st.markdown("""
    <style type="text/css">
    * {
        border-radius: 0px !important;
    }

    div[role="slider"] {
        border-radius: 50% !important;
    }

    .stSpinner i {
        border-radius: 50% !important;
    }

    .block-container {
        padding-bottom: 1rem !important;
    }
    </style>
""", unsafe_allow_html=True)

col1, col2 = st.columns([2, 1], gap="medium")
col1.subheader("Research Paper Title Generator")
col2.subheader("Generated Title")

with col1.form(key='research_paper_title_generation_parameters'):
    paper_contents = st.text_area('Paper Contents', key="paper_contents", value="", height=220, placeholder="Paste Paper Contents Here")
    preferred_model = st.selectbox("Select Preferred Model", AVAILABLE_MODELS, key="preferred_model")
    maximum_tokens = st.slider("Maximum Tokens", key="maximum_tokens", value=10, min_value=3, max_value=20, step=1)
    submitted = st.form_submit_button(label='Generate Title', type="primary")
    if submitted: 
        with col2:
            col2.write("")
            with st.spinner(text="In progress..."):
                generated_title = generate_title(paper_contents, preferred_model, maximum_tokens)
            col2.write(generated_title)