import rclpy
from rclpy.node import Node

import cv2
import numpy as np

from sensor_msgs.msg import Image
from geometry_msgs.msg import PointStamped
from geometry_msgs.msg import PoseStamped
from cv_bridge import CvBridge
from sensor_msgs.msg import CameraInfo
import mediapipe as mp
from scipy.spatial.transform import Rotation as R
import time
# from tf_transformations import quaternion_from_euler, euler_from_quaternion


class StereoHandDepthNode(Node):

    def __init__(self):
        super().__init__('stereo_hand_depth_node')
        self.marker_timeout = 0.3   # 秒，建议 0.2 ~ 0.5
        self.last_valid_pos = None
        self.last_valid_quat = None
        self.target_pos = np.array([0.0, 0.0, 0.0])
        self.target_quat = np.array([1.0, 0.0, 0.0, 0.0])

        # ---------- 参数 ----------
        self.fx = 600.0   # 焦距（像素）👉 后面用标定值替换
        self.fy = 600.0
        self.cx = 320.0
        self.cy = 240.0
        self.baseline = 0.06  # 双目基线（米）👉 换成你真实的

        # ---------- OpenCV ----------
        self.numDisparities = 128  # 必须是16的倍数
        self.bridge = CvBridge()

        self.stereo = cv2.StereoSGBM_create(
            minDisparity=0,
            numDisparities=self.numDisparities,
            blockSize=7,
            P1=8 * 3 * 7 ** 2,
            P2=32 * 3 * 7 ** 2,
            disp12MaxDiff=1,
            uniquenessRatio=10,
            speckleWindowSize=100,
            speckleRange=32
        )

        self.aruco_dict = cv2.aruco.getPredefinedDictionary(
            cv2.aruco.DICT_4X4_50
        )
        self.aruco_params = cv2.aruco.DetectorParameters()

        self.marker_size = 0.058  # 米

        # ---------- ROS ----------
        self.left_info = None
        self.right_info = None
        self.sub = self.create_subscription(
            Image,
            '/image_1',
            self.image_callback,
            10
        )
        self.sub = self.create_subscription(
            CameraInfo,
            '/stereo_camera/left/camera_info',
            self.left_camera_info_callback,
            10
        )
        self.sub = self.create_subscription(
            CameraInfo,
            '/stereo_camera/right/camera_info',
            self.right_camera_info_callback,
            10
        )

        self.pub = self.create_publisher(
            PoseStamped,
            '/hand_3d',
            10
        )

        # self.mp_hands = mp.solutions.hands.Hands(
        #     static_image_mode=False,
        #     max_num_hands=1,
        #     min_detection_confidence=0.4,
        #     min_tracking_confidence=0.4
        # )
        self.get_logger().info('StereoHandDepthNode started')

    def init_rectify_maps(self, left_info, right_info, image_size):
        self.K1 = np.array(left_info.k).reshape(3, 3)
        self.D1 = np.array(left_info.d)
        self.K2 = np.array(right_info.k).reshape(3, 3)
        self.D2 = np.array(right_info.d)

        self.R1 = np.array(left_info.r).reshape(3, 3)
        self.R2 = np.array(right_info.r).reshape(3, 3)

        self.P1 = np.array(left_info.p).reshape(3, 4)
        self.P2 = np.array(right_info.p).reshape(3, 4)

        map1_l, map2_l = cv2.initUndistortRectifyMap(
            self.K1, self.D1, self.R1, self.P1[:3, :3], image_size, cv2.CV_32FC1
        )

        map1_r, map2_r = cv2.initUndistortRectifyMap(
            self.K2, self.D2, self.R2, self.P2[:3, :3], image_size, cv2.CV_32FC1
        )   
        return map1_l, map2_l, map1_r, map2_r

    def left_camera_info_callback(self, msg: CameraInfo):
        self.left_info = msg

    def right_camera_info_callback(self, msg: CameraInfo):
        self.right_info = msg

    def image_callback(self, msg: Image):

        if self.left_info is None or self.right_info is None:
            return
        self.map1_l, self.map2_l, self.map1_r, self.map2_r = self.init_rectify_maps(self.left_info, self.right_info, (640, 480))
        self.fx = self.left_info.p[0]
        self.fy = self.left_info.p[5]
        self.cx = self.left_info.p[2]
        self.cy = self.left_info.p[6]

        self.baseline = -self.right_info.p[3] / self.fx
        # self.get_logger().info(f'baseline = {self.baseline:.4f} m')
        # 转 OpenCV
        frame = self.bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')

        h, w, _ = frame.shape
        mid = w // 2
        # print("Frame size:", w, h)

        # 左右拆分
        left = frame[:, :mid]
        right = frame[:, mid:]
        left_rect = cv2.remap(left, self.map1_l, self.map2_l, cv2.INTER_LINEAR)
        # right_rect = cv2.remap(right, self.map1_r, self.map2_r, cv2.INTER_LINEAR)
        # results = self.mp_hands.process(left_rect)
        # p = PointStamped()
        # if results.multi_hand_landmarks:
        #     wrist = results.multi_hand_landmarks[0].landmark[0]

        #     u = int(wrist.x * (w / 2))  # 注意这里是左图宽度
        #     v = int(wrist.y * h)
        #     p.header = msg.header
        #     p.point.x = float(u)
        #     p.point.y = float(v)
        #     p.point.z = 0.0  # 暂时不用

        # 转灰度
        # left_gray = cv2.cvtColor(left_rect, cv2.COLOR_BGR2GRAY)
        # right_gray = cv2.cvtColor(right_rect, cv2.COLOR_BGR2GRAY)


        # 计算视差
        # disparity = self.stereo.compute(left_gray, right_gray).astype(np.float32) / 16.0

        corners, ids, _ = cv2.aruco.detectMarkers(
            left_rect, self.aruco_dict, parameters=self.aruco_params
        )
        msg_out = PoseStamped()
        if ids is not None:
            cv2.aruco.drawDetectedMarkers(left_rect, corners, ids)
            # c = corners[0][0]
            # center = c.mean(axis=0)
            # u, v = int(center[0]), int(center[1])
            # d = disparity[v, u]

            # cv2.circle(left_gray, (u, v), 5, (0, 0, 255), -1)
            # cv2.putText(
            #     left_gray, f"disp={d:.2f}", (u+5, v-5),
            #     cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 1
            # )
            for e in corners:
                ret, rvec, tvec = cv2.solvePnP(
                    objectPoints=np.array([
                        [-self.marker_size/2,  self.marker_size/2, 0],
                        [ self.marker_size/2,  self.marker_size/2, 0],
                        [ self.marker_size/2, -self.marker_size/2, 0],
                        [-self.marker_size/2, -self.marker_size/2, 0]
                    ], dtype=np.float32),
                    imagePoints=e,
                    cameraMatrix=self.K1,
                    distCoeffs=self.D1
                )
                if(ret == True):
                    cv2.drawFrameAxes(left_rect, self.K1, self.D1, rvec, tvec, 0.1, 2)
                    R_cv, _ = cv2.Rodrigues(rvec)
                    R_cv_to_ros = np.array([
                        [ 0,  0,  -1],
                        [ 1,  0,  0],
                        [ 0, -1,  0]
                    ])
                    R_ros = R_cv_to_ros @ R_cv
                    t_ros = R_cv_to_ros @ tvec.reshape(3)
                    # t_ros = tvec.reshape(3)
                    rot = R.from_matrix(R_ros)
                    
                    qx, qy, qz, qw = rot.as_quat()
                    roll, pitch, yaw = rot.as_euler('xyz')
                    self.get_logger().info(f"quat: ({qx}, {qy}, {qz}, {qw})")
                    self.get_logger().info(f"euler: ({roll*180/3.14:.1f}, {pitch*180/3.14:.1f}, {yaw*180/3.14:.1f})")

                    # 记录时间
                    self.last_marker_time = time.time()

                    # 记录最后一次有效目标
                    self.last_valid_pos = np.array([t_ros[0] + 0.7, t_ros[1], t_ros[2]])
                    self.last_valid_quat = np.array([qx, qy, qz, qw])
                    
            # u = int(u * (w / 2))  # 注意这里是左图宽度
            # v = int(v * h)
            # p.header = msg.header
            # p.point.x = float(u)
            # p.point.y = float(v)
            # p.point.z = 0.0  # 暂时不用
        cv2.imshow("left", left_rect)
        # ----------- 这里先用“假手点”（画面中心）-----------
        # x = int(p.point.x)
        # y = int(p.point.y)
        # cv2.circle(left_gray, (x, y), 5, (255, 0, 0), -1)
        # print("left_gray.shape:", left_gray.shape)
        # cv2.imshow('left', left_gray)
        # print(f'Wrist pixel: ({x}, {y})')
        # ROI 取中值，稳很多
        # roi = disparity[y-3:y+3, x-3:x+3]
        # d = np.nanmedian(roi)
        # print(f'Disparity at wrist: {roi}')

        # if d <= 0:
        #     return

        # 深度
        # Z = (self.fx * self.baseline) / d
        # X = (x - self.cx) * Z / self.fx
        # Y = (y - self.cy) * Z / self.fy
        now = time.time()
        marker_valid = (
            self.last_valid_pos is not None and
            (now - self.last_marker_time) < self.marker_timeout
        )

        if marker_valid:
            self.target_pos = self.last_valid_pos
            self.target_quat = self.last_valid_quat
        else:
            # 🔥 marker 丢失 → 保持最后位姿
            if self.last_valid_pos is None:
                cv2.waitKey(1)
                return  # 连一次都没看到 marker，啥也不做
            self.target_pos = self.last_valid_pos
            self.target_quat = self.last_valid_quat
        # 发布
        
        msg_out.header.stamp = self.get_clock().now().to_msg()
        msg_out.header.frame_id = 'camera_link'  
        # msg_out.pose.position.x = float(Z)
        # msg_out.pose.position.y = float(X)
        # msg_out.pose.position.z = -float(Y)
        # self.get_logger().info(f"X:{X:.3f}, Y:{Y:.3f}, Z:{Z:.3f}")
        msg_out.pose.orientation.x = float(self.target_quat[0])
        msg_out.pose.orientation.y = float(self.target_quat[1])
        msg_out.pose.orientation.z = float(self.target_quat[2])
        msg_out.pose.orientation.w = float(self.target_quat[3])
        msg_out.pose.position.x = float(self.target_pos[0])
        msg_out.pose.position.y = float(self.target_pos[1])
        msg_out.pose.position.z = float(self.target_pos[2])
        self.pub.publish(msg_out)

        # ----------- 可视化（调试用）-----------
        # disp_vis = cv2.normalize(disparity, None, 0, 255, cv2.NORM_MINMAX)
        # disp_vis = disp_vis.astype(np.uint8)
        # cv2.circle(disp_vis, (x, y), 5, 255, -1)
        # disp_vis = disp_vis[:, self.numDisparities:]
        # cv2.imshow('disparity', disp_vis)
        cv2.waitKey(1)


def main(args=None):
    rclpy.init(args=args)
    node = StereoHandDepthNode()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()


if __name__ == '__main__':
    main()
