# Source setup of ROS2 before running this script!

import rclpy
from rclpy.node import Node


class LoggingSubscriber(Node):
    def __init__(self, node_name, topic, msg_type):
        super().__init__(node_name)
        self.subscription = self.create_subscription(
            msg_type,
            topic,
            self.listener_callback,
            10)
        self.subscription  # prevent unused variable warning

    def listener_callback(self, msg):
        self.get_logger().info('I heard: "%s"' % msg.data)


import numpy as np
import cv2
class ImageSubscriber(Node):
    def __init__(self, node_name, topic, msg_type):
        super().__init__(node_name)
        self.subscription = self.create_subscription(
            msg_type,
            topic,
            self.listener_callback,
            10)
        self.subscription  # prevent unused variable warning

    def listener_callback(self, msg):
        dt = np.dtype('float32')
        dt = dt.newbyteorder('<')
        current_frame = np.ndarray(shape=(msg.height, msg.width, 1), dtype = dt, buffer = bytearray(msg.data))

        clipmin = 0.3
        clipmax = 2.0
        skewmin = np.amin(current_frame)
        skewmax = np.amax(current_frame)

        new_frame = current_frame.copy()

        np.clip(new_frame, clipmin, clipmax, out=new_frame)
        new_frame -= skewmin
        new_frame /= (skewmax-skewmin)/255.0
        new_frame = new_frame.astype(dtype=np.uint8, copy=False)

        im_color = cv2.applyColorMap(new_frame, cv2.COLORMAP_WINTER)
      
        cv2.imshow("camera", im_color)
        cv2.waitKey(1)


def main(args=None):
    from std_msgs.msg import Float32
    from sensor_msgs.msg import Image
    
    rclpy.init(args=args)

    sub = ImageSubscriber("logger", "rt/depth_processor", Image)

    rclpy.spin(sub)

    sub.destroy_node()
    rclpy.shutdown()


if __name__ == '__main__':
    main()
