import pandas as pd
import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
from datasets import load_dataset

# Load the model and tokenizer
model_name = "modelSamLowe/roberta-base-go_emotions"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Define the emotion labels (based on the GoEmotions dataset)
emotion_labels = ["admiration", "amusement", "anger", "annoyance", "approval",
                  "caring", "confusion", "curiosity", "desire", "disappointment",
                  "disapproval", "disgust", "embarrassment", "excitement", "fear",
                  "gratitude", "grief", "joy", "love", "nervousness", "optimism",
                  "pride", "realization", "relief", "remorse", "sadness", "surprise",
                  "neutral"]

# Function to classify emotion
def classify_emotion(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    outputs = model(**inputs)
    logits = outputs.logits
    predicted_class_id = torch.argmax(logits, dim=-1).item()
    return emotion_labels[predicted_class_id]

# Streamlit interface
st.title("Enron Emails Emotion Analysis")

# Button to run the inference script
if st.button("Run Inference"):
    # Load the Enron dataset
    with st.spinner('Loading dataset...'):
        dataset = load_dataset("Hellisotherpeople/enron_emails_parsed")
        enron_data = pd.DataFrame(dataset['train'])

    # Apply emotion classification to the email content
    with st.spinner('Running inference...'):
        enron_data['emotion'] = enron_data['body'].apply(classify_emotion)

    # Save the results to a CSV file
    enron_data.to_csv("enron_emails_with_emotions.csv", index=False)
    st.success("Inference completed and results saved!")

# Check if the results file exists and load it
try:
    enron_data = pd.read_csv("enron_emails_with_emotions.csv")

    # Dropdown for selecting an emotion
    selected_emotion = st.selectbox("Select Emotion", emotion_labels)

    # Filter emails based on the selected emotion
    filtered_emails = enron_data[enron_data['emotion'] == selected_emotion].head(10)

    # Display the filtered emails in a table
    if not filtered_emails.empty:
        st.write("Top 10 emails with emotion:", selected_emotion)
        st.table(filtered_emails[['From', 'To', 'body', 'emotion']])
    else:
        st.write("No emails found with the selected emotion.")
except FileNotFoundError:
    st.warning("Run inference first by clicking the 'Run Inference' button.")