import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import fitz
import os

# Load the model and tokenizer
model = AutoModelForSequenceClassification.from_pretrained("REEM-ALRASHIDI/LongFormer-Paper-Citaion-Classifier")
tokenizer = AutoTokenizer.from_pretrained("allenai/longformer-base-4096")

def extract_text_from_pdf(file_path):
    text = ''
    with fitz.open(file_path) as pdf_document:
        for page_number in range(pdf_document.page_count):
            page = pdf_document.load_page(page_number)
            text += page.get_text()
    return text

def predict_class(text):
    try:
        # Truncate text to maximum length of 4096 tokens
        max_length = 4096
        truncated_text = text[:max_length]

        inputs = tokenizer(truncated_text, return_tensors="pt", padding=True, truncation=True, max_length=max_length)
        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits
            predicted_class = torch.argmax(logits, dim=1).item()
        return predicted_class
    except Exception as e:
        st.error(f"Error during prediction: {e}")
        return None

# Create a directory to store uploaded files
uploaded_files_dir = "uploaded_files"
os.makedirs(uploaded_files_dir, exist_ok=True)

# Define colors for different classes
class_colors = {
    0: "#1f77b4",  # Level 1
    1: "#ff7f0e",  # Level 2
    2: "#2ca02c",  # Level 3
    3: "#d62728"   # Level 4
}

st.set_page_config("Paper Citation Calssifier")
st.title("Paper Citation Classifier")

option = st.radio("Select input type:", ("Text", "PDF"))

if option == "Text":
    # Input text boxes for abstract, full text, and affiliations
    title_input = st.text_area("Enter Title:")
    abstract_input = st.text_area("Enter Abstract:")
    full_text_input = st.text_area("Enter Full Text:")
    affiliations_input = st.text_area("Enter Affiliations:")
    options=["Nursing", "Physics", "Maths", "Chemical", "Nuclear" ,"Other"]
    categories = pills("Select a category", options)
    
    # categories = st.multiselect("Select WoS categories:", ["Nursing", "Physics", "Maths", "Chemical", "Nuclear" ,"Other"])

    # Combine selected categories with [SEP]
    combined_text = f"{title_input} [SEP] {abstract_input} [SEP] {full_text_input} [SEP] {affiliations_input} [SEP] {' [SEP] '.join(categories)}"

    if st.button("Predict"):
        with st.spinner("Predicting..."):
            predicted_class = predict_class(combined_text)
            if predicted_class is not None:
                class_labels = ["Level 1 (Highly Cited Paper)", "Level 2 (Average Cited Paper)", "Level 3 (More Cited Paper)", "Level 4 (Low Cited Paper)"]

                st.text("Predicted Class:")
                for i, label in enumerate(class_labels):
                    if i == predicted_class:
                        st.markdown(
                            f'<div style="background-color: {class_colors[predicted_class]}; padding: 10px; border-radius: 5px; color: white; font-weight: bold;">{label}</div>',
                            unsafe_allow_html=True
                        )
                    else:
                        st.text(label)

elif option == "PDF":
    uploaded_file = st.file_uploader("Upload a PDF file", type=["pdf"])

    if uploaded_file is not None:
        with st.spinner("Processing PDF..."):
            file_path = os.path.join(uploaded_files_dir, uploaded_file.name)
            with open(file_path, "wb") as f:
                f.write(uploaded_file.getbuffer())
            st.success("File uploaded successfully.")
            st.text(f"File Path: {file_path}")
            
            file_text = extract_text_from_pdf(file_path)
            st.text("Extracted Text:")
            st.text(file_text)

            # Provide an option to predict from PDF text
            if st.button("Predict from PDF Text"):
                with st.spinner("Predicting..."):
                    predicted_class = predict_class(file_text)
                    if predicted_class is not None:
                        class_labels = ["Level 1 (Highly Cited Paper)", "Level 2 (Average Cited Paper)", "Level 3 (More Cited Paper)", "Level 4 (Low Cited Paper)"]
                        st.text("Predicted Class:")
                        for i, label in enumerate(class_labels):
                            if i == predicted_class:
                                st.markdown(
                                    f'<div style="background-color: {class_colors[predicted_class]}; padding: 10px; border-radius: 5px; color: white; font-weight: bold;">{label}</div>',
                                    unsafe_allow_html=True
                                )
                            else:
                                st.text(label)