import streamlit as st
import arxiv
from inference import top_95_labels, load_model  # Replace with your actual module

@st.cache_resource
def get_model():
    return load_model("kalinin-a-i/ml2-hw4", "class2name.joblib")

def get_arxiv_article_info(url):
    """Extracts title and abstract from an arXiv article link."""
    paper_id = url.split("/abs/")[-1]
    search = arxiv.Search(id_list=[paper_id])
    
    for result in search.results():
        return result.title, result.summary
    return None, None

# Load model once
pipe, class2name = get_model()

st.title("ArXiv Article Classifier")

# Input method selection
input_method = st.radio("Select input method:", ("ArXiv Link", "Manual Input"))

title, abstract = None, None

if input_method == "ArXiv Link":
    arxiv_url = st.text_input("Enter ArXiv article link:")
    if st.button("Classify"):
        if arxiv_url:
            title, abstract = get_arxiv_article_info(arxiv_url)
        else:
            st.error("Please enter a valid ArXiv link.")
elif input_method == "Manual Input":
    title = st.text_input("Enter Article Title:")
    abstract = st.text_area("Enter Article Abstract:")
    if st.button("Classify"):
        if not title:
            st.error("Please provide title")

# Classification and output
if title:
    categories = top_95_labels(pipe, class2name, title, abstract)  # Assuming classify returns a list of labels
    st.write(f"### Title: {title}")
    st.write("**Predicted Categories:**")
    for category in categories:
        st.write(f"- {category}")