#!/opt/homebrew/bin/python3.11
"""
MediaPipe Holistic Processor - Full body keypoint extraction

Purpose:
1. Extract Face Mesh (468 keypoints) → eye/mouth actions
2. Extract Pose (33 keypoints) → arm/leg/feet actions
3. Extract Hands (21 keypoints × 2) → hand gestures

Output structure:
{
  "metadata": {...},
  "frames": {
    "frame_num": {
      "persons": [
        {
          "person_id": 0,
          "bbox": {...},
          "face_mesh": {
            "landmarks": [[x,y,z], ...],  # 468 points
            "eye_features": {...},
            "mouth_features": {...},
          },
          "pose": {
            "landmarks": [[x,y,z,visibility], ...],  # 33 points
            "arm_features": {...},
            "leg_features": {...},
          },
          "hands": {
            "left": {
              "landmarks": [[x,y,z], ...],  # 21 points
              "gesture": "...",
            },
            "right": {
              "landmarks": [[x,y,z], ...],  # 21 points
              "gesture": "...",
            },
          },
        }
      ]
    }
  }
}
"""

import sys
import json
import argparse
import cv2
import numpy as np
import mediapipe as mp
from pathlib import Path
from typing import Dict, List, Optional
from collections import defaultdict


class MediaPipeHolisticProcessor:
    """
    Process video with MediaPipe Holistic (Face + Pose + Hands)
    """
    
    def __init__(
        self,
        model_complexity: int = 1,  # 0, 1, 2
        refine_face_landmarks: bool = True,
        enable_segmentation: bool = False,
        min_detection_confidence: float = 0.5,
        min_tracking_confidence: float = 0.5,
    ):
        """
        Initialize MediaPipe Holistic
        
        Args:
            model_complexity: 0 (lite), 1 (full), 2 (heavy)
            refine_face_landmarks: Enable iris detection
            enable_segmentation: Enable segmentation mask
            min_detection_confidence: Detection confidence threshold
            min_tracking_confidence: Tracking confidence threshold
        """
        self.mp_holistic = mp.solutions.holistic
        self.mp_drawing = mp.solutions.drawing_utils
        self.mp_drawing_styles = mp.solutions.drawing_styles
        
        self.holistic = self.mp_holistic.Holistic(
            static_image_mode=False,  # Video mode
            model_complexity=model_complexity,
            smooth_landmarks=True,  # Smooth landmarks across frames
            enable_segmentation=enable_segmentation,
            smooth_segmentation=True,
            refine_face_landmarks=refine_face_landmarks,
            min_detection_confidence=min_detection_confidence,
            min_tracking_confidence=min_tracking_confidence,
        )
        
        # Eye landmark indices (Face Mesh)
        self.LEFT_EYE_INDICES = [33, 133, 159, 145, 158, 144]  # 6 points
        self.RIGHT_EYE_INDICES = [362, 263, 386, 374, 385, 373]
        
        # Iris indices
        self.LEFT_IRIS_CENTER = 468
        self.RIGHT_IRIS_CENTER = 473
        
        # Mouth indices
        self.MOUTH_TOP = 13
        self.MOUTH_BOTTOM = 14
        self.MOUTH_LEFT = 61
        self.MOUTH_RIGHT = 291
        
        # Pose key indices
        self.POSE_KEYPOINTS = {
            "nose": 0,
            "left_shoulder": 11,
            "right_shoulder": 12,
            "left_elbow": 13,
            "right_elbow": 14,
            "left_wrist": 15,
            "right_wrist": 16,
            "left_hip": 23,
            "right_hip": 24,
            "left_knee": 25,
            "right_knee": 26,
            "left_ankle": 27,
            "right_ankle": 28,
        }
        
        # Hand key indices
        self.HAND_KEYPOINTS = {
            "wrist": 0,
            "thumb_cmc": 1,
            "thumb_mcp": 2,
            "thumb_ip": 3,
            "thumb_tip": 4,
            "index_mcp": 5,
            "index_pip": 6,
            "index_dip": 7,
            "index_tip": 8,
            "middle_mcp": 9,
            "middle_pip": 10,
            "middle_dip": 11,
            "middle_tip": 12,
            "ring_mcp": 13,
            "ring_pip": 14,
            "ring_dip": 15,
            "ring_tip": 16,
            "pinky_mcp": 17,
            "pinky_pip": 18,
            "pinky_dip": 19,
            "pinky_tip": 20,
        }
    
    def process_frame(self, frame: np.ndarray) -> Dict:
        """
        Process single frame
        
        Args:
            frame: BGR image
        
        Returns:
            Dict with face_mesh, pose, hands data
        """
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        
        results = self.holistic.process(frame_rgb)
        
        person_data = {
            "person_id": 0,
            "bbox": None,
            "face_mesh": None,
            "pose": None,
            "hands": {"left": None, "right": None},
        }
        
        # Extract face mesh
        if results.face_landmarks:
            person_data["face_mesh"] = self._extract_face_mesh(results.face_landmarks)
        
        # Extract pose
        if results.pose_landmarks:
            person_data["pose"] = self._extract_pose(results.pose_landmarks)
        
        # Extract hands
        if results.left_hand_landmarks:
            person_data["hands"]["left"] = self._extract_hand(results.left_hand_landmarks, "left")
        
        if results.right_hand_landmarks:
            person_data["hands"]["right"] = self._extract_hand(results.right_hand_landmarks, "right")
        
        # Calculate bbox from pose landmarks
        if results.pose_landmarks:
            landmarks = results.pose_landmarks.landmark
            x_coords = [lm.x for lm in landmarks if lm.visibility > 0.5]
            y_coords = [lm.y for lm in landmarks if lm.visibility > 0.5]
            
            if x_coords and y_coords:
                x_min, x_max = min(x_coords), max(x_coords)
                y_min, y_max = min(y_coords), max(y_coords)
                
                height, width = frame.shape[:2]
                
                person_data["bbox"] = {
                    "x": int(x_min * width),
                    "y": int(y_min * height),
                    "width": int((x_max - x_min) * width),
                    "height": int((y_max - y_min) * height),
                }
        
        return person_data
    
    def _extract_face_mesh(self, face_landmarks) -> Dict:
        """
        Extract face mesh landmarks and calculate features
        
        Args:
            face_landmarks: MediaPipe face landmarks
        
        Returns:
            Dict with landmarks, eye_features, mouth_features
        """
        landmarks = []
        for lm in face_landmarks.landmark:
            landmarks.append([lm.x, lm.y, lm.z])
        
        # Eye Aspect Ratio (EAR)
        def calculate_ear(eye_indices):
            # Get eye points
            p1 = face_landmarks.landmark[eye_indices[0]]
            p2 = face_landmarks.landmark[eye_indices[1]]
            p3 = face_landmarks.landmark[eye_indices[2]]
            p4 = face_landmarks.landmark[eye_indices[3]]
            p5 = face_landmarks.landmark[eye_indices[4]]
            p6 = face_landmarks.landmark[eye_indices[5]]
            
            # Vertical distances
            vertical_1 = np.linalg.norm([p3.x - p5.x, p3.y - p5.y])
            vertical_2 = np.linalg.norm([p4.x - p6.x, p4.y - p6.y])
            
            # Horizontal distance
            horizontal = np.linalg.norm([p1.x - p2.x, p1.y - p2.y])
            
            ear = (vertical_1 + vertical_2) / (2 * horizontal) if horizontal > 0 else 0
            return ear
        
        left_ear = calculate_ear(self.LEFT_EYE_INDICES)
        right_ear = calculate_ear(self.RIGHT_EYE_INDICES)
        avg_ear = (left_ear + right_ear) / 2
        
        # Iris position (if refined landmarks enabled)
        left_iris_x = None
        right_iris_x = None
        
        if len(face_landmarks.landmark) > 477:
            left_iris = face_landmarks.landmark[self.LEFT_IRIS_CENTER]
            right_iris = face_landmarks.landmark[self.RIGHT_IRIS_CENTER]
            
            # Normalize iris position relative to eye
            left_eye_center_x = (face_landmarks.landmark[33].x + face_landmarks.landmark[133].x) / 2
            right_eye_center_x = (face_landmarks.landmark[362].x + face_landmarks.landmark[263].x) / 2
            
            left_eye_width = abs(face_landmarks.landmark[33].x - face_landmarks.landmark[133].x)
            right_eye_width = abs(face_landmarks.landmark[362].x - face_landmarks.landmark[263].x)
            
            left_iris_x = (left_iris.x - left_eye_center_x) / left_eye_width if left_eye_width > 0 else 0
            right_iris_x = (right_iris.x - right_eye_center_x) / right_eye_width if right_eye_width > 0 else 0
        
        # Eye action detection
        eye_action = "unknown"
        if avg_ear < 0.15:
            eye_action = "closed"
        elif avg_ear > 0.4:
            eye_action = "wide_open"
        elif 0.15 <= avg_ear < 0.25:
            eye_action = "squint"
        else:
            eye_action = "normal"
        
        # Gaze direction
        gaze_direction = "center"
        if left_iris_x and right_iris_x:
            avg_iris_x = (left_iris_x + right_iris_x) / 2
            if avg_iris_x < -0.2:
                gaze_direction = "left"
            elif avg_iris_x > 0.2:
                gaze_direction = "right"
        
        # Mouth Aspect Ratio (MAR)
        mouth_top = face_landmarks.landmark[self.MOUTH_TOP]
        mouth_bottom = face_landmarks.landmark[self.MOUTH_BOTTOM]
        mouth_left = face_landmarks.landmark[self.MOUTH_LEFT]
        mouth_right = face_landmarks.landmark[self.MOUTH_RIGHT]
        
        mouth_height = np.linalg.norm([mouth_top.x - mouth_bottom.x, mouth_top.y - mouth_bottom.y])
        mouth_width = np.linalg.norm([mouth_left.x - mouth_right.x, mouth_left.y - mouth_right.y])
        
        mar = mouth_height / mouth_width if mouth_width > 0 else 0
        
        # Mouth corner distance (for smile detection)
        mouth_center_y = (mouth_top.y + mouth_bottom.y) / 2
        corner_lift = (mouth_center_y - mouth_left.y) + (mouth_center_y - mouth_right.y)
        
        # Mouth action detection
        mouth_action = "unknown"
        if mar > 0.7:
            mouth_action = "yawn"
        elif mar > 0.5:
            mouth_action = "open"
        elif mar < 0.2:
            if corner_lift > 0.02:
                mouth_action = "smile"
            else:
                mouth_action = "closed"
        else:
            mouth_action = "slightly_open"
        
        return {
            "landmarks": landmarks,
            "num_landmarks": len(landmarks),
            "eye_features": {
                "left_ear": round(left_ear, 4),
                "right_ear": round(right_ear, 4),
                "avg_ear": round(avg_ear, 4),
                "left_iris_x": round(left_iris_x, 4) if left_iris_x else None,
                "right_iris_x": round(right_iris_x, 4) if right_iris_x else None,
                "eye_action": eye_action,
                "gaze_direction": gaze_direction,
            },
            "mouth_features": {
                "mar": round(mar, 4),
                "mouth_height": round(mouth_height, 4),
                "mouth_width": round(mouth_width, 4),
                "corner_lift": round(corner_lift, 4),
                "mouth_action": mouth_action,
            },
        }
    
    def _extract_pose(self, pose_landmarks) -> Dict:
        """
        Extract pose landmarks and calculate features
        
        Args:
            pose_landmarks: MediaPipe pose landmarks
        
        Returns:
            Dict with landmarks, arm_features, leg_features
        """
        landmarks = []
        for lm in pose_landmarks.landmark:
            landmarks.append([lm.x, lm.y, lm.z, lm.visibility])
        
        # Helper function to calculate angle
        def calculate_angle(p1_idx, p2_idx, p3_idx):
            p1 = pose_landmarks.landmark[p1_idx]
            p2 = pose_landmarks.landmark[p2_idx]
            p3 = pose_landmarks.landmark[p3_idx]
            
            v1 = np.array([p1.x, p1.y]) - np.array([p2.x, p2.y])
            v2 = np.array([p3.x, p3.y]) - np.array([p2.x, p2.y])
            
            angle = np.arccos(np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2)))
            return np.degrees(angle)
        
        # Arm features
        left_elbow_angle = calculate_angle(11, 13, 15)  # shoulder-elbow-wrist
        right_elbow_angle = calculate_angle(12, 14, 16)
        
        # Check if arms raised
        left_wrist = pose_landmarks.landmark[15]
        left_elbow = pose_landmarks.landmark[13]
        left_shoulder = pose_landmarks.landmark[11]
        
        right_wrist = pose_landmarks.landmark[16]
        right_elbow = pose_landmarks.landmark[14]
        right_shoulder = pose_landmarks.landmark[12]
        
        left_arm_raised = left_wrist.y < left_elbow.y < left_shoulder.y
        right_arm_raised = right_wrist.y < right_elbow.y < right_shoulder.y
        
        # Arm action detection
        left_arm_action = "unknown"
        if left_arm_raised:
            left_arm_action = "raise_left"
        elif left_elbow_angle > 150:
            left_arm_action = "extend_left"
        elif left_elbow_angle < 90:
            left_arm_action = "fold_left"
        else:
            left_arm_action = "neutral_left"
        
        right_arm_action = "unknown"
        if right_arm_raised:
            right_arm_action = "raise_right"
        elif right_elbow_angle > 150:
            right_arm_action = "extend_right"
        elif right_elbow_angle < 90:
            right_arm_action = "fold_right"
        else:
            right_arm_action = "neutral_right"
        
        # Cross arms detection
        cross_arms = False
        if left_wrist.x > right_wrist.x and right_wrist.x < left_shoulder.x:
            cross_arms = True
        
        # Leg features
        left_knee_angle = calculate_angle(23, 25, 27)  # hip-knee-ankle
        right_knee_angle = calculate_angle(24, 26, 28)
        
        # Check standing/sitting
        left_hip = pose_landmarks.landmark[23]
        left_knee = pose_landmarks.landmark[25]
        left_ankle = pose_landmarks.landmark[27]
        
        right_hip = pose_landmarks.landmark[24]
        right_knee = pose_landmarks.landmark[26]
        right_ankle = pose_landmarks.landmark[28]
        
        hip_avg_y = (left_hip.y + right_hip.y) / 2
        knee_avg_y = (left_knee.y + right_knee.y) / 2
        
        # Standing: hip < knee < ankle (y increases downward)
        standing = left_hip.y < left_knee.y < left_ankle.y and right_hip.y < right_knee.y < right_ankle.y
        
        # Sitting: hip ≈ knee height
        sitting = abs(hip_avg_y - knee_avg_y) < 0.1
        
        # Leg action detection
        leg_action = "unknown"
        if sitting:
            leg_action = "sit"
        elif standing:
            if left_knee_angle < 120 or right_knee_angle < 120:
                leg_action = "knee_bend"
            else:
                leg_action = "stand"
        
        return {
            "landmarks": landmarks,
            "num_landmarks": len(landmarks),
            "arm_features": {
                "left_elbow_angle": round(left_elbow_angle, 2),
                "right_elbow_angle": round(right_elbow_angle, 2),
                "left_arm_raised": left_arm_raised,
                "right_arm_raised": right_arm_raised,
                "left_arm_action": left_arm_action,
                "right_arm_action": right_arm_action,
                "cross_arms": cross_arms,
            },
            "leg_features": {
                "left_knee_angle": round(left_knee_angle, 2),
                "right_knee_angle": round(right_knee_angle, 2),
                "standing": standing,
                "sitting": sitting,
                "leg_action": leg_action,
            },
        }
    
    def _extract_hand(self, hand_landmarks, hand_type: str) -> Dict:
        """
        Extract hand landmarks and detect gesture
        
        Args:
            hand_landmarks: MediaPipe hand landmarks
            hand_type: "left" or "right"
        
        Returns:
            Dict with landmarks, gesture
        """
        landmarks = []
        for lm in hand_landmarks.landmark:
            landmarks.append([lm.x, lm.y, lm.z])
        
        # Check finger extensions
        def is_finger_extended(tip_idx, pip_idx):
            tip = hand_landmarks.landmark[tip_idx]
            pip = hand_landmarks.landmark[pip_idx]
            
            # Finger is extended if tip is higher (lower y) than pip
            return tip.y < pip.y
        
        thumb_extended = is_finger_extended(4, 3)
        index_extended = is_finger_extended(8, 6)
        middle_extended = is_finger_extended(12, 10)
        ring_extended = is_finger_extended(16, 14)
        pinky_extended = is_finger_extended(20, 18)
        
        extensions = {
            "thumb": thumb_extended,
            "index": index_extended,
            "middle": middle_extended,
            "ring": ring_extended,
            "pinky": pinky_extended,
        }
        
        # Gesture detection
        gesture = "unknown"
        
        num_extended = sum(extensions.values())
        
        if num_extended == 5:
            gesture = "open_hand"
        elif num_extended == 0:
            gesture = "fist"
        elif thumb_extended and num_extended == 1:
            gesture = "thumbs_up"
        elif index_extended and middle_extended and num_extended == 2:
            gesture = "peace_sign"
        elif index_extended and num_extended == 1:
            gesture = "pointing"
        elif thumb_extended and index_extended and not any([middle_extended, ring_extended, pinky_extended]):
            # Check thumb-index distance for OK gesture
            thumb_tip = hand_landmarks.landmark[4]
            index_tip = hand_landmarks.landmark[8]
            
            distance = np.linalg.norm([thumb_tip.x - index_tip.x, thumb_tip.y - index_tip.y])
            
            if distance < 0.05:
                gesture = "ok_sign"
            else:
                gesture = "grab"
        
        return {
            "landmarks": landmarks,
            "num_landmarks": len(landmarks),
            "finger_extensions": extensions,
            "num_fingers_extended": num_extended,
            "gesture": gesture,
            "hand_type": hand_type,
        }
    
    def process_video(
        self,
        video_path: str,
        output_path: str,
        sample_interval: int = 1,
    ) -> Dict:
        """
        Process entire video
        
        Args:
            video_path: Path to video file
            output_path: Path to output JSON
            sample_interval: Process every N frames
        
        Returns:
            Dict with all processed data
        """
        cap = cv2.VideoCapture(video_path)
        
        if not cap.isOpened():
            print(f"❌ Cannot open video: {video_path}")
            return {}
        
        fps = cap.get(cv2.CAP_PROP_FPS)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        
        print(f"\nVideo: {video_path}")
        print(f"  FPS: {fps}")
        print(f"  Resolution: {width}x{height}")
        print(f"  Total frames: {total_frames}")
        print(f"  Sample interval: {sample_interval}")
        print()
        
        output_data = {
            "metadata": {
                "video_path": video_path,
                "fps": fps,
                "width": width,
                "height": height,
                "total_frames": total_frames,
                "sample_interval": sample_interval,
                "processor": "mediapipe_holistic",
                "model_complexity": 1,
                "refine_face_landmarks": True,
            },
            "frames": {},
        }
        
        frame_count = 0
        processed_count = 0
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            frame_count += 1
            
            if frame_count % sample_interval != 0:
                continue
            
            # Process frame
            person_data = self.process_frame(frame)
            
            # Only save if landmarks detected
            if person_data["face_mesh"] or person_data["pose"] or person_data["hands"]["left"] or person_data["hands"]["right"]:
                timestamp = frame_count / fps if fps > 0 else 0
                
                output_data["frames"][str(frame_count)] = {
                    "frame_number": frame_count,
                    "timestamp": round(timestamp, 3),
                    "persons": [person_data],
                }
                
                processed_count += 1
                
                if processed_count % 100 == 0:
                    print(f"  Processed {processed_count} frames (frame {frame_count})")
        
        cap.release()
        
        # Update metadata
        output_data["metadata"]["processed_frames"] = processed_count
        
        # Save output
        with open(output_path, "w") as f:
            json.dump(output_data, f, indent=2)
        
        print(f"\n✅ Processed {processed_count} frames")
        print(f"✅ Output saved to: {output_path}")
        
        return output_data
    
    def close(self):
        """Close MediaPipe model"""
        self.holistic.close()


def main():
    parser = argparse.ArgumentParser(description="MediaPipe Holistic Processor")
    parser.add_argument("--video", required=True, help="Path to video file")
    parser.add_argument("--output", required=True, help="Path to output JSON")
    parser.add_argument("--sample-interval", type=int, default=1, help="Process every N frames")
    parser.add_argument("--model-complexity", type=int, default=1, choices=[0, 1, 2], help="Model complexity")
    parser.add_argument("--test-frame", type=int, help="Test single frame only")
    args = parser.parse_args()
    
    print("=" * 70)
    print("MediaPipe Holistic Processor")
    print("=" * 70)
    
    processor = MediaPipeHolisticProcessor(
        model_complexity=args.model_complexity,
        refine_face_landmarks=True,
    )
    
    if args.test_frame:
        # Test single frame
        print(f"\nTesting frame {args.test_frame}...")
        
        cap = cv2.VideoCapture(args.video)
        cap.set(cv2.CAP_PROP_POS_FRAMES, args.test_frame - 1)
        
        ret, frame = cap.read()
        cap.release()
        
        if ret:
            person_data = processor.process_frame(frame)
            
            print("\n=== Results ===")
            
            if person_data["face_mesh"]:
                face = person_data["face_mesh"]
                print(f"\nFace Mesh: {face['num_landmarks']} landmarks")
                print(f"  Eye: {face['eye_features']['eye_action']} (EAR: {face['eye_features']['avg_ear']})")
                print(f"  Gaze: {face['eye_features']['gaze_direction']}")
                print(f"  Mouth: {face['mouth_features']['mouth_action']} (MAR: {face['mouth_features']['mar']})")
            
            if person_data["pose"]:
                pose = person_data["pose"]
                print(f"\nPose: {pose['num_landmarks']} keypoints")
                print(f"  Left arm: {pose['arm_features']['left_arm_action']} (angle: {pose['arm_features']['left_elbow_angle']}°)")
                print(f"  Right arm: {pose['arm_features']['right_arm_action']} (angle: {pose['arm_features']['right_elbow_angle']}°)")
                print(f"  Cross arms: {pose['arm_features']['cross_arms']}")
                print(f"  Leg: {pose['leg_features']['leg_action']}")
            
            if person_data["hands"]["left"]:
                hand = person_data["hands"]["left"]
                print(f"\nLeft hand: {hand['num_landmarks']} keypoints")
                print(f"  Gesture: {hand['gesture']}")
                print(f"  Fingers extended: {hand['num_fingers_extended']}")
            
            if person_data["hands"]["right"]:
                hand = person_data["hands"]["right"]
                print(f"\nRight hand: {hand['num_landmarks']} keypoints")
                print(f"  Gesture: {hand['gesture']}")
                print(f"  Fingers extended: {hand['num_fingers_extended']}")
        else:
            print("❌ Cannot read frame")
    else:
        # Process entire video
        processor.process_video(
            args.video,
            args.output,
            args.sample_interval,
        )
    
    processor.close()


if __name__ == "__main__":
    main()