#!/usr/bin/env python3
#
# Convert a depth image (16UC1 mm or 32FC1 m) + CameraInfo into a sensor_msgs/PointCloud2.
# This replaces depth_image_proc/point_cloud_xyz for environments where that package
# is not available or its executable is missing.

import numpy as np
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import CameraInfo, Image, PointCloud2, PointField
from std_msgs.msg import Header


class DepthToPointCloud(Node):
    def __init__(self):
        super().__init__('depth_to_pointcloud')

        self.declare_parameter('depth_topic', '/intel_realsense_r200_depth/depth/image_raw')
        self.declare_parameter('camera_info_topic', '/intel_realsense_r200_depth/camera_info')
        self.declare_parameter('points_topic', '/intel_realsense_r200_depth/depth/color/points')
        self.declare_parameter('max_depth_m', 10.0)

        depth_topic = self.get_parameter('depth_topic').value
        camera_info_topic = self.get_parameter('camera_info_topic').value
        points_topic = self.get_parameter('points_topic').value
        self.max_depth_m = float(self.get_parameter('max_depth_m').value)

        self.camera_info = None

        self.info_sub = self.create_subscription(
            CameraInfo, camera_info_topic, self.info_callback, qos_profile=1)
        self.depth_sub = self.create_subscription(
            Image, depth_topic, self.depth_callback, qos_profile=1)
        self.cloud_pub = self.create_publisher(PointCloud2, points_topic, qos_profile=1)

        self.get_logger().info(
            f'Waiting for camera info on {camera_info_topic}; '
            f'publishing points on {points_topic}'
        )

    def info_callback(self, msg: CameraInfo):
        self.camera_info = msg

    def depth_callback(self, msg: Image):
        if self.camera_info is None:
            self.get_logger().warn(
                'CameraInfo not received yet; skipping depth frame',
                throttle_duration_sec=5)
            return

        # Decode depth image according to its encoding.
        # Gz Sim may publish depth as either 16-bit unsigned mm (16UC1/mono16)
        # or 32-bit float meters (32FC1), depending on sensor/bridge version.
        if msg.encoding in ('16UC1', 'mono16'):
            depth_raw = np.frombuffer(msg.data, dtype=np.uint16).reshape(msg.height, msg.width)
            depth_m = depth_raw.astype(np.float32) * 0.001
        elif msg.encoding == '32FC1':
            depth_m = np.frombuffer(msg.data, dtype=np.float32).reshape(msg.height, msg.width)
        else:
            self.get_logger().error(
                f'Unsupported depth encoding: {msg.encoding}; expected 16UC1/mono16 or 32FC1',
                throttle_duration_sec=5)
            return

        # Intrinsics from CameraInfo
        K = self.camera_info.k
        fx, fy = float(K[0]), float(K[4])
        cx, cy = float(K[2]), float(K[5])

        # Pixel coordinates
        u = np.arange(msg.width, dtype=np.float32)
        v = np.arange(msg.height, dtype=np.float32)
        u, v = np.meshgrid(u, v)

        # Filter valid depth pixels
        valid = (depth_m > 0.0) & (depth_m < self.max_depth_m)
        u = u[valid]
        v = v[valid]
        z = depth_m[valid]

        if z.size == 0:
            return

        # Back-project to 3D (in the camera optical frame)
        x = (u - cx) * z / fx
        y = (v - cy) * z / fy

        points = np.stack((x, y, z), axis=-1).astype(np.float32)

        header = Header()
        header.stamp = msg.header.stamp
        header.frame_id = self.camera_info.header.frame_id

        cloud_msg = PointCloud2()
        cloud_msg.header = header
        cloud_msg.height = 1
        cloud_msg.width = int(points.shape[0])
        cloud_msg.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),
        ]
        cloud_msg.is_bigendian = False
        cloud_msg.point_step = 12
        cloud_msg.row_step = cloud_msg.point_step * cloud_msg.width
        cloud_msg.is_dense = True
        cloud_msg.data = points.tobytes()

        self.cloud_pub.publish(cloud_msg)


def main(args=None):
    rclpy.init(args=args)
    node = DepthToPointCloud()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
