import streamlit as st
import cv2
import numpy as np
from transformers import pipeline

@st.cache_resource
def load_emotion_model():
    # Load emotion recognition pipeline
    return pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base")

# Load the model
emotion_model = load_emotion_model()

# Streamlit UI
st.title("Real-Time Stress Detection")
st.write("This app uses a Hugging Face emotion recognition model to estimate stress levels based on webcam input.")

# Activate webcam
run = st.checkbox("Run Webcam")
frame_placeholder = st.empty()

if run:
    cap = cv2.VideoCapture(0)  # Start webcam
    while True:
        ret, frame = cap.read()
        if not ret:
            st.error("Failed to access webcam.")
            break

        # Process frame for display
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frame_placeholder.image(frame_rgb, channels="RGB")

        # Mock text-based emotion input for testing
        # Replace this with input from a more suitable sensor or feature extractor
        emotion_input = "I'm feeling anxious and stressed."

        # Predict emotion
        predictions = emotion_model(emotion_input)
        predicted_emotion = predictions[0]["label"]
        st.write(f"Predicted Emotion: {predicted_emotion}")

        # Map emotion to stress level (example logic)
        stress_mapping = {
            "joy": "Low Stress",
            "sadness": "High Stress",
            "anger": "High Stress",
            "fear": "High Stress",
            "neutral": "Moderate Stress"
        }
        stress_level = stress_mapping.get(predicted_emotion, "Unknown")
        st.write(f"Estimated Stress Level: {stress_level}")

        # Break on user input
        if st.button("Stop"):
            break

    cap.release()
    st.write("Webcam stopped.")