"""
depth_confidence_filter_node.py — Depth confidence filter for D435i.

Subscribes:
  /camera/depth/image_rect_raw      sensor_msgs/Image  float32 depth (m)
  /camera/depth/confidence          sensor_msgs/Image  uint8 confidence (0-255)

Publishes:
  /camera/depth/filtered            sensor_msgs/Image  float32 depth, low-confidence
                                                        pixels zeroed

The D435i outputs a per-pixel confidence map alongside the depth image.
Values below the threshold (default 0.5, normalised from 0–255) are replaced
with 0.0 (no-return) so downstream nodes (VO, RTAB-Map) treat them as invalid.

Parameters:
  confidence_threshold  float  0.5    (0.0–1.0; raw uint8 value = threshold * 255)
  queue_size            int    5
"""

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy
import message_filters
import numpy as np
from cv_bridge import CvBridge
from sensor_msgs.msg import Image


_SENSOR_QOS = QoSProfile(
    reliability=ReliabilityPolicy.BEST_EFFORT,
    history=HistoryPolicy.KEEP_LAST,
    depth=5,
)


class DepthConfidenceFilterNode(Node):

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

        self.declare_parameter('confidence_threshold', 0.5)
        self.declare_parameter('queue_size', 5)

        thr_norm   = self.get_parameter('confidence_threshold').value
        queue_size = self.get_parameter('queue_size').value

        # Raw uint8 threshold (D435i confidence map is 0–255)
        self._thr = int(thr_norm * 255)

        self._bridge = CvBridge()

        depth_sub = message_filters.Subscriber(
            self, Image, '/camera/depth/image_rect_raw', qos_profile=_SENSOR_QOS
        )
        conf_sub = message_filters.Subscriber(
            self, Image, '/camera/depth/confidence', qos_profile=_SENSOR_QOS
        )
        self._sync = message_filters.ApproximateTimeSynchronizer(
            [depth_sub, conf_sub], queue_size=queue_size, slop=0.01
        )
        self._sync.registerCallback(self._on_frame)

        self._pub = self.create_publisher(Image, '/camera/depth/filtered', 10)

        self.get_logger().info(
            f'depth_confidence_filter ready — threshold={thr_norm:.2f} ({self._thr}/255)'
        )

    def _on_frame(self, depth_msg: Image, conf_msg: Image) -> None:
        try:
            depth = self._bridge.imgmsg_to_cv2(depth_msg, desired_encoding='32FC1')
            conf  = self._bridge.imgmsg_to_cv2(conf_msg,  desired_encoding='mono8')
        except Exception as exc:
            self.get_logger().warning(f'decode error: {exc}')
            return

        filtered = np.where(conf >= self._thr, depth, 0.0).astype(np.float32)

        out = self._bridge.cv2_to_imgmsg(filtered, encoding='32FC1')
        out.header = depth_msg.header
        self._pub.publish(out)


def main(args=None):
    rclpy.init(args=args)
    node = DepthConfidenceFilterNode()
    try:
        rclpy.spin(node)
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
