import cv2
import numpy as np
import mediapipe as mp
import gradio as gr

# Initialize Mediapipe Pose
mp_pose = mp.solutions.pose
mp_drawing = mp.solutions.drawing_utils
pose = mp_pose.Pose(static_image_mode=True, min_detection_confidence=0.5)

def calculate_angle(a, b, c):
    """Calculate angle between three points (for body extension)."""
    a = np.array(a)  # Point A
    b = np.array(b)  # Joint (Point B)
    c = np.array(c)  # Point C

    ba = a - b
    bc = c - b

    cosine_angle = np.dot(ba, bc) / (np.linalg.norm(ba) * np.linalg.norm(bc))
    angle = np.arccos(cosine_angle)
    return np.degrees(angle)

def detect_pose(image):
    """Detect pose in an image and return annotated image + example angle."""
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    results = pose.process(image_rgb)

    annotated = image.copy()
    angle_text = "No person detected"

    if results.pose_landmarks:
        # Draw pose landmarks
        mp_drawing.draw_landmarks(
            annotated,
            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)
        )

        # Example: Calculate elbow angle (shoulder-elbow-wrist)
        h, w, _ = image.shape
        landmarks = results.pose_landmarks.landmark
        shoulder = [landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].x * w,
                    landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].y * h]
        elbow = [landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value].x * w,
                 landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value].y * h]
        wrist = [landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].x * w,
                 landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].y * h]

        angle = calculate_angle(shoulder, elbow, wrist)
        angle_text = f"Left Elbow Angle: {int(angle)}°"

        cv2.putText(annotated, angle_text, (20, 40),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 0), 2)

    return annotated[:, :, ::-1], angle_text  # Convert BGR → RGB for Gradio

# Gradio Interface
demo = gr.Interface(
    fn=detect_pose,
    inputs=gr.Image(type="numpy", label="Upload Image"),
    outputs=[gr.Image(type="numpy", label="Pose Estimation"), gr.Textbox(label="Body Extension Angle")],
    title="Human Pose Estimation with MediaPipe",
    description="Upload an image with a person doing exercise. The app will detect body pose and calculate joint angles."
)

if __name__ == "__main__":
    demo.launch()