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

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

# Load BERT model and tokenizer
bert_model = AutoModelForSequenceClassification.from_pretrained("Mhassanen/bert-class-test")
bert_tokenizer = AutoTokenizer.from_pretrained("Mhassanen/bert-class-test")

# Function to extract text from PDF
def extract_text_from_pdf(file_path):
    text = ''
    try:
        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()
    except Exception as e:
        st.error(f"Error reading PDF file: {e}")
    return text

# Function to predict the class of the text using a specified model and tokenizer
def predict_class(text, model, tokenizer):
    try:
        max_length = 4096 if "longformer" in str(model) else 512
        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

# Setup for uploaded files directory
uploaded_files_dir = "uploaded_files"
os.makedirs(uploaded_files_dir, exist_ok=True)

# Color mapping for class levels
class_colors = {
    0: "#d62728",  # Level 1
    1: "#ff7f0e",  # Level 2
    2: "#2ca02c",  # Level 3
    3: "#1f77b4"   # Level 4
}

# Streamlit page configuration
st.set_page_config(page_title="Paper Citation Classifier", page_icon="logo.png")

# Sidebar content
with st.sidebar:
    st.image("logo.png", width=70)
    st.markdown('<div style="position: absolute; left: 5px;"></div>', unsafe_allow_html=True)
    st.markdown("# Paper Citation Classifier")
    st.markdown("---")
    st.markdown("## About")
    st.markdown('''
    This tool classifies paper citations into different levels based on their number of citations.
    Powered by Fine-Tuned [Longformer model](https://huggingface.co/REEM-ALRASHIDI/LongFormer-Paper-Citaion-Classifier) and BERT model with custom data.
    ''')
    st.markdown("### Class Levels:")
    st.markdown("- Level 1: Highly cited papers")
    st.markdown("- Level 2: Average cited papers")
    st.markdown("- Level 3: More cited papers")
    st.markdown("- Level 4: Low cited papers")
    st.markdown("---")
    st.markdown('Tabuk University')

st.title("Check Your Paper Now!")

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

if option == "Text":
    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:")
    keywords_input = st.text_area("Enter Keywords:")
    options = ['cs', "AI"]

    selected_category = st.selectbox("Select WoS categories:", options)
    if selected_category == "Other":
        custom_category = st.text_input("Enter custom category:")
        selected_category = custom_category if custom_category else "Other"

    combined_text = f"{title_input} [SEP] {keywords_input} [SEP] {abstract_input} [SEP] {selected_category} [SEP] {affiliations_input} [SEP] {full_text_input}"

    if st.button("Predict"):
        if not any([title_input, abstract_input, keywords_input, full_text_input, affiliations_input]):
            st.warning("Please enter paper text.")
        else:
            with st.spinner("Predicting..."):
                longformer_class = predict_class(combined_text, longformer_model, longformer_tokenizer)
                bert_class = predict_class(combined_text, bert_model, bert_tokenizer)
                if longformer_class is not None and bert_class is not None:
                    class_labels = ["Level 1", "Level 2", "Level 3", "Level 4"]

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

                    st.text("BERT Predicted Class:")
                    for i, label in enumerate(class_labels):
                        if i == bert_class:
                            st.markdown(
                                f'<div style="background-color: {class_colors[bert_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)

            if st.button("Predict from PDF Text"):
                if not file_text.strip():
                    st.warning("Please upload a PDF with text content.")
                else:
                    with st.spinner("Predicting..."):
                        longformer_class = predict_class(file_text, longformer_model, longformer_tokenizer)
                        bert_class = predict_class(file_text, bert_model, bert_tokenizer)
                        if longformer_class is not None and bert_class is not None:
                            class_labels = ["Level 1", "Level 2", "Level 3", "Level 4"]
                            st.text("**Longformer Predicted Class:**")
                            for i, label in enumerate(class_labels):
                                if i == longformer_class:
                                    st.markdown(
                                        f'<div style="background-color: {class_colors[longformer_class]}; padding: 10px; border-radius: 5px; color: white; font-weight: bold;">{label}</div>',
                                        unsafe_allow_html=True
                                    )
                                else:
                                    st.text(label)

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