# Import necessary libraries
import streamlit as st
from transformers import pipeline
import torch
from datasets import load_dataset

# Load the T5-based Emotion Classifier model
@st.cache_resource
def load_model():
    try:
        st.write("Loading the emotion analysis model...")
        emotion_analyzer = pipeline("text-classification", model="suryakiran786/T5-emotion")
        st.write("Model loaded successfully!")
        return emotion_analyzer
    except Exception as e:
        st.write(f"Error loading the model: {e}")
        return None

# Initialize the model (with caching to prevent reloads)
emotion_analyzer = load_model()

# Load the dataset if needed for any additional logic
@st.cache_data
def load_data():
    try:
        # For demonstration purposes, let's load a sentiment analysis dataset from Hugging Face
        dataset = load_dataset("glue", "sst2")
        st.write("Dataset loaded successfully!")
        return dataset
    except Exception as e:
        st.write(f"Error loading dataset: {e}")
        return None

# Load data (just to show usage, not used in emotion analysis directly)
dataset = load_data()

# Function to predict emotion for a single response
def predict_emotion_single(response):
    if not emotion_analyzer:
        return {"Error": "Emotion analyzer model not initialized. Please check model loading."}
    try:
        response = response.strip()
        result = emotion_analyzer([response])
        return {res["label"]: round(res["score"], 4) for res in result}
    except Exception as e:
        return {"Error": str(e)}

# Streamlit App Layout
st.title("Behavior Prediction App")
st.write("Enter your thoughts or feelings, and let the app predict your emotional states.")

# Define questions for the user
questions = [
    "How are you feeling today?",
    "Describe your mood in a few words.",
    "What was the most significant emotion you felt this week?",
    "How do you handle stress or challenges?",
    "What motivates you the most right now?"
]

# Initialize a dictionary to store responses
responses = {}

# Ask each question and get response
for i, question in enumerate(questions, start=1):
    user_response = st.text_input(f"Question {i}: {question}")
    if user_response:
        analysis = predict_emotion_single(user_response)
        responses[question] = (user_response, analysis)
        st.write(f"**Your Response**: {user_response}")
        st.write(f"**Emotion Analysis**: {analysis}")

# Provide button to clear input fields
if st.button("Clear Responses"):
    st.experimental_rerun()

# Display results once all responses are filled
if st.button("Submit Responses"):
    if responses:
        st.write("-- Emotion Analysis Results ---")
        for i, (question, (response, analysis)) in enumerate(responses.items(), start=1):
            st.write(f"\n**Question {i}:** {question}")
            st.write(f"Your Response: {response}")
            st.write(f"Emotion Analysis: {analysis}")
    else:
        st.write("Please answer all the questions before submitting.")