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

mp_pose = mp.solutions.pose
mp_drawing = mp.solutions.drawing_utils

def _to_uint8_rgb(img: np.ndarray) -> np.ndarray:
    """Ensure image is uint8 RGB."""
    if img is None:
        return None

    img = np.asarray(img)

    # Gradio usually provides RGB. Ensure dtype is uint8.
    if img.dtype != np.uint8:
        img = np.clip(img, 0, 255).astype(np.uint8)

    # If someone passes grayscale, convert to RGB.
    if img.ndim == 2:
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)

    # If RGBA, drop alpha.
    if img.ndim == 3 and img.shape[2] == 4:
        img = img[:, :, :3]

    return img

def _px_dist(a, b, w, h):
    ax, ay = a.x * w, a.y * h
    bx, by = b.x * w, b.y * h
    return float(np.hypot(ax - bx, ay - by))

def analyze_body(image):
    image = _to_uint8_rgb(image)
    if image is None:
        return None, "No image captured."

    h, w = image.shape[:2]

    # Create Pose per-call (safer with Gradio concurrency)
    with mp_pose.Pose(
        static_image_mode=True,          # single-frame analysis
        model_complexity=1,
        min_detection_confidence=0.5
    ) as pose:
        results = pose.process(image)

    if not results.pose_landmarks:
        return image, "ML Note: No body detected. Stand further back and keep your full body in frame."

    annotated = image.copy()
    mp_drawing.draw_landmarks(
        annotated,
        results.pose_landmarks,
        mp_pose.POSE_CONNECTIONS
    )

    lm = results.pose_landmarks.landmark
    ls = lm[mp_pose.PoseLandmark.LEFT_SHOULDER]
    rs = lm[mp_pose.PoseLandmark.RIGHT_SHOULDER]
    lh = lm[mp_pose.PoseLandmark.LEFT_HIP]
    rh = lm[mp_pose.PoseLandmark.RIGHT_HIP]

    shoulder_width_px = _px_dist(ls, rs, w, h)
    hip_width_px = _px_dist(lh, rh, w, h)

    ratio = (shoulder_width_px / hip_width_px) if hip_width_px > 1e-6 else 0.0

    key_vis = [ls.visibility, rs.visibility, lh.visibility, rh.visibility]
    avg_vis = float(np.mean(key_vis))

    metrics = (
        f"✨ ML Alignment Data:\n"
        f"- Shoulder Width (px): {shoulder_width_px:.1f}\n"
        f"- Hip Width (px): {hip_width_px:.1f}\n"
        f"- Shoulder-to-Hip Ratio: {ratio:.2f}\n"
        f"- Avg Keypoint Visibility (0-1): {avg_vis:.2f}\n"
        f"Tip: Keep the same camera distance + full body in frame for consistent comparisons."
    )

    return annotated, metrics

with gr.Blocks() as demo:
    gr.Markdown("# 📸 ML Body Progress Tracker")

    with gr.Row():
        input_img = gr.Image(sources=["webcam"], type="numpy", label="Webcam Feed")
        output_img = gr.Image(type="numpy", label="ML Analysis (Skeleton)")

    stats = gr.Textbox(label="Body Proportion Data", lines=7)
    btn = gr.Button("Analyze Pose")

    btn.click(fn=analyze_body, inputs=input_img, outputs=[output_img, stats])

demo.launch()