import numpy as np
import cv2

# 左眼
LEFT_EYE = [33, 160, 158, 133, 153, 144]
# 右眼
RIGHT_EYE = [362, 385, 387, 263, 373, 380]
# 左眼虹膜关键点索引
LEFT_EYE_GAZE_IDXS = [33, 133, 159, 145, 468]
# 右眼虹膜关键点索引
RIGHT_EYE_GAZE_IDXS = [263, 362, 386, 374, 473] 
# 嘴唇 (内圈)
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_iris_pos(landmarks, indices, width, height):
    p_left = np.array([landmarks[indices[0]].x * width, landmarks[indices[0]].y * height])
    p_right = np.array([landmarks[indices[1]].x * width, landmarks[indices[1]].y * height])
    p_top = np.array([landmarks[indices[2]].x * width, landmarks[indices[2]].y * height])
    p_bottom = np.array([landmarks[indices[3]].x * width, landmarks[indices[3]].y * height])
    p_iris = np.array([landmarks[indices[4]].x * width, landmarks[indices[4]].y * height])

    # 修改为欧几里得距离计算
    eye_width = _euclidean_distance(p_left, p_right)
    iris_left = _euclidean_distance(p_iris, p_left)
    if eye_width < 1e-3: 
        raw_x = 0.5
    else :
        raw_x = iris_left / eye_width
    eye_height = _euclidean_distance(p_bottom, p_top)
    iris_top = _euclidean_distance(p_iris, p_top)
    if eye_height < 1e-3:
        raw_y = 0.5
    else:
        raw_y = iris_top / eye_height

    # x_min = 0.4
    # x_max = 0.6
    # ratio_x = abs(raw_x - x_min) / (x_max - x_min)
    # y_min = 0.2
    # y_max = 0.8
    # ratio_y = abs(raw_y - y_min) / (y_max - y_min)
    ratio_x = raw_x
    ratio_y = raw_y
    
    return max(0.0, min(1.0, ratio_x)), max(0.0, min(1.0, ratio_y))


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