"""头部姿态估计（cv2.solvePnP）+ 正面照校验。

用通用 3D 头模与 6 个 MediaPipe 关键点求解欧拉角（yaw/pitch/roll，单位：度），
阈值即可写成业务可读的「yaw>15° 拒绝」，并把角度返回前端做拍照引导。
详见技术方案 §9。
"""
import os

import cv2
import numpy as np

from face_analysis.face_mesh_landmarks import PNP_INDICES

# 正面照判定阈值（度），可由环境变量覆盖，便于上线后按真实数据标定（见技术方案 §11）。
# ⚠️ 标定说明：基于通用 6 点 3D 头模 + solvePnP，对明显正面但相机略带俯仰/个体
#    脸型差异的真实照片，解出的 yaw/pitch 常落在 15~25°（roll 较稳定，多在 5° 内）。
#    因此默认阈值放宽到 30°，只拦截明显侧脸（真实侧脸 yaw 通常 40°+），
#    避免误杀正常上传图。生产可通过环境变量随时收紧/放宽，无需改代码。
YAW_THRESHOLD = float(os.getenv("FRONTAL_YAW_THR", "30"))
PITCH_THRESHOLD = float(os.getenv("FRONTAL_PITCH_THR", "30"))
ROLL_THRESHOLD = float(os.getenv("FRONTAL_ROLL_THR", "30"))

# 通用 3D 头部模型（单位 mm，近似），与 PNP_INDICES 一一对应：
# 鼻尖(1) / 下巴(152) / 左眼外角(33) / 右眼外角(263) / 左嘴角(61) / 右嘴角(291)
# ⚠️ 采用「相机坐标系」约定：x 向右、y 向下、z 向场景内（远离观察者）。
#    与 MediaPipe 像素坐标（y 下）一致，且 +z 指向人脸背面，
#    这样正面照解出的旋转矩阵≈单位阵，欧拉角≈0。
#    若只翻 y 不翻 z（或都不翻），会残留 ~180° 翻转使正面图被误判。
_MODEL_POINTS = np.array([
    (0.0,    0.0,    0.0),     # 鼻尖
    (0.0,    63.6,   12.5),    # 下巴（在鼻尖下方 → y 正）
    (-43.3, -32.7,   26.0),    # 左眼外角（在鼻尖上方 → y 负，且凹于鼻尖 → z 正）
    (43.3,  -32.7,   26.0),    # 右眼外角
    (-28.9,  28.9,   24.1),    # 左嘴角
    (28.9,   28.9,   24.1),    # 右嘴角
], dtype=np.float64)


def estimate_head_pose(landmarks, image_width, image_height):
    """求解头部欧拉角，返回 (yaw, pitch, roll)（度）。solvePnP 失败返回 None。"""
    lm = landmarks.landmark if hasattr(landmarks, "landmark") else landmarks
    image_points = np.array([
        (lm[i].x * image_width, lm[i].y * image_height)
        for i in PNP_INDICES
    ], dtype=np.float64)

    focal = float(image_width)  # 近似焦距
    cam_matrix = np.array([[focal, 0, image_width / 2],
                           [0, focal, image_height / 2],
                           [0, 0, 1]], dtype=np.float64)
    dist = np.zeros((4, 1))  # 假设无畸变

    success, rvec, _tvec = cv2.solvePnP(
        _MODEL_POINTS, image_points, cam_matrix, dist,
        flags=cv2.SOLVEPNP_ITERATIVE,
    )
    if not success:
        return None
    rot, _ = cv2.Rodrigues(rvec)
    # 在「相机坐标系」(x右 y下 z内) 下抽取 Tait-Bryan 欧拉角，物理含义对齐：
    #   yaw   = 绕 Y（竖轴）转 → 左右扭头
    #   pitch = 绕 X（横轴）转 → 上下点头
    #   roll  = 绕 Z（光轴）转 → 面内倾斜
    sy = (rot[0, 0] ** 2 + rot[1, 0] ** 2) ** 0.5
    yaw = float(np.degrees(np.arctan2(-rot[2, 0], sy)))
    pitch = float(np.degrees(np.arctan2(rot[2, 1], rot[2, 2])))
    roll = float(np.degrees(np.arctan2(rot[1, 0], rot[0, 0])))
    return yaw, pitch, roll


def check_frontal_face(landmarks, image_width, image_height,
                       yaw_thr=YAW_THRESHOLD, pitch_thr=PITCH_THRESHOLD,
                       roll_thr=ROLL_THRESHOLD):
    """正面照判定：yaw/pitch/roll 绝对值均在阈值内才算正面。

    solvePnP 解算失败时返回 True（不拦截，交由后续逻辑），避免误杀。
    """
    pose = estimate_head_pose(landmarks, image_width, image_height)
    if pose is None:
        return True
    yaw, pitch, roll = pose
    return abs(yaw) <= yaw_thr and abs(pitch) <= pitch_thr and abs(roll) <= roll_thr


if __name__ == "__main__":
    import sys
    from face_analysis.detector import detector

    path = sys.argv[1] if len(sys.argv) > 1 else "tests/fixtures/frontal.jpg"
    img = cv2.imread(path)
    if img is None:
        print(f"无法读取图片: {path}")
        sys.exit(1)
    h, w = img.shape[:2]
    lms = detector.detect(img)
    if lms is None:
        print("未检出人脸")
        sys.exit(1)
    yaw, pitch, roll = estimate_head_pose(lms, w, h)
    frontal = check_frontal_face(lms, w, h)
    print(f"yaw={yaw:.2f} pitch={pitch:.2f} roll={roll:.2f} frontal={frontal}")