import numpy as np
import cv2

# 左眼
LEFT_EYE = [33, 160, 158, 133, 153, 144]
# 右眼
RIGHT_EYE = [362, 385, 387, 263, 373, 380]
# 嘴唇 (内圈)
LIPS = [78, 95, 88, 178, 87, 14, 317, 402, 318, 324, 308, 415, 310, 311, 312, 13]

def _euclidean_distance(point1, point2):
    return np.linalg.norm(point1 - point2)

def calculate_ear(landmarks, width, height):
    """计算眼睛纵横比 EAR"""
    # 坐标转换
    points = np.array([(p.x * width, p.y * height) for p in landmarks])
    
    # 垂直距离
    v1 = _euclidean_distance(points[1], points[5])
    v2 = _euclidean_distance(points[2], points[4])
    # 水平距离
    h = _euclidean_distance(points[0], points[3])
    
    ear = (v1 + v2) / (2.0 * h)
    return ear

def calculate_mar(landmarks, width, height):
    """计算嘴巴纵横比 MAR"""
    points = np.array([(p.x * width, p.y * height) for p in landmarks])
    pass 

def calculate_mar_simple(top, bottom, left, right):
    h = _euclidean_distance(top, bottom)
    w = _euclidean_distance(left, right)
    return h / w

# geometry_utils.py 中的 estimate_head_pose 函数替换为以下内容

def estimate_head_pose(landmarks, width, height):
    """
    计算头部姿态 (Pitch, Yaw, Roll)
    返回单位：角度 (Degree)
    """
    # 3D 模型点 (标准人脸模型)
    model_points = np.array([
        (0.0, 0.0, 0.0),             # Nose tip
        (0.0, -330.0, -65.0),        # Chin
        (-225.0, 170.0, -135.0),     # Left eye left corner
        (225.0, 170.0, -135.0),      # Right eye right corner
        (-150.0, -150.0, -125.0),    # Left Mouth corner
        (150.0, -150.0, -125.0)      # Right mouth corner
    ])

    # MediaPipe 对应的关键点索引
    idx_list = [1, 152, 33, 263, 61, 291]
    
    image_points = []
    for idx in idx_list:
        p = landmarks[idx]
        image_points.append((p.x * width, p.y * height))
    
    image_points = np.array(image_points, dtype="double")

    focal_length = width
    center = (width / 2, height / 2)
    camera_matrix = np.array(
        [[focal_length, 0, center[0]],
         [0, focal_length, center[1]],
         [0, 0, 1]], dtype="double"
    )
    dist_coeffs = np.zeros((4, 1)) 

    # 求解PnP
    success, rotation_vector, translation_vector = cv2.solvePnP(
        model_points, image_points, camera_matrix, dist_coeffs
    )

    rmat, _ = cv2.Rodrigues(rotation_vector)
    angles, mtxR, mtxQ, Qx, Qy, Qz = cv2.RQDecomp3x3(rmat)

    pitch = angles[0]
    yaw   = angles[1]
    roll  = angles[2]

    if pitch < -180: pitch += 360
    if pitch > 180:  pitch -= 360
    pitch = 180 - pitch if pitch > 0 else -pitch - 180
    if yaw < -180: yaw += 360
    if yaw > 180:  yaw -= 360
    if roll < -180: roll += 360
    if roll > 180:  roll -= 360

    return pitch, yaw, roll