"""Face pose detection (yaw, pitch, roll) using RetinaFace landmarks"""

import os
import numpy as np
from math import atan2, degrees
from typing import Dict, Tuple, Optional, List

# Skip RetinaFace import during tests to avoid illegal instruction errors
if os.getenv("SKIP_DEEPFACE_IN_TESTS") == "1":
    RETINAFACE_AVAILABLE = False
    RetinaFace = None
else:
    try:
        from retinaface import RetinaFace
        RETINAFACE_AVAILABLE = True
    except ImportError:
        RETINAFACE_AVAILABLE = False
        RetinaFace = None


class PoseDetector:
    """Detect face pose (yaw, pitch, roll) using RetinaFace landmarks"""
    
    # Thresholds for pose detection (in degrees)
    PROFILE_YAW_THRESHOLD = 30.0  # Faces with |yaw| >= 30° are considered profile
    EXTREME_YAW_THRESHOLD = 60.0  # Faces with |yaw| >= 60° are extreme profile
    
    PITCH_THRESHOLD = 20.0  # Faces with |pitch| >= 20° are looking up/down
    EXTREME_PITCH_THRESHOLD = 45.0  # Faces with |pitch| >= 45° are extreme
    
    ROLL_THRESHOLD = 15.0  # Faces with |roll| >= 15° are tilted
    EXTREME_ROLL_THRESHOLD = 45.0  # Faces with |roll| >= 45° are extreme
    
    def __init__(self, 
                 yaw_threshold: float = None,
                 pitch_threshold: float = None,
                 roll_threshold: float = None):
        """Initialize pose detector
        
        Args:
            yaw_threshold: Yaw angle threshold for profile detection (degrees)
                          Default: 30.0
            pitch_threshold: Pitch angle threshold for up/down detection (degrees)
                           Default: 20.0
            roll_threshold: Roll angle threshold for tilt detection (degrees)
                          Default: 15.0
        """
        if not RETINAFACE_AVAILABLE:
            raise RuntimeError("RetinaFace not available")
        
        self.yaw_threshold = yaw_threshold or self.PROFILE_YAW_THRESHOLD
        self.pitch_threshold = pitch_threshold or self.PITCH_THRESHOLD
        self.roll_threshold = roll_threshold or self.ROLL_THRESHOLD
    
    @staticmethod
    def detect_faces_with_landmarks(img_path: str) -> Dict:
        """Detect faces using RetinaFace directly
        
        Returns:
            Dictionary with face keys and landmark data:
            {
                'face_1': {
                    'facial_area': {'x': x, 'y': y, 'w': w, 'h': h},
                    'landmarks': {
                        'left_eye': (x, y),
                        'right_eye': (x, y),
                        'nose': (x, y),
                        'left_mouth': (x, y),
                        'right_mouth': (x, y)
                    },
                    'confidence': 0.95
                }
            }
        """
        if not RETINAFACE_AVAILABLE:
            return {}
        
        faces = RetinaFace.detect_faces(img_path)
        return faces
    
    @staticmethod
    def calculate_face_width_from_landmarks(landmarks: Dict) -> Optional[float]:
        """Calculate face width (eye distance) from facial landmarks.
        
        Face width is the horizontal distance between the two eyes.
        For profile faces, this distance is very small (< 20-30 pixels).
        
        Args:
            landmarks: Dictionary with landmark positions:
                      {
                          'left_eye': (x, y),
                          'right_eye': (x, y),
                          ...
                      }
        
        Returns:
            Face width in pixels, or None if landmarks invalid
        """
        if not landmarks:
            return None
        
        left_eye = landmarks.get('left_eye')
        right_eye = landmarks.get('right_eye')
        
        if not all([left_eye, right_eye]):
            return None
        
        # Calculate face width (eye distance)
        face_width = abs(right_eye[0] - left_eye[0])
        
        return face_width if face_width > 0 else None
    
    @staticmethod
    def calculate_yaw_from_landmarks(landmarks: Dict) -> Optional[float]:
        """Calculate yaw angle from facial landmarks
        
        Args:
            landmarks: Dictionary with landmark positions:
                      {
                          'left_eye': (x, y),
                          'right_eye': (x, y),
                          'nose': (x, y),
                          'left_mouth': (x, y),
                          'right_mouth': (x, y)
                      }
        
        Returns:
            Yaw angle in degrees (-90 to +90):
            - Negative: face turned left (left profile visible)
            - Positive: face turned right (right profile visible)
            - Zero: frontal face
            - None: if landmarks invalid
        """
        if not landmarks:
            return None
        
        left_eye = landmarks.get('left_eye')
        right_eye = landmarks.get('right_eye')
        nose = landmarks.get('nose')
        
        if not all([left_eye, right_eye, nose]):
            return None
        
        # Calculate eye midpoint
        eye_mid_x = (left_eye[0] + right_eye[0]) / 2
        eye_mid_y = (left_eye[1] + right_eye[1]) / 2
        
        # Calculate horizontal distance from nose to eye midpoint
        nose_x = nose[0]
        eye_midpoint_x = eye_mid_x
        
        # Calculate face width (eye distance)
        face_width = abs(right_eye[0] - left_eye[0])
        
        if face_width == 0:
            return None
        
        # Calculate horizontal offset
        horizontal_offset = nose_x - eye_midpoint_x
        
        # Calculate yaw angle using atan2
        # Normalize by face width to get angle
        yaw_radians = atan2(horizontal_offset, face_width)
        yaw_degrees = degrees(yaw_radians)
        
        return yaw_degrees
    
    @staticmethod
    def calculate_pitch_from_landmarks(landmarks: Dict) -> Optional[float]:
        """Calculate pitch angle from facial landmarks (up/down tilt)
        
        Args:
            landmarks: Dictionary with landmark positions
        
        Returns:
            Pitch angle in degrees (-90 to +90):
            - Positive: looking up
            - Negative: looking down
            - None: if landmarks invalid
        """
        if not landmarks:
            return None
        
        left_eye = landmarks.get('left_eye')
        right_eye = landmarks.get('right_eye')
        # RetinaFace uses 'mouth_left' and 'mouth_right', not 'left_mouth' and 'right_mouth'
        left_mouth = landmarks.get('mouth_left') or landmarks.get('left_mouth')
        right_mouth = landmarks.get('mouth_right') or landmarks.get('right_mouth')
        nose = landmarks.get('nose')
        
        if not all([left_eye, right_eye, left_mouth, right_mouth, nose]):
            return None
        
        # Eye midpoint
        eye_mid_y = (left_eye[1] + right_eye[1]) / 2
        # Mouth midpoint
        mouth_mid_y = (left_mouth[1] + right_mouth[1]) / 2
        # Nose vertical position
        nose_y = nose[1]
        
        # Expected nose position (typically 60% down from eyes to mouth)
        expected_nose_y = eye_mid_y + (mouth_mid_y - eye_mid_y) * 0.6
        face_height = abs(mouth_mid_y - eye_mid_y)
        
        if face_height == 0:
            return None
        
        # Vertical offset from expected position
        vertical_offset = nose_y - expected_nose_y
        
        # Calculate pitch angle
        pitch_radians = atan2(vertical_offset, face_height)
        pitch_degrees = degrees(pitch_radians)
        
        return pitch_degrees
    
    @staticmethod
    def calculate_roll_from_landmarks(landmarks: Dict) -> Optional[float]:
        """Calculate roll angle from facial landmarks (rotation around face axis)
        
        Args:
            landmarks: Dictionary with landmark positions
        
        Returns:
            Roll angle in degrees (-90 to +90):
            - Positive: tilted right (clockwise)
            - Negative: tilted left (counterclockwise)
            - None: if landmarks invalid
        """
        if not landmarks:
            return None
        
        left_eye = landmarks.get('left_eye')
        right_eye = landmarks.get('right_eye')
        
        if not all([left_eye, right_eye]):
            return None
        
        # Calculate angle of eye line
        dx = right_eye[0] - left_eye[0]
        dy = right_eye[1] - left_eye[1]
        
        if dx == 0:
            return 90.0 if dy > 0 else -90.0  # Vertical line
        
        # Roll angle - atan2 returns [-180, 180], normalize to [-90, 90]
        roll_radians = atan2(dy, dx)
        roll_degrees = degrees(roll_radians)
        
        # Normalize to [-90, 90] range for head tilt
        # If angle is > 90°, subtract 180°; if < -90°, add 180°
        if roll_degrees > 90.0:
            roll_degrees = roll_degrees - 180.0
        elif roll_degrees < -90.0:
            roll_degrees = roll_degrees + 180.0
        
        return roll_degrees
    
    @staticmethod
    def classify_pose_mode(yaw: Optional[float], 
                          pitch: Optional[float], 
                          roll: Optional[float],
                          face_width: Optional[float] = None) -> str:
        """Classify face pose mode from all three angles and optionally face width
        
        Args:
            yaw: Yaw angle in degrees
            pitch: Pitch angle in degrees
            roll: Roll angle in degrees
            face_width: Face width in pixels (eye distance). Used as indicator for profile detection.
                       If face_width < 25px, indicates profile view. When yaw is available but < 30°,
                       face_width can override yaw if it suggests profile (face_width < 25px).
        
        Returns:
            Pose mode classification string:
            - 'frontal': frontal, level, upright
            - 'profile_left', 'profile_right': profile views
            - 'looking_up', 'looking_down': pitch variations
            - 'tilted_left', 'tilted_right': roll variations
            - Combined modes: e.g., 'profile_left_looking_up'
        """
        # Default to frontal if angles unknown
        yaw_original = yaw
        if yaw is None:
            yaw = 0.0
        if pitch is None:
            pitch = 0.0
        if roll is None:
            roll = 0.0
        
        # Face width threshold for profile detection (in pixels)
        # Profile faces have very small eye distance (< 25 pixels typically)
        PROFILE_FACE_WIDTH_THRESHOLD = 10.0 #25.0
        
        # Yaw classification - PRIMARY INDICATOR
        # Use yaw angle as the primary indicator (30° threshold)
        abs_yaw = abs(yaw)
        
        # Primary classification based on yaw angle
        if abs_yaw < 30.0:
            # Yaw indicates frontal view
            # Trust yaw when it's available and reasonable (< 30°)
            # Only use face_width as fallback when yaw is unavailable (None)
            if yaw_original is None:
                # Yaw unavailable - use face_width as fallback
                if face_width is not None:
                    if face_width < PROFILE_FACE_WIDTH_THRESHOLD:
                        # Face width suggests profile view - use it when yaw is unavailable
                        yaw_mode = "profile_left"  # Default direction when yaw unavailable
                    else:
                        # Face width is normal (>= 25px) - likely frontal
                        yaw_mode = "frontal"
                else:
                    # Both yaw and face_width unavailable - cannot determine reliably
                    # This usually means landmarks are incomplete (missing nose and/or eyes)
                    # For extreme profile views, both eyes might not be visible, which would
                    # cause face_width to be None. In this case, we cannot reliably determine
                    # pose without additional indicators (like face bounding box aspect ratio).
                    # Default to frontal (conservative approach), but this might misclassify
                    # some extreme profile faces.
                    yaw_mode = "frontal"
            else:
                # Yaw is available and < 30° - but still check face_width
                # If face_width is very small (< 25px), it suggests profile even with small yaw
                if face_width is not None:
                    if face_width < PROFILE_FACE_WIDTH_THRESHOLD:
                        # Face width suggests profile view - override yaw
                        yaw_mode = "profile_left"  # Default direction when yaw is small
                    else:
                        # Face width is normal (>= 25px) - trust yaw, classify as frontal
                        yaw_mode = "frontal"
                else:
                    # No face_width provided - trust yaw, classify as frontal
                    yaw_mode = "frontal"
        elif yaw <= -30.0:
            # abs_yaw >= 30.0 and yaw is negative - profile left
            yaw_mode = "profile_left"  # Negative yaw = face turned left = left profile visible
        elif yaw >= 30.0:
            # abs_yaw >= 30.0 and yaw is positive - profile right
            yaw_mode = "profile_right"  # Positive yaw = face turned right = right profile visible
        else:
            # This should never be reached, but handle edge case
            yaw_mode = "slight_yaw"
        
        # Pitch classification
        abs_pitch = abs(pitch)
        if abs_pitch < 20.0:
            pitch_mode = "level"
        elif pitch > 20.0:
            pitch_mode = "looking_up"
        elif pitch < -20.0:
            pitch_mode = "looking_down"
        else:
            pitch_mode = "slight_pitch"
        
        # Roll classification
        abs_roll = abs(roll)
        if abs_roll < 15.0:
            roll_mode = "upright"
        elif roll > 15.0:
            roll_mode = "tilted_right"
        elif roll < -15.0:
            roll_mode = "tilted_left"
        else:
            roll_mode = "slight_roll"
        
        # Combine modes - simple case first
        if yaw_mode == "frontal" and pitch_mode == "level" and roll_mode == "upright":
            return "frontal"
        
        # Build combined mode string
        modes = []
        if yaw_mode != "frontal":
            modes.append(yaw_mode)
        if pitch_mode != "level":
            modes.append(pitch_mode)
        if roll_mode != "upright":
            modes.append(roll_mode)
        
        return "_".join(modes) if modes else "frontal"
    
    def detect_pose_faces(self, img_path: str) -> List[Dict]:
        """Detect all faces and classify pose status (all angles)
        
        Args:
            img_path: Path to image file
        
        Returns:
            List of face dictionaries with pose information:
            [{
                'facial_area': {'x': x, 'y': y, 'w': w, 'h': h},
                'landmarks': {...},
                'confidence': 0.95,
                'yaw_angle': -45.2,
                'pitch_angle': 10.5,
                'roll_angle': -5.2,
                'pose_mode': 'profile_right_level_upright'
            }, ...]
        """
        faces = self.detect_faces_with_landmarks(img_path)
        
        results = []
        for face_key, face_data in faces.items():
            landmarks = face_data.get('landmarks', {})
            
            # Calculate all three angles
            yaw_angle = self.calculate_yaw_from_landmarks(landmarks)
            pitch_angle = self.calculate_pitch_from_landmarks(landmarks)
            roll_angle = self.calculate_roll_from_landmarks(landmarks)
            
            # Calculate face width (eye distance) for profile detection
            face_width = self.calculate_face_width_from_landmarks(landmarks)
            
            # Classify pose mode (using face width as additional indicator)
            pose_mode = self.classify_pose_mode(yaw_angle, pitch_angle, roll_angle, face_width)
            
            # Normalize facial_area format (RetinaFace returns list [x, y, w, h] or dict)
            facial_area_raw = face_data.get('facial_area', {})
            if isinstance(facial_area_raw, list) and len(facial_area_raw) >= 4:
                # Convert list [x, y, w, h] to dict format
                facial_area = {
                    'x': facial_area_raw[0],
                    'y': facial_area_raw[1],
                    'w': facial_area_raw[2],
                    'h': facial_area_raw[3]
                }
            elif isinstance(facial_area_raw, dict):
                # Already in dict format
                facial_area = facial_area_raw
            else:
                # Default to empty dict
                facial_area = {}
            
            result = {
                'facial_area': facial_area,
                'landmarks': landmarks,
                'confidence': face_data.get('confidence', 0.0),
                'yaw_angle': yaw_angle,
                'pitch_angle': pitch_angle,
                'roll_angle': roll_angle,
                'face_width': face_width,  # Eye distance in pixels
                'pose_mode': pose_mode
            }
            results.append(result)
        
        return results

