import streamlit as st
import requests
import os
import time  # Import time for sleep

API_URL = "https://api-inference.huggingface.co/models/ehcalabres/wav2vec2-lg-xlsr-en-speech-emotion-recognition"
headers = {"Authorization": f"Bearer {os.getenv('HF_API_TOKEN')}"}

emotion_to_neg_rate = {
    'angry': 70.0,
    'calm': 5.0,
    'disgust': 60.0,
    'fearful': 50.0,
    'happy': 5.0,
    'neutral': 20.0,
    'sad': 80.0,
    'surprised': 10.0
}

def get_emotion(audio_file):
    for _ in range(5):  # Retry up to 5 times
        response = requests.post(API_URL, headers=headers, files={"file": audio_file})
        if response.status_code == 200:
            return response.json()
        elif response.status_code == 503:  # Service Unavailable
            error_info = response.json()
            if "error" in error_info and "loading" in error_info["error"]:
                st.write("Model is currently loading. Please wait...")
                time.sleep(10)  # Wait for 10 seconds before retrying
            else:
                return {"error": response.text}
        else:
            return {"error": response.text}
    return {"error": "Model is taking too long to load or an unexpected error occurred."}

def get_neg_rate(emotion):
    return emotion_to_neg_rate.get(emotion, 30.0)

def evaluate_depression(neg_rate):
    if neg_rate <= 30.0:
        if neg_rate <= 5.0:
            return "just chill and relax"
        else:
            return "***   MILD DEPRESSION  ***"
        if neg_rate <= 10.0:
            return ("Everything is a-okay! There's absolutely nothing wrong 😁😁\n"
                    "You're probably cuddling a fluffy kitten right now")
        elif neg_rate > 10.0 and neg_rate <= 20.0:
            return ("You are a bit frustrated and disappointed\n"
                    "But you're easily distracted and cheered with little effort 😁😁")
        elif neg_rate > 20.0 and neg_rate <= 30.0:
            return ("Things are bothering you but you're coping up\n"
                    "You might be over tired and hungry\n"
                    "The emotional equivalent of a headache 😁😁")
    elif neg_rate > 30.0 and neg_rate <= 60.0:
        if neg_rate > 30.0 and neg_rate <= 40.0:
            return ("Today is slightly a bad day for you.\n"
                    "You still have the skills to get through it, but be gentle with yourself\n"
                    "Use self-care strategies 😁😁")
        elif neg_rate > 40.0 and neg_rate <= 50.0:
            return ("Your mental health is starting to impact your everyday life.\n"
                    "Easy things are becoming difficult")
        elif neg_rate > 50.0 and neg_rate <= 60.0:
            return ("You are not able to do things the way you usually do them due to your mental health.\n"
                    "Impulsive and compulsive thoughts might be difficult to cope with")
    elif neg_rate > 60.0 and neg_rate <= 100.0:
        if neg_rate > 60.0 and neg_rate <= 70.0:
            return ("You are losing interest in the activities that used to be enjoyable.\n"
                    "You should definitely seek help\n"
                    "This is becoming serious ☹☹")
        elif neg_rate > 70.0 and neg_rate <= 80.0:
            return ("You can't ignore your struggles now; it's HIGH time!\n"
                    "You may have issues sleeping, eating, having fun, socializing, and work/study\n"
                    "Your mental health is affecting almost all parts of your life ☹☹")
        elif neg_rate > 80.0 and neg_rate <= 90.0:
            return ("You are at a critical point!!!\n"
                    "You aren't functioning anymore and need urgent help.\n"
                    "You may be a risk to yourself or others if left untreated ☹☹")
        elif neg_rate > 90.0 and neg_rate <= 100.0:
            return ("The worst mental and emotional distress possible.\n"
                    "You can't imagine things getting any better now and you might think it's all over for you (SUICIDE).\n"
                    "Contact a crisis line or get started with treatment immediately. ☹☹")

st.title("Emotion Detection and Depression Analysis")

audio_file = st.file_uploader("Upload an audio file", type=["wav", "mp3"])

if audio_file:
    st.audio(audio_file, format='audio/wav')
    emotion_response = get_emotion(audio_file)
    
    if "error" not in emotion_response:
        emotion = emotion_response.get("emotion", "neutral")
        neg_rate = get_neg_rate(emotion)
        result = evaluate_depression(neg_rate)
        st.write(result)
    else:
        st.error(f"Error fetching emotion from API: {emotion_response.get('error')}")