"""
MediaPipe detector untuk mengambil landmark tangan dan wajah.

Mengambil landmark tangan dan wajah menggunakan MediaPipe Hands dan Face Mesh.
Feature size: 21 (left hand) + 21 (right hand) + 60 (face) = 306
"""

import cv2
import mediapipe as mp
import numpy as np
from typing import Tuple


class MediaPipeDetector:
    """
    Kelas untuk mendeteksi landmark tangan dan wajah menggunakan MediaPipe.

    Feature breakdown:
    - Left hand: 21 landmarks x 3 coords = 63
    - Right hand: 21 landmarks x 3 coords = 63
    - Face: 60 landmarks x 3 coords = 180
    - Total: 306
    """

    def __init__(
        self,
        static_image_mode: bool = False,
        max_num_hands: int = 2,
        max_num_faces: int = 1,
        min_detection_confidence: float = 0.5,
        min_tracking_confidence: float = 0.5
    ):
        """
        Inisialisasi MediaPipe detector.
        """
        self.hands = mp.solutions.hands.Hands(
            static_image_mode=static_image_mode,
            max_num_hands=max_num_hands,
            min_detection_confidence=min_detection_confidence,
            min_tracking_confidence=min_tracking_confidence
        )

        self.face_mesh = mp.solutions.face_mesh.FaceMesh(
            static_image_mode=static_image_mode,
            max_num_faces=max_num_faces,
            min_detection_confidence=min_detection_confidence,
            min_tracking_confidence=min_tracking_confidence,
            refine_landmarks=True
        )

        # 60 Face landmarks indices dari: dagu, mulut, bibir, hidung, mata, pipi
        # Ini adalah indeks dari MediaPipe Face Mesh 468 landmarks
        self.face_indices = [
            # Dagu (Chin) - 8 landmarks
            152, 234, 454, 358, 465, 452, 253, 389,
            # Alis (Eyebrows) - 8 landmarks
            336, 296, 334, 293, 300, 276, 283, 282,
            # Mata (Eyes) - 12 landmarks
            33, 160, 158, 133, 153, 144, 362, 385, 387, 263, 373, 249,
            # Hidung (Nose) - 12 landmarks
            1, 2, 98, 327, 4, 75, 221, 189, 417, 62, 305, 293,
            # Mulut luar (Mouth outer) - 8 landmarks
            13, 14, 17, 84, 185, 61, 291, 0,
            # Bibir (Mouth inner/Lips) - 8 landmarks
            78, 95, 88, 178, 87, 317, 402, 310,
            # Pipi (Cheeks) - 4 landmarks
            50, 187, 411, 93
        ]

    def detect_hands(self, frame: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        Deteksi landmark tangan kiri dan kanan.

        Args:
            frame: Frame gambar (BGR)

        Returns:
            Tuple dari (left_hand_landmarks, right_hand_landmarks)
            Setiap landmark adalah array shape (21, 3)
        """
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        results = self.hands.process(rgb_frame)

        left_hand = np.zeros((21, 3))
        right_hand = np.zeros((21, 3))

        if results.multi_hand_landmarks:
            for hand_landmarks, handedness in zip(
                results.multi_hand_landmarks,
                results.multi_handedness
            ):
                hand_type = handedness.classification[0].label
                landmarks = np.array([
                    [lm.x, lm.y, lm.z]
                    for lm in hand_landmarks.landmark
                ])

                if hand_type == "Left":
                    left_hand = landmarks
                else:
                    right_hand = landmarks

        return left_hand, right_hand

    def detect_face(self, frame: np.ndarray) -> np.ndarray:
        """
        Deteksi 60 landmark wajah dari dagu, mulut, bibir, hidung, mata, pipi.
    
        Args:
            frame: Frame gambar (BGR)

        Returns:
            Array landmark wajah shape (60, 3)
        """
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        results = self.face_mesh.process(rgb_frame)

        # Default: zeros
        face = np.zeros((60, 3))

        if results.multi_face_landmarks:
            face_landmarks = results.multi_face_landmarks[0]
            face = np.array([
                [face_landmarks.landmark[idx].x,
                 face_landmarks.landmark[idx].y,
                 face_landmarks.landmark[idx].z]
                for idx in self.face_indices
            ])

        return face

    def detect(self, frame: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Deteksi semua landmark (kedua tangan dan wajah).

        Args:
            frame: Frame gambar (BGR)

        Returns:
            Tuple (left_hand, right_hand, face)
            - left_hand: array shape (21, 3)
            - right_hand: array shape (21, 3)
            - face: array shape (60, 3)
        """
        left_hand, right_hand = self.detect_hands(frame)
        face = self.detect_face(frame)

        return left_hand, right_hand, face

    def close(self):
        """Tutup semua MediaPipe resources."""
        self.hands.close()
        self.face_mesh.close()