import streamlit as st
from transformers import pipeline
import torch

# Function to load the model with error handling and ensure it's available
@st.cache_resource
def load_model():
    try:
        # Attempt to load a model
        st.write("Attempting to load the emotion analysis model...")
        emotion_analyzer = pipeline("text-classification", model="bhadresh-savani/bert-base-uncased-emotion", device=0 if torch.cuda.is_available() else -1)
        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()

# Check if the model is loaded
if emotion_analyzer is None:
    st.warning("The emotion analysis model could not be loaded. Please try again.")
else:
    st.success("Emotion model is ready for predictions!")

# Function to predict emotion for a single response
def predict_emotion_single(response):
    if emotion_analyzer is None:
        st.error("Model not loaded. Please try reloading the app.")
        return {"Error": "Emotion analyzer model not initialized. Please check model loading."}
    
    try:
        # Analyze the emotion of the response
        result = emotion_analyzer([response])
        return {res["label"]: round(res["score"], 4) for res in result}
    except Exception as e:
        st.error(f"Prediction failed: {e}")
        return {"Error": f"Prediction failed: {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.")