# MediaPipe Pose for AR overlay tab

from PIL import Image
import mediapipe as mp
import numpy as np

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

LANDMARK = mp_pose.PoseLandmark
TORSO_INDICES = [
  LANDMARK.LEFT_SHOULDER.value,
  LANDMARK.RIGHT_SHOULDER.value,
  LANDMARK.LEFT_HIP.value,
  LANDMARK.RIGHT_HIP.value,
]


def extract_pose_mediapipe(image_pil: Image.Image):
  img_rgb = np.array(image_pil.convert("RGB"))
  height, width = img_rgb.shape[:2]

  with mp_pose.Pose(
    static_image_mode=True,
    model_complexity=1,
    enable_segmentation=False,
    min_detection_confidence=0.5,
    min_tracking_confidence=0.5,
  ) as pose:
    results = pose.process(img_rgb)

  if not results.pose_landmarks:
    return None

  landmarks_raw = results.pose_landmarks.landmark

  landmarks = {}
  for idx, name in enumerate(LANDMARK):
    px = int(landmarks_raw[idx].x * width)
    py = int(landmarks_raw[idx].y * height)
    landmarks[name.name] = (px, py)

  skeleton_canvas = np.zeros_like(img_rgb)
  mp_drawing.draw_landmarks(
    skeleton_canvas,
    results.pose_landmarks,
    mp_pose.POSE_CONNECTIONS,
    landmark_drawing_spec=mp_drawing.DrawingSpec(
      color=(0, 255, 200), thickness=2, circle_radius=4
    ),
    connection_drawing_spec=mp_drawing.DrawingSpec(color=(100, 180, 255), thickness=2),
  )

  torso_x = [landmarks_raw[idx].x * width for idx in TORSO_INDICES]
  torso_y = [landmarks_raw[idx].y * height for idx in TORSO_INDICES]

  pad_x = int((max(torso_x) - min(torso_x)) * 0.12)
  pad_y_top = int((max(torso_y) - min(torso_y)) * 0.08)
  pad_y_bottom = int((max(torso_y) - min(torso_y)) * 0.05)

  x1 = max(0, int(min(torso_x)) - pad_x)
  y1 = max(0, int(min(torso_y)) - pad_y_top)
  x2 = min(width, int(max(torso_x)) + pad_x)
  y2 = min(height, int(max(torso_y)) + pad_y_bottom)

  return {
    "landmarks": landmarks,
    "skeleton_pil": Image.fromarray(skeleton_canvas),
    "torso_box": (x1, y1, x2, y2),
    "image_size": (width, height),
  }


def draw_skeleton_overlay(image_pil: Image.Image, pose_data: dict, alpha: float = 0.6):
  base = np.array(image_pil.convert("RGB")).astype(float)
  skeleton = np.array(pose_data["skeleton_pil"].convert("RGB")).astype(float)

  skeleton_mask = (skeleton.sum(axis=2) > 0).astype(float)[:, :, np.newaxis]
  blended = base * (1 - skeleton_mask * alpha) + skeleton * skeleton_mask * alpha

  return Image.fromarray(blended.astype(np.uint8))