import streamlit as st
from transformers import pipeline
import torch

# ---- Page Configuration ----
st.set_page_config(
    page_title="Emotion Prediction App",
    page_icon="🌟",
    layout="centered",
    initial_sidebar_state="expanded",
)

# ---- App Title ----
st.title("🌟 Emotion Prediction App 🌈")
st.subheader("Understand your emotions better with AI-powered predictions!")

# ---- Function to Load Emotion Analysis Model ----
@st.cache_resource
def load_emotion_model():
    try:
        st.info("⏳ Loading the emotion analysis model, please wait...")
        # Using a publicly available model for emotion analysis
        emotion_analyzer = pipeline(
            "text-classification",
            model="bhadresh-savani/distilbert-base-uncased-emotion",  # A valid public model
            device=0 if torch.cuda.is_available() else -1,  # Use GPU if available
        )
        st.success("✅ Model loaded successfully!")
        return emotion_analyzer
    except Exception as e:
        st.error(f"⚠️ Error loading model: {e}")
        return None

# ---- Load the Model ----
emotion_analyzer = load_emotion_model()

# ---- Function for Predicting Emotion ----
def predict_emotion(text):
    if emotion_analyzer is None:
        st.error("⚠️ Model not loaded. Please reload the app.")
        return {"Error": "Emotion analyzer model not initialized. Please try again later."}
    
    try:
        # Analyze emotions
        result = emotion_analyzer([text])
        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}"}

# ---- Suggesting Activities Based on Emotional State ----
def suggest_activity(emotion_analysis):
    # Suggest activities based on emotional state
    max_emotion = max(emotion_analysis, key=emotion_analysis.get) if emotion_analysis else None
    if max_emotion == 'sadness':
        return "It's okay to feel sad sometimes. Try taking a 5-minute mindfulness break or a short walk outside to clear your mind."
    elif max_emotion == 'joy':
        return "You’re feeling happy! Keep that positive energy going. How about practicing some deep breathing exercises to maintain your balance?"
    elif max_emotion == 'fear':
        return "Feeling anxious? It might help to do a quick breathing exercise or take a walk to ease your mind."
    elif max_emotion == 'anger':
        return "It seems like you're angry. Try taking a few deep breaths, or engage in a relaxing mindfulness exercise to calm your nerves."
    elif max_emotion == 'surprise':
        return "You’re surprised! Take a moment to breathe deeply and reflect. A walk might help clear your thoughts."
    elif max_emotion == 'disgust':
        return "If you’re feeling disgusted, a change of environment might help. Go for a walk or try a mindfulness technique to relax."
    elif max_emotion == 'sadness':
        return "It’s okay to feel sad. Try grounding yourself with some mindfulness or a breathing exercise."
    else:
        return "Keep doing great! If you feel overwhelmed, consider taking a deep breath or going for a short walk."

# ---- User Input Section ----
st.write("### 🌺 Let's Get Started!")
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?",
]

responses = {}

# ---- Ask Questions and Analyze Responses ----
for i, question in enumerate(questions, start=1):
    st.write(f"#### ❓ Question {i}: {question}")
    user_response = st.text_input(f"Your answer to Q{i}:", key=f"q{i}")
    
    if user_response:
        with st.spinner("Analyzing emotion... 🎭"):
            analysis = predict_emotion(user_response)
        responses[question] = {"Response": user_response, "Analysis": analysis}
        
        # Display Emotion Analysis
        st.success(f"🎯 Emotion Analysis: {analysis}")
        
        # Display Activity Suggestion
        if analysis:
            max_emotion = max(analysis, key=analysis.get)
            activity_suggestion = suggest_activity(analysis)
            st.write(f"### 🧘 Suggested Activity: {activity_suggestion}")

# ---- Display Results ----
if st.button("Submit Responses"):
    st.write("### 📊 Emotion Analysis Results")
    if responses:
        for i, (question, details) in enumerate(responses.items(), start=1):
            st.write(f"#### Question {i}: {question}")
            st.write(f"**Your Response:** {details['Response']}")
            st.write(f"**Emotion Analysis:** {details['Analysis']}")
            activity_suggestion = suggest_activity(details["Analysis"])
            st.write(f"**Suggested Activity:** {activity_suggestion}")
    else:
        st.warning("Please answer at least one question before submitting!")

# ---- Footer ----
st.markdown(
    """
    ---
    **Developed using 🤗 Transformers**  
    Designed for a fun and intuitive experience! 🌟  
    """
)

# ---- Error Handling and User Suggestions ----
if emotion_analyzer is None:
    st.error("⚠️ We couldn't load the emotion analysis model. Please check your internet connection or try again later.")
    st.markdown("🔧 **Troubleshooting Steps:**")
    st.markdown("1. Ensure you have a stable internet connection.")
    st.markdown("2. If the issue persists, please refresh the page and try again.")
    st.markdown("3. Check if the model has been updated or is temporarily unavailable.")