import cv2
import dlib
import numpy as np
import urx
import time
import math
from datetime import datetime
from math import cos, sin

# 配置参数
UR_ROBOT_IP = "192.168.10.99"
CAMERA_INDEX = 0  # 默认摄像头索引
FILTER_WINDOW = 5  # 移动平均窗口大小
DAMPING = 0.1  # 阻尼系数
SAFE_DISTANCE = 0.8  # 相机到人脸的安全距离(米)
Y_CONSTRAINT = 0.4  # y轴固定位置(米)
CONTROL_FREQ = 0.05  # 控制周期(秒)

# 3D面部特征点模型（世界坐标系）
MODEL_POINTS = np.array([
    (0.0, 0.0, 0.0),            # 鼻尖
    (0.0, -330.0, -65.0),       # 下巴
    (-225.0, 170.0, -135.0),    # 左眼角
    (225.0, 170.0, -135.0),     # 右眼角
    (-150.0, -150.0, -125.0),   # 左嘴角
    (150.0, -150.0, -125.0)     # 右嘴角
])

class HeadPoseEstimator:
    def __init__(self):
        # 初始化dlib人脸检测器和特征点预测器
        self.detector = dlib.get_frontal_face_detector()
        self.predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
        self.camera = cv2.VideoCapture(CAMERA_INDEX)
        
        # 相机内参（需要根据实际相机校准）
        self.camera_matrix = np.array(
            [[640, 0, 320],
             [0, 640, 240],
             [0, 0, 1]], dtype="double"
        )
        self.dist_coeffs = np.zeros((4, 1))  # 畸变系数

    def get_head_pose(self):
        """获取头部姿态：返回(旋转角, 位置)"""
        ret, frame = self.camera.read()
        if not ret:
            return None
            
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = self.detector(gray)
        
        if len(faces) == 0:
            return None
            
        # 取第一个检测到的人脸
        shape = self.predictor(gray, faces[0])
        image_points = self._shape_to_np(shape)
        
        # 提取用于姿态估计的6个特征点
        selected_points = [30, 8, 36, 45, 48, 54]  # 对应MODEL_POINTS的索引
        image_points = image_points[selected_points]
        
        # 求解PnP问题
        success, rotation_vector, translation_vector = cv2.solvePnP(
            MODEL_POINTS, image_points, self.camera_matrix, self.dist_coeffs,
            flags=cv2.SOLVEPNP_ITERATIVE
        )
        
        # 旋转向量转欧拉角（弧度）
        rotation_matrix, _ = cv2.Rodrigues(rotation_vector)
        euler_angles = self._rotation_matrix_to_euler_angles(rotation_matrix)
        
        # 返回(旋转角, 位置)：(yaw, pitch, roll, x, y, z)
        return np.concatenate([euler_angles, translation_vector.flatten() / 1000.0])  # 毫米转米

    def _shape_to_np(self, shape):
        """将dlib特征点转换为numpy数组"""
        return np.array([(p.x, p.y) for p in shape.parts()], dtype="double")

    def _rotation_matrix_to_euler_angles(self, R):
        """旋转矩阵转欧拉角(Z-Y-X顺序)"""
        sy = math.sqrt(R[0,0] * R[0,0] + R[1,0] * R[1,0])
        singular = sy < 1e-6

        if not singular:
            x = math.atan2(R[2,1], R[2,2])
            y = math.atan2(-R[2,0], sy)
            z = math.atan2(R[1,0], R[0,0])
        else:
            x = math.atan2(-R[1,2], R[1,1])
            y = math.atan2(-R[2,0], sy)
            z = 0

        return np.array([z, y, x])  # yaw, pitch, roll

    def release(self):
        self.camera.release()

class URHeadTracker:
    def __init__(self):
        self.robot = urx.Robot(UR_ROBOT_IP)
        self.estimator = HeadPoseEstimator()
        self.prev_deltas = []  # 用于移动平均滤波
        self.current_pose = self.robot.getl()
        print(f"机械臂初始位姿: {self.current_pose}")

    def compute_face_direction(self, head_rot):
        """计算面部方向向量"""
        yaw, pitch, roll = head_rot
        # 方向向量计算 (单位向量)
        dir_x = sin(yaw) * cos(pitch)
        dir_z = cos(yaw) * cos(pitch)
        return np.array([dir_x, 0, dir_z])  # y方向约束为0

    def damping_least_squares(self, delta, current_joints):
        """阻尼最小二乘法优化关节运动"""
        # 简化实现：基于关节当前位置添加阻尼
        joint_weights = np.array([1.0, 1.0, 1.0, 0.5, 0.5, 0.5])  # 末端关节权重降低
        damping = DAMPING * np.eye(6)
        return delta / (np.linalg.norm(joint_weights) + damping.diagonal())

    def moving_average_filter(self, delta):
        """移动平均滤波"""
        self.prev_deltas.append(delta)
        if len(self.prev_deltas) > FILTER_WINDOW:
            self.prev_deltas.pop(0)
        return np.mean(self.prev_deltas, axis=0)

    def compute_target_pose(self, head_pose):
        """计算相机目标位姿"""
        head_rot = head_pose[:3]  # yaw, pitch, roll
        head_pos = head_pose[3:]  # x, y, z (米)

        # 计算目标位置
        face_dir = self.compute_face_direction(head_rot)
        target_x = head_pos[0] - face_dir[0] * SAFE_DISTANCE
        target_z = head_pos[2] - face_dir[2] * SAFE_DISTANCE

        # 计算目标姿态（相机对准人脸）
        target_rx = -head_rot[1]  # 俯仰角反向
        target_ry = 0  # 消除侧倾
        target_rz = -head_rot[0]  # 偏航角反向

        return np.array([target_x, Y_CONSTRAINT, target_z, target_rx, target_ry, target_rz])

    def run(self):
        try:
            while True:
                # 获取头部姿态
                head_pose = self.estimator.get_head_pose()
                if head_pose is None:
                    print("未检测到人脸，等待...")
                    time.sleep(CONTROL_FREQ)
                    continue

                # 计算目标位姿
                target_pose = self.compute_target_pose(head_pose)
                
                # 计算相对位移
                current = np.array(self.current_pose)
                delta = target_pose - current
                
                # 应用阻尼最小二乘法
                current_joints = self.robot.getj()
                delta = self.damping_least_squares(delta, current_joints)
                
                # 移动平均滤波
                delta_smoothed = self.moving_average_filter(delta)
                
                # 约束y轴不移动
                delta_smoothed[1] = 0
                
                # 执行运动
                new_pose = current + delta_smoothed
                self.robot.movel(new_pose.tolist(), acc=0.05, vel=0.05, relative=False)
                self.current_pose = self.robot.getl()
                
                # 打印状态
                print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]")
                print(f"目标位姿: {target_pose}")
                print(f"当前位姿: {self.current_pose}\n")
                
                time.sleep(CONTROL_FREQ)

        except KeyboardInterrupt:
            print("用户终止程序")
        finally:
            self.robot.close()
            self.estimator.release()
            cv2.destroyAllWindows()

if __name__ == "__main__":
    # 提示：需要先下载dlib特征点模型
    # 下载地址：http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2
    tracker = URHeadTracker()
    tracker.run()