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

mp_pose = mp.solutions.pose
pose = mp_pose.Pose(
    min_detection_confidence=0.5,
    min_tracking_confidence=0.5,
    model_complexity=1
)
mp_drawing = mp.solutions.drawing_utils

def analyze_posture(image):
    """Analyze posture on the image and return the image with keypoints and analysis text"""
    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 = "キーポイントが検出されませんでした (Key points not detected)"
    
    return annotated_image, posture_status

def check_posture(landmarks, image_shape):
    """Analyze posture and return a text 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/standing)
    sitting = left_hip.y < left_shoulder.y + 0.1 or right_hip.y < right_shoulder.y + 0.1
    
    messages = []
    
    # Check for forward head posture
    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("• 頭が前に傾いています (テキストネック) (Head tilted forward - text neck)")
    
    # Check for rounded shoulders
    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)")
    
    # Check for side tilt
    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("• 体が横に傾いています (Asymmetrical posture)")
    
    # Check pelvis position
    if sitting and (left_hip.y < left_shoulder.y + 0.15 or right_hip.y < right_shoulder.y + 0.15):
        messages.append("• 骨盤が前に傾いています (Pelvis tilted forward)")
    
    # Generate final report
    if messages:
        report = [
            f"**{'座り姿勢' if sitting else '立ち姿勢'} - 問題が検出されました ({'Sitting' if sitting else 'Standing'} - problems detected):**",
            *messages,
            "\n**アドバイス (Recommendations):**",
            "• 頭をまっすぐに保ち、耳が肩の上にくるように (Keep your head straight, ears over shoulders)",
            "• 肩を後ろに引き下げて (Pull shoulders back and down)",
            "• 背中をまっすぐに保ち、横に傾かないように (Keep your back straight, avoid side tilting)",
            "• 座るときは坐骨で支えるように (When sitting, support your weight on sitting bones)"
        ]
    else:
        report = [
            f"**完璧な姿勢です ({'座り姿勢' if sitting else '立ち姿勢'})! (Perfect posture {'sitting' if sitting else 'standing'})**",
            "すべてのキーポイントが正しい位置にあります (All key points are in correct position)",
            "\n**アドバイス (Advice):**",
            "• 一日中姿勢に気を付けてください (Continue to monitor your posture throughout the day)"
        ]
    
    return "\n\n".join(report)

def main():
    st.set_page_config(layout="wide")
    st.title("📷 姿勢分析アプリ (Posture Analyzer)")
    
    col1, col2 = st.columns([2, 1])
    
    with col1:
        st.header("カメラビュー (Camera View)")
        
        if st.button("カメラアクセスを許可 (Allow camera access)"):
            st.session_state.camera_allowed = True
            
        if not st.session_state.get('camera_allowed', False):
            st.warning("⚠️ カメラを使用するには許可が必要です (Camera access requires permission)")
            st.image("demo_pose.jpg")  # Заглушка
            run = False
        else:
            run = st.checkbox("カメラを起動 (Enable camera)", value=True)
            FRAME_WINDOW = st.image([])
    
    with col2:
        st.header("姿勢分析 (Posture Analysis)")
        status_placeholder = st.empty()
        
        uploaded_file = st.file_uploader("または画像をアップロード (Or upload image)", 
                                      type=["jpg", "png", "jpeg"])
    
    if run and st.session_state.camera_allowed:
        camera = cv2.VideoCapture(0)
        if not camera.isOpened():
            st.error("カメラエラー: デバイスに接続できません (Camera error: Device not accessible)")
            run = False
            
        while run:
            ret, frame = camera.read()
            if ret:
                analyzed_frame, posture_status = analyze_posture(frame)
                FRAME_WINDOW.image(analyzed_frame)
                status_placeholder.markdown(posture_status)
            time.sleep(0.1)
        camera.release()
    
    elif uploaded_file:
        file_bytes = np.frombuffer(uploaded_file.read(), np.uint8)
        frame = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
        analyzed_frame, posture_status = analyze_posture(frame)
        col1.image(analyzed_frame)
        status_placeholder.markdown(posture_status)

if __name__ == "__main__":
    main()