#!/usr/bin/env python3

import rospy
import pyrealsense2 as rs
import numpy as np

from sensor_msgs.msg import Image, CameraInfo
from cv_bridge import CvBridge


def main():
    rospy.init_node("depth_publisher")

    bridge = CvBridge()

    depth_pub = rospy.Publisher(
        "/camera/depth/image_raw",
        Image,
        queue_size=1
    )

    color_pub = rospy.Publisher(
        "/camera/color/image_raw",
        Image,
        queue_size=1
    )

    info_pub = rospy.Publisher(
        "/camera/depth/camera_info",
        CameraInfo,
        queue_size=1
    )

    # RealSense pipeline
    pipeline = rs.pipeline()
    config = rs.config()

    width = 424
    height = 240
    fps = 30

    config.enable_stream(
        rs.stream.depth,
        width,
        height,
        rs.format.z16,
        fps
    )
    config.enable_stream(
        rs.stream.color,
        640,
        480,
        rs.format.rgb8,
        fps
    )

    profile = pipeline.start(config)

    # Lấy intrinsic của camera
    depth_stream = profile.get_stream(rs.stream.depth)
    intr = depth_stream.as_video_stream_profile().get_intrinsics()

    rospy.loginfo("Depth camera started.")

    rate = rospy.Rate(fps)

    while not rospy.is_shutdown():

        frames = pipeline.wait_for_frames()
        depth = frames.get_depth_frame()
        color = frames.get_color_frame()

        if not depth or not color:
            continue

        depth_image = np.asanyarray(depth.get_data())
        color_image = np.asanyarray(color.get_data())

        # Image message
        img_msg = bridge.cv2_to_imgmsg(depth_image, encoding="16UC1")
        img_msg.header.stamp = rospy.Time.now()
        img_msg.header.frame_id = "camera_depth_optical_frame"

        color_msg = bridge.cv2_to_imgmsg(color_image, encoding="rgb8")
        color_msg.header.stamp = img_msg.header.stamp
        color_msg.header.frame_id = "camera_color_optical_frame"

        # CameraInfo message
        info_msg = CameraInfo()
        info_msg.header = img_msg.header

        info_msg.width = intr.width
        info_msg.height = intr.height

        info_msg.distortion_model = "plumb_bob"

        # Thông số méo (D)
        info_msg.D = list(intr.coeffs)

        # Camera matrix (K)
        info_msg.K = [
            intr.fx, 0.0, intr.ppx,
            0.0, intr.fy, intr.ppy,
            0.0, 0.0, 1.0
        ]

        # Rectification matrix (R)
        info_msg.R = [
            1.0, 0.0, 0.0,
            0.0, 1.0, 0.0,
            0.0, 0.0, 1.0
        ]

        # Projection matrix (P)
        info_msg.P = [
            intr.fx, 0.0, intr.ppx, 0.0,
            0.0, intr.fy, intr.ppy, 0.0,
            0.0, 0.0, 1.0, 0.0
        ]

        depth_pub.publish(img_msg)
        color_pub.publish(color_msg)
        info_pub.publish(info_msg)

        rate.sleep()

    pipeline.stop()


if __name__ == "__main__":
    main()