import streamlit as st
from transformers import Wav2Vec2ForSequenceClassification, Wav2Vec2Tokenizer
import torch
import numpy as np
import io

# Load the model and tokenizer
model_name = "harshit345/xlsr-wav2vec-speech-emotion-recognition"
model = Wav2Vec2ForSequenceClassification.from_pretrained(model_name)
tokenizer = Wav2Vec2Tokenizer.from_pretrained(model_name)

def analyze_emotion(audio_bytes):
    audio_input = tokenizer(audio_bytes, return_tensors="pt", padding="longest")
    with torch.no_grad():
        logits = model(**audio_input).logits
    prob = torch.sigmoid(logits).numpy().flatten()
    neg_rate = prob[0] * 100  # Assuming the first element is the negative emotion probability
    
    if neg_rate <= 30.0:
        if neg_rate <= 5.0:
            return "just chill and relax"
        elif neg_rate <= 10.0:
            return (f"***   MILD DEPRESSION  ***\n"
                    f"\nDepression Scale Detected : {round(neg_rate) / 10}\n"
                    f"My Response :\tEverything is a-okay! There's absolutely nothing wrong 😁😁\n"
                    f"\t\tYou're probably cuddling a fluffy kitten right now")
        elif neg_rate <= 20.0:
            return (f"\nDepression Scale Detected : {round(neg_rate) / 10}\n"
                    f"My Response :\tYou are a bit frustrated and disappointed\n"
                    f"\t\tBut you're easily distracted and cheered with little effort 😁😁")
        elif neg_rate <= 30.0:
            return (f"\nDepression Scale Detected : {round(neg_rate) / 10}\n"
                    f"My Response :\tThings are bothering you but you're coping up\n"
                    f"\t\tYou might be over tired and hungry\n"
                    f"\t\tThe emotional equivalent of a headache 😁😁")
    elif neg_rate <= 60.0:
        if neg_rate <= 40.0:
            return (f"***   MODERATE DEPRESSION  ***\n"
                    f"\nDepression Scale Detected : {round(neg_rate) / 10}\n"
                    f"My Response :\tToday is slightly bad for you.\n"
                    f"\t\tYou still have the skills to get through it, but be gentle with yourself\n"
                    f"\t\tUse self-care strategies 😁😁")
        elif neg_rate <= 50.0:
            return (f"\nDepression Scale Detected : {round(neg_rate) / 10}\n"
                    f"My Response :\tYour mental health is starting to impact your everyday life.\n"
                    f"\t\tEasy things are becoming difficult")
        elif neg_rate <= 60.0:
            return (f"\nDepression Scale Detected : {round(neg_rate) / 10}\n"
                    f"My Response :\tYou are not able to do things the way usually you do them due to your mental health.\n"
                    f"\t\tImpulsive and compulsive thoughts might be difficult to cope with")
    elif neg_rate <= 100.0:
        if neg_rate <= 70.0:
            return (f"***   SEVERE DEPRESSION  ***\n"
                    f"\nDepression Scale Detected : {round(neg_rate) / 10}\n"
                    f"My Response :\tYou are losing interest in the activities that used to be enjoyable.\n"
                    f"\t\tYou should definitely seek help\n"
                    f"\t\tThis is becoming serious ☹☹")
        elif neg_rate <= 80.0:
            return (f"\nDepression Scale Detected : {round(neg_rate) / 10}\n"
                    f"My Response :\tYou can't ignore your struggles now, It's HIGH time!.\n"
                    f"\t\tYou may have issues sleeping, eating, having fun, socializing, and work/study\n"
                    f"\t\tYour mental health is affecting almost all parts of your life ☹☹")
        elif neg_rate <= 90.0:
            return (f"\nDepression Scale Detected : {round(neg_rate) / 10}\n"
                    f"My Response :\tYou are at a critical point !!!\n"
                    f"\t\tYou aren't functioning anymore and need urgent help.\n"
                    f"\t\tYou may be a risk to yourself or others if left untreated ☹☹")
        elif neg_rate <= 100.0:
            return (f"\nDepression Scale Detected : {round(neg_rate) / 10}\n"
                    f"My Response :\tThe worst mental and emotional distress possible.\n"
                    f"\t\tYou can't imagine things getting any better now and you might think it's all over for you (SUICIDE).\n"
                    f"\t\tContact crisis line or get started for the treatment immediately. ☹☹")

st.title("Emotion Detection App")

uploaded_file = st.file_uploader("Choose an audio file", type=["wav"])

if uploaded_file is not None:
    audio_bytes = uploaded_file.read()
    result = analyze_emotion(audio_bytes)
    st.write(result)