import streamlit as st
import cv2
import numpy as np
import mediapipe as mp
from PIL import Image

# Initialize MediaPipe
mp_pose = mp.solutions.pose
mp_drawing = mp.solutions.drawing_utils

# Configure page
st.set_page_config(layout="wide")
st.title("📷 Posture Analysis App")

# Initialize session state
if 'camera_on' not in st.session_state:
    st.session_state.camera_on = False
if 'posture_status' not in st.session_state:
    st.session_state.posture_status = "Please enable camera"

def analyze_posture(image):
    """Analyze posture using MediaPipe"""
    with mp_pose.Pose(
        min_detection_confidence=0.5,
        min_tracking_confidence=0.5,
        model_complexity=1
    ) as pose:
        
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        results = pose.process(image_rgb)
        
        annotated_image = image.copy()
        if results.pose_landmarks:
            mp_drawing.draw_landmarks(
                annotated_image, 
                results.pose_landmarks, 
                mp_pose.POSE_CONNECTIONS,
                mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=2, circle_radius=2),
                mp_drawing.DrawingSpec(color=(0, 0, 255), thickness=2)
            )
            posture_status = check_posture(results.pose_landmarks, image.shape)
        else:
            posture_status = "No pose detected"
        
        return annotated_image, posture_status

def check_posture(landmarks, image_shape):
    """Analyze body landmarks and generate posture report"""
    h, w, _ = image_shape
    
    # Get key points
    left_shoulder = landmarks.landmark[mp_pose.PoseLandmark.LEFT_SHOULDER]
    right_shoulder = landmarks.landmark[mp_pose.PoseLandmark.RIGHT_SHOULDER]
    left_hip = landmarks.landmark[mp_pose.PoseLandmark.LEFT_HIP]
    right_hip = landmarks.landmark[mp_pose.PoseLandmark.RIGHT_HIP]
    left_ear = landmarks.landmark[mp_pose.PoseLandmark.LEFT_EAR]
    right_ear = landmarks.landmark[mp_pose.PoseLandmark.RIGHT_EAR]
    nose = landmarks.landmark[mp_pose.PoseLandmark.NOSE]
    
    # Determine posture
    sitting = left_hip.y < left_shoulder.y + 0.1 or right_hip.y < right_shoulder.y + 0.1
    
    messages = []
    
    # Forward head check
    head_forward = (left_ear.y > left_shoulder.y + 0.1 or right_ear.y > right_shoulder.y + 0.1) and \
                   (nose.y > left_shoulder.y or nose.y > right_shoulder.y)
    if head_forward:
        messages.append("• Forward head tilt detected")
    
    # Shoulders check
    shoulders_rounded = left_shoulder.x > left_hip.x + 0.05 or right_shoulder.x < right_hip.x - 0.05
    if shoulders_rounded:
        messages.append("• Rounded shoulders detected")
    
    # Side tilt check
    shoulder_diff = abs(left_shoulder.y - right_shoulder.y)
    hip_diff = abs(left_hip.y - right_hip.y)
    if shoulder_diff > 0.05 or hip_diff > 0.05:
        messages.append("• Body leaning to one side")
    
    # Generate report
    if messages:
        report = [
            f"**{'Sitting' if sitting else 'Standing'} posture issues:**",
            *messages,
            "\n**Recommendations:**",
            "• Keep head straight",
            "• Pull shoulders back",
            "• Maintain straight back",
            "• Sit on sitting bones"
        ]
    else:
        report = [
            f"**Excellent {'sitting' if sitting else 'standing'} posture!**",
            "\n**Tips:**",
            "• Continue monitoring your posture"
        ]
    
    return "\n\n".join(report)

def main():
    col1, col2 = st.columns([2, 1])
    
    with col1:
        st.header("Camera View")
        
        # Camera toggle
        if st.button("Enable Camera" if not st.session_state.camera_on else "Disable Camera"):
            st.session_state.camera_on = not st.session_state.camera_on
            st.rerun()
        
        if st.session_state.camera_on:
            # Placeholder for local camera implementation
            demo_image = np.zeros((480, 640, 3), dtype=np.uint8)
            st.image(demo_image, caption="Camera feed placeholder")
            st.warning("For full camera functionality, please run locally")
            
            # Simulate analysis
            if st.button("Analyze Current Frame"):
                _, posture_status = analyze_posture(demo_image)
                st.session_state.posture_status = posture_status
    
    with col2:
        st.header("Analysis Results")
        st.markdown(st.session_state.posture_status)

if __name__ == "__main__":
    main()