import streamlit as st
import cv2
import numpy as np
from deepface import DeepFace
from PIL import Image
import io

st.set_page_config(
    page_title="✨ Age & Gender Predictor",
    page_icon=":sparkles:",
    layout="centered",
)

st.title("✨ Age & Gender Predictor")
st.write(
    """
    Welcome to the future of facial analysis!  
    **Take a snapshot with your webcam** and let our cutting-edge AI reveal your age and gender with impressive precision.
    
    **No data is stored**.
    """
)

# Initialize session state for analysis results
if 'age' not in st.session_state:
    st.session_state.age = None
if 'gender' not in st.session_state:
    st.session_state.gender = None
if 'gender_confidence' not in st.session_state:
    st.session_state.gender_confidence = None

# Streamlit's built-in webcam capture
img_file_buffer = st.camera_input("Take a picture with your webcam")

# If an image was captured
if img_file_buffer is not None:
    # Convert the image buffer to a CV2 image
    bytes_data = img_file_buffer.getvalue()
    img_array = np.frombuffer(bytes_data, np.uint8)
    img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
    
    # Display a spinner while analyzing
    with st.spinner("Analyzing your image with advanced AI models..."):
        try:
            # Analyze the image using DeepFace
            results = DeepFace.analyze(
                img,
                actions=['age', 'gender'],
                detector_backend='retinaface',
                enforce_detection=True,
                align=True
            )
            
            # Process results
            if isinstance(results, list) and len(results) > 0:
                results = sorted(results, key=lambda x: x.get('face_confidence', 0), reverse=True)
                main_result = results[0]
            else:
                main_result = results
            
            # Store results in session state
            st.session_state.age = main_result['age']
            st.session_state.gender = main_result['gender']
            
            # Handle gender confidence if available
            if isinstance(main_result['gender'], dict):
                dominant_gender = max(main_result['gender'], key=main_result['gender'].get)
                st.session_state.gender = dominant_gender
                st.session_state.gender_confidence = main_result['gender'][dominant_gender]
            
            # Display success message
            st.success("Analysis complete! Here's what we found:")
            
            # Display detailed results
            st.write("## Detailed Results")
            st.write(f"**Predicted Age:** {st.session_state.age} years")
            
            if st.session_state.gender_confidence:
                st.write(f"**Predicted Gender:** {st.session_state.gender} ({st.session_state.gender_confidence:.2f}% confidence)")
            else:
                st.write(f"**Predicted Gender:** {st.session_state.gender}")
            
        except Exception as e:
            st.error(f"Analysis failed: {str(e)}")
            st.info(
                "For best results, please try the following tips:\n"
                "- Ensure good lighting conditions\n"
                "- Position your face clearly in the frame\n"
                "- Move closer to the camera if needed"
            )

st.markdown("---")
st.markdown(
    """
    **Powered by DeepFace & RetinaFace**    
    """
)