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 cv_bridge import CvBridge


class StereoHandDepthNode(Node):

    def __init__(self):
        super().__init__('stereo_hand_depth_node')

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

        # ---------- OpenCV ----------
        self.bridge = CvBridge()

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

        # ---------- ROS ----------
        # self.left = self.create_subscription(
        #     Image,
        #     '/camera/left/image_raw',
        #     self.image_left_callback,
        #     10
        # )
        # self.right = self.create_subscription(
        #     Image,
        #     '/camera/right/image_raw',
        #     self.image_right_callback,
        #     10
        # )
        self.image = self.create_subscription(
            Image,
            '/image',
            self.image_callback,
            10
        )
        self.pub = self.create_publisher(
            PointStamped,
            '/hand_3d',
            10
        )

        self.get_logger().info('StereoHandDepthNode started')

    def image_callback(self, msg: Image):
        # 转 OpenCV
        frame = self.bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')

        h, w, _ = frame.shape
        mid = w // 2

        # 左右拆分
        left = frame[:, :mid]
        right = frame[:, mid:]
        cv2.imshow("left", left)
        cv2.imshow("right", right)

        # 转灰度
        left_gray = cv2.cvtColor(left, cv2.COLOR_BGR2GRAY)
        right_gray = cv2.cvtColor(right, cv2.COLOR_BGR2GRAY)

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

        # ----------- 这里先用“假手点”（画面中心）-----------
        x = int(self.cx)
        y = int(self.cy)

        # ROI 取中值，稳很多
        roi = disparity[y-3:y+3, x-3:x+3]
        d = np.nanmedian(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

        # 发布
        msg_out = PointStamped()
        msg_out.header.stamp = self.get_clock().now().to_msg()
        msg_out.header.frame_id = 'camera_link'
        msg_out.point.x = float(X)
        msg_out.point.y = float(Y)
        msg_out.point.z = float(Z)

        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)

        cv2.imshow('disparity', disp_vis)
        cv2.waitKey(10)


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


if __name__ == '__main__':
    main()
