import cv2
import numpy as np
import mediapipe as mp
from typing import List, Tuple, Optional

class FaceDetector:
    def __init__(self, min_detection_confidence=0.5):
        self.mp_face_detection = mp.solutions.face_detection
        self.mp_face_mesh = mp.solutions.face_mesh
        self.face_detection = self.mp_face_detection.FaceDetection(
            model_selection=1,
            min_detection_confidence=min_detection_confidence
        )
        self.face_mesh = self.mp_face_mesh.FaceMesh(
            max_num_faces=1,
            min_detection_confidence=min_detection_confidence,
            min_tracking_confidence=min_detection_confidence
        )
    
    def detect_faces(self, image: np.ndarray) -> List[dict]:
        """
        Detect faces in image
        """
        rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        results = self.face_detection.process(rgb_image)
        
        faces = []
        if results.detections:
            for detection in results.detections:
                bbox = detection.location_data.relative_bounding_box
                faces.append({
                    'bbox': {
                        'x': bbox.xmin,
                        'y': bbox.ymin,
                        'width': bbox.width,
                        'height': bbox.height
                    },
                    'confidence': detection.score[0]
                })
        
        return faces
    
    def get_face_landmarks(self, image: np.ndarray) -> Optional[List[Tuple[float, float]]]:
        """
        Get face landmarks for detailed analysis
        """
        rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        results = self.face_mesh.process(rgb_image)
        
        if results.multi_face_landmarks:
            landmarks = []
            h, w, _ = image.shape
            for landmark in results.multi_face_landmarks[0].landmark:
                landmarks.append((landmark.x * w, landmark.y * h))
            return landmarks
        
        return None
    
    def get_face_regions(self, image: np.ndarray) -> dict:
        """
        Extract different face regions for analysis
        """
        landmarks = self.get_face_landmarks(image)
        if not landmarks:
            return {}
        
        h, w, _ = image.shape
        
        # Define regions based on landmarks (simplified)
        regions = {
            'forehead': (int(w * 0.3), int(h * 0.1), int(w * 0.4), int(h * 0.2)),
            'left_cheek': (int(w * 0.1), int(h * 0.3), int(w * 0.2), int(h * 0.3)),
            'right_cheek': (int(w * 0.7), int(h * 0.3), int(w * 0.2), int(h * 0.3)),
            'nose': (int(w * 0.4), int(h * 0.3), int(w * 0.2), int(h * 0.2)),
            'chin': (int(w * 0.4), int(h * 0.6), int(w * 0.2), int(h * 0.2)),
        }
        
        return regions
    
    def crop_face(self, image: np.ndarray, padding: float = 0.2) -> Optional[np.ndarray]:
        """
        Crop face from image with padding
        """
        faces = self.detect_faces(image)
        if not faces:
            return None
        
        face = faces[0]
        h, w, _ = image.shape
        
        x = int(face['bbox']['x'] * w)
        y = int(face['bbox']['y'] * h)
        width = int(face['bbox']['width'] * w)
        height = int(face['bbox']['height'] * h)
        
        # Add padding
        pad_x = int(width * padding)
        pad_y = int(height * padding)
        
        x1 = max(0, x - pad_x)
        y1 = max(0, y - pad_y)
        x2 = min(w, x + width + pad_x)
        y2 = min(h, y + height + pad_y)
        
        return image[y1:y2, x1:x2]
    
    def get_face_angle(self, image: np.ndarray) -> float:
        """
        Estimate face angle (for pose detection)
        """
        landmarks = self.get_face_landmarks(image)
        if not landmarks:
            return 0.0
        
        # Use eye positions to estimate angle
        left_eye = landmarks[33]  # Left eye index
        right_eye = landmarks[263]  # Right eye index
        
        dx = right_eye[0] - left_eye[0]
        dy = right_eye[1] - left_eye[1]
        
        angle = np.arctan2(dy, dx) * 180 / np.pi
        return angle