import streamlit as st
from transformers import pipeline

# Load emotion classification model
@st.cache_resource
def load_model():
    try:
        emotion_classifier = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base")
        return emotion_classifier
    except Exception as e:
        st.error(f"Error loading model: {str(e)}")  # Error handling
        return None

emotion_classifier = load_model()

# Well-being suggestions based on emotions
def get_well_being_suggestions(emotion):
    suggestions = {
        "joy": {
            "text": "You're feeling joyful! Keep the positivity going.",
            "links": [
                "https://www.nih.gov/health-information/emotional-wellness-toolkit",
                "https://www.health.harvard.edu/health-a-to-z",
                "https://www.helpguide.org/mental-health/meditation/mindful-breathing-meditation"
            ],
            "videos": [
                "https://youtu.be/m1vaUGtyo-A",
                "https://youtu.be/MIc299Flibs"
            ]
        },
        "anger": {
            "text": "You're feeling angry. Take a moment to calm down.",
            "links": [
                "https://www.helpguide.org/mental-health/anxiety/tips-for-dealing-with-anxiety",
                "https://www.helpguide.org/mental-health/meditation/mindful-breathing-meditation"
            ],
            "videos": [
                "https://youtu.be/m1vaUGtyo-A",
                "https://www.youtube.com/shorts/fwH8Ygb0K60?feature=share"
            ]
        },
        "sadness": {
            "text": "You're feeling sad. It's okay to take a break.",
            "links": [
                "https://www.nih.gov/health-information/emotional-wellness-toolkit",
                "https://www.helpguide.org/mental-health/anxiety/tips-for-dealing-with-anxiety"
            ],
            "videos": [
                "https://youtu.be/-e-4Kx5px_I",
                "https://youtu.be/Y8HIFRPU6pM"
            ]
        },
        "fear": {
            "text": "You're feeling fearful. Try some relaxation techniques.",
            "links": [
                "https://www.helpguide.org/mental-health/anxiety/tips-for-dealing-with-anxiety",
                "https://www.health.harvard.edu/health-a-to-z"
            ],
            "videos": [
                "https://www.youtube.com/shorts/Tq49ajl7c8Q?feature=share",
                "https://youtu.be/yGKKz185M5o"
            ]
        },
        "disgust": {
            "text": "You're feeling disgusted. Take a deep breath and refocus.",
            "links": [
                "https://www.health.harvard.edu/health-a-to-z",
                "https://www.helpguide.org/mental-health/anxiety/tips-for-dealing-with-anxiety"
            ],
            "videos": [
                "https://youtu.be/MIc299Flibs",
                "https://youtu.be/-e-4Kx5px_I"
            ]
        },
        # New addition for "Surprise"
        "surprise": {
            "text": "You're feeling surprised. Take a deep breath and ground yourself.",
            "links": [
                "https://www.helpguide.org/mental-health/anxiety/tips-for-dealing-with-anxiety",
                "https://www.health.harvard.edu/health-a-to-z",
                "https://www.psychologytoday.com/us/blog/mindful-anger/201908/5-ways-to-deal-with-unexpected-surprises"
            ],
            "videos": [
                "https://youtu.be/MIc299Flibs",
                "https://www.youtube.com/shorts/Tq49ajl7c8Q?feature=share",
                "https://youtu.be/m1vaUGtyo-A"
            ]
        },
    }
    return suggestions.get(emotion, {
        "text": "Feeling neutral? That's okay! Take care of your mental health.",
        "links": [],
        "videos": []
    })

# Streamlit UI
def main():
    # Add the background image
    st.markdown("""
    <style>
    .stApp {
        background-image: url('https://drive.google.com/uc?export=view&id=1TydbjDRL5fXLMxQgaPkIg9oGOJlU2IWt');
        background-size: cover;
        background-position: center;
        color: white;  /* Set text color to white */
    }
    
    h1, h2, h3 {
        color: white;
        font-family: 'Arial', sans-serif;
    }
    
    .stTextArea textarea {
        background-color: #f2f2f2;  /* Light grey text area */
        border-radius: 8px;
        color: #333333;
    }
    
    .stButton button {
        background-color: #9b4dca;  /* Purple button */
        color: white;
        font-weight: bold;
        border-radius: 5px;
    }

    .stButton button:hover {
        background-color: #7a33a2;  /* Darker purple on hover */
    }

    .stMarkdown, .stWrite {
        color: white !important;
    }

    /* Style for links to make them more visible */
    a {
        color: #ffffff;
        text-decoration: underline;
    }
    a:hover {
        color: #ffccff;  /* Light purple hover effect */
    }
    </style>
    """, unsafe_allow_html=True)
    
    # Title of the app
    st.title("Emotion Prediction and Well-being Suggestions")
    
    # User input for emotional state
    st.header("Tell us how you're feeling today!")
    
    user_input = st.text_area("Enter a short sentence about your current mood:", "")
    
    if user_input:
        # Display Enter button only after user has entered input
        enter_button = st.button("Enter")
        
        if enter_button:
            # Clean the input text (stripping unnecessary spaces, lowercasing)
            clean_input = user_input.strip().lower()
            
            # Use the model to predict emotion
            try:
                result = emotion_classifier(clean_input)
                st.write(f"Raw Model Result: {result}")  # Debug output to see raw result
                
                emotion = result[0]['label'].lower()
                
                st.subheader(f"Emotion Detected: {emotion.capitalize()}")
                
                # Get well-being suggestions based on emotion
                suggestions = get_well_being_suggestions(emotion)
                
                # Display text suggestions
                st.write(suggestions["text"])

                # Display links
                if suggestions["links"]:
                    st.write("Useful Resources:")
                    for link in suggestions["links"]:
                        st.markdown(f"[{link}]({link})", unsafe_allow_html=True)

                # Display video links
                if suggestions["videos"]:
                    st.write("Relaxation Videos:")
                    for video in suggestions["videos"]:
                        st.markdown(f"[Watch here]({video})", unsafe_allow_html=True)
            
                # Add a button for a summary
                if st.button('Summary'):
                    st.write(f"Emotion detected: {emotion.capitalize()}. Here are your well-being suggestions to enhance your mood.")
                    st.write("Explore the links and videos to improve your emotional health!")
            
            except Exception as e:
                st.error(f"Error predicting emotion: {str(e)}")

# Run the Streamlit app
if __name__ == "__main__":
    main()