#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from rclpy.time import Time
from message_filters import Subscriber, ApproximateTimeSynchronizer  # 时间同步工具
from sensor_msgs.msg import Image, CameraInfo, PointCloud2, PointField
from std_msgs.msg import Header
from rclpy.executors import MultiThreadedExecutor #全线程执行器
import cv2
import numpy as np
import open3d as o3d
from cv_bridge import CvBridge
import struct

# -------------------------- 工具函数：Open3D 点云 ↔ ROS 2 PointCloud2 转换 --------------------------
def o3d_pointcloud_to_ros2(o3d_pcd, frame_id="camera_depth_optical_frame", stamp=None):
    """
    将 Open3D 点云转换为 ROS 2 PointCloud2 消息（仅黑白深度点云）
    :param o3d_pcd: Open3D PointCloud 对象
    :param frame_id: 坐标系ID（D435i 深度相机默认坐标系）
    :param stamp: 时间戳（与图像消息同步）
    :return: ROS 2 PointCloud2 消息
    """
    # 提取点云数据（仅 xyz，不含颜色）
    points = np.asarray(o3d_pcd.points)

    # 构造 PointCloud2 消息头部
    header = Header()
    header.frame_id = frame_id
    header.stamp = stamp if stamp else Time().to_msg()

    # 定义 PointField（仅 xyz 为 float32）
    fields = [
        PointField(name="x", offset=0, datatype=PointField.FLOAT32, count=1),
        PointField(name="y", offset=4, datatype=PointField.FLOAT32, count=1),
        PointField(name="z", offset=8, datatype=PointField.FLOAT32, count=1),
    ]

    # 构造点云数据缓冲区（每个点 12 字节：仅 xyz）
    point_step = 12
    num_points = len(points)
    data = bytearray(num_points * point_step)

    for i in range(num_points):
        # 写入 xyz 坐标（float32）
        struct.pack_into('<fff', data, i*point_step, points[i][0], points[i][1], points[i][2])

    # 构造 PointCloud2 消息
    pcl_msg = PointCloud2()
    pcl_msg.header = header
    pcl_msg.height = 1
    pcl_msg.width = num_points
    pcl_msg.fields = fields
    pcl_msg.is_bigendian = False
    pcl_msg.point_step = point_step
    pcl_msg.row_step = num_points * point_step
    pcl_msg.data = data
    pcl_msg.is_dense = True  # 无无效点（NaN）

    return pcl_msg

# -------------------------- ROS 2 节点：3D 点云生成与发布 --------------------------
class Open3DPointCloudNode(Node):
    def __init__(self):
        super().__init__("open3d_pointcloud_node")

        # 1. 初始化工具
        self.bridge = CvBridge()
        self.o3d_intrinsic = None  # Open3D 相机内参对象
        self.frame_id = "camera_depth_optical_frame"  # D435i 深度相机坐标系（与驱动保持一致）

        # 1.1 点云下采样参数
        self.voxel_size = 0.05  # 体素下采样大小：3cm
        self.random_sampling_ratio = 1.0  # 随机下采样保留比例：100%
        self.uniform_downsample_k = 3  # 均匀下采样：每隔3个点选1个
        self.depth_trunc = 5.0  # 深度截断：只生成5米以内的点云

        # 2. 订阅话题（深度图、彩色图、深度内参）
        self.sub_depth = Subscriber(self, Image, "/camera/camera/depth/image_rect_raw")
        self.sub_color = Subscriber(self, Image, "/camera/camera/color/image_raw")
        self.sub_camera_info = Subscriber(self, CameraInfo, "/camera/camera/depth/camera_info")

        # 3. 时间同步器（近似同步，容忍 100ms 时间差，适配 D435i 数据发布）
        self.ts = ApproximateTimeSynchronizer(
            [self.sub_depth, self.sub_color, self.sub_camera_info],
            queue_size=10,
            slop=0.1  # 时间同步容错率（秒）
        )
        self.ts.registerCallback(self.sync_callback)

        # 4. 发布 3D 点云话题（ROS 2 标准 PointCloud2）
        self.pub_pointcloud = self.create_publisher(
            PointCloud2,
            "/vision_3d/color_pointcloud",  # 自定义 3D 视觉话题名
            10
        )

        self.get_logger().info("Open3D 3D 点云节点已启动！")
        self.get_logger().info(f"体素下采样大小：{self.voxel_size}m，随机下采样保留比例：{self.random_sampling_ratio}，均匀下采样：每隔{self.uniform_downsample_k}个点选1个, 深度截断：{self.depth_trunc}m")
        self.get_logger().info("发布话题：/vision_3d/color_pointcloud")
    
    def generate_point_cloud(self, color_img_np, depth_img_np, depth_scale=1000.0, depth_trunc=3.0):
        """
        生成 3D 点云（修正尺寸不匹配问题）
        :param color_img_np: 彩色图像（numpy 数组）
        :param depth_img_np: 深度图像（numpy 数组）
        :param depth_scale: 深度缩放因子
        :param depth_trunc: 深度截断值
        :return: 点云对象
        """
        try:
            # 1. 将 numpy 数组转换为 Open3D Image 对象
            color_o3d = o3d.geometry.Image(color_img_np.astype(np.uint8))
            depth_o3d = o3d.geometry.Image(depth_img_np.astype(np.uint16))

            # 2. 核心修正：获取两张图像的尺寸，进行缩放对齐
            # 获取尺寸（open3d Image 尺寸格式：(高度, 宽度)）
            color_h, color_w = color_o3d.height, color_o3d.width
            depth_h, depth_w = depth_o3d.height, depth_o3d.width

            self.get_logger().info(f"原始尺寸：彩色图({color_h}×{color_w})，深度图({depth_h}×{depth_w})")

            # 3. 选择目标尺寸（这里选择深度图尺寸作为目标，也可以换成彩色图尺寸）
            target_h, target_w = depth_h, depth_w

            # 4. 对尺寸不匹配的图像进行缩放
            if (color_h, color_w) != (target_h, target_w):
                self.get_logger().info(f"正在将彩色图缩放至 {target_h}×{target_w}")
                # Open3D Image.resize() 方法：参数 (宽度, 高度)，插值方式默认线性插值
                color_o3d = color_o3d.resize((target_w, target_h))
            
            # （可选）如果想缩放深度图到彩色图尺寸，取消下面注释
            # if (depth_h, depth_w) != (color_h, color_w):
            #     self.get_logger().info(f"正在将深度图缩放至 {color_h}×{color_w}")
            #     depth_o3d = depth_o3d.resize((color_w, color_h))

            # 5. 生成 RGBD 图像（此时两张图像尺寸已一致）
            rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(
                color_o3d,
                depth_o3d,
                depth_scale=depth_scale,  # 深度值缩放（通常深度图像素值是实际距离的 1000 倍，如 1000 对应 1 米）
                depth_trunc=depth_trunc,  # 超过该距离的深度值会被截断
                convert_rgb_to_intensity=False  # 保留 RGB 颜色，不转换为灰度图
            )

            # 6. 生成点云（需要相机内参，这里用默认内参，你需要替换为自己相机的实际内参）
            pinhole_camera_intrinsic = o3d.camera.PinholeCameraIntrinsic(
                o3d.camera.PinholeCameraIntrinsicParameters.PrimeSenseDefault
            )
            # 若有实际内参，替换为如下格式（宽、高、fx、fy、cx、cy）
            # pinhole_camera_intrinsic = o3d.camera.PinholeCameraIntrinsic(
            #     target_w, target_h, fx=550.0, fy=550.0, cx=target_w/2, cy=target_h/2
            # )

            point_cloud = o3d.geometry.PointCloud.create_from_rgbd_image(
                rgbd_image,
                pinhole_camera_intrinsic
            )

            self.get_logger().info("3D 点云生成成功")
            return point_cloud

        except Exception as e:
            self.get_logger().error(f"3D 点云生成失败：{str(e)}")
            return None

    def sync_callback(self, depth_msg, color_msg, camera_info_msg):
        """
        同步回调函数：处理深度图、内参，生成并发布 3D 点云（黑白深度点云）
        :param depth_msg: 深度图 ROS 消息
        :param color_msg: 彩色图 ROS 消息（不使用，仅保留兼容性）
        :param camera_info_msg: 深度相机内参 ROS 消息
        """
        try:
            # -------------------------- Step 1：ROS 消息 → 数据格式转换 --------------------------
            # 1.1 深度图：ROS Image → OpenCV → Open3D（16UC1 格式，单位毫米 → 米）
            cv_depth = self.bridge.imgmsg_to_cv2(depth_msg, desired_encoding="16UC1")
            o3d_depth = o3d.geometry.Image(cv_depth.astype(np.uint16))
            # 转换深度单位：毫米（mm）→ 米（m）
            o3d_depth_np = np.asarray(o3d_depth).astype(np.float32) / 1000.0

            # 1.2 相机内参：ROS CameraInfo → Open3D PinholeCameraIntrinsic
            if self.o3d_intrinsic is None:
                # 从 CameraInfo 提取内参（焦距 fx/fy，主点 cx/cy，图像尺寸 width/height）
                fx = camera_info_msg.k[0]  # K 矩阵：[fx, 0, cx; 0, fy, cy; 0, 0, 1]
                fy = camera_info_msg.k[4]
                cx = camera_info_msg.k[2]
                cy = camera_info_msg.k[5]
                width = camera_info_msg.width
                height = camera_info_msg.height

                # 构造 Open3D 内参对象（只初始化一次，提升性能）
                self.o3d_intrinsic = o3d.camera.PinholeCameraIntrinsic(
                    width, height, fx, fy, cx, cy
                )
                self.get_logger().info(f"相机内参初始化完成：fx={fx:.2f}, fy={fy:.2f}, cx={cx:.2f}, cy={cy:.2f}")

            # -------------------------- Step 2：Open3D 生成 3D 点云（仅深度，黑白） --------------------------
            # 2.1 从深度图生成点云（只生成3米以内的点云，不使用彩色信息）
            rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(
                o3d.geometry.Image(o3d_depth_np),  # 使用深度图作为颜色（灰度）
                o3d.geometry.Image(o3d_depth_np),
                depth_scale=1.0,  # 已转换为米，缩放因子为 1
                depth_trunc=self.depth_trunc,  # 截断深度（只生成3米以内的点云）
                convert_rgb_to_intensity=True  # 转换为灰度，不保留彩色信息
            )

            # 2.2 生成黑白点云
            pcd = o3d.geometry.PointCloud.create_from_rgbd_image(
                rgbd_image,
                self.o3d_intrinsic
            )

            # 2.3 点云坐标校正（D435i 深度图与彩色图配准，翻转点云使其朝向正确）
            pcd.transform([[1, 0, 0, 0],
                           [0, -1, 0, 0],
                           [0, 0, -1, 0],
                           [0, 0, 0, 1]])

            # 2.4 点云下采样（多种方法组合使用）
            original_points = len(pcd.points)
            
            # 2.4.1 体素下采样（1cm）
            pcd = pcd.voxel_down_sample(voxel_size=self.voxel_size)
            voxel_points = len(pcd.points)
            
            # 2.4.2 随机下采样（保留70%）
            pcd = pcd.random_down_sample(sampling_ratio=self.random_sampling_ratio)
            random_points = len(pcd.points)
            
            # 2.4.3 均匀下采样（每隔3个点选1个）
            pcd = pcd.uniform_down_sample(every_k_points=self.uniform_downsample_k)
            final_points = len(pcd.points)
            
            self.get_logger().debug(f"点云下采样统计：原始{original_points}点 → 体素{voxel_points}点 → 随机{random_points}点 → 最终{final_points}点")

            # -------------------------- Step 3：Open3D 点云 → ROS 2 PointCloud2 发布 --------------------------
            # 转换为 ROS 2 消息（同步时间戳）
            pcl_msg = o3d_pointcloud_to_ros2(
                pcd,
                frame_id=self.frame_id,
                stamp=depth_msg.header.stamp  # 沿用深度图时间戳
            )

            # 发布 3D 点云话题
            self.pub_pointcloud.publish(pcl_msg)

        except Exception as e:
            self.get_logger().error(f"3D 点云生成失败：{str(e)}")

# -------------------------- 主函数 --------------------------
def main(args=None):
    rclpy.init(args=args)
    node = Open3DPointCloudNode()
    executor = MultiThreadedExecutor(num_threads=8)
    executor.add_node(node)
   # 多线程自旋节点（替代原来的 rclpy.spin(node)）
    try:
        executor.spin()
    finally:
        executor.shutdown()
        node.destroy_node()
        rclpy.shutdown()
    # 清理资源
    node.destroy_node()
    rclpy.shutdown()

if __name__ == "__main__":
    main()