import os
import sys
import time
from tqdm import tqdm
import numpy as np
import cv2

sys.path.append('/path/dgp')
from dgp.datasets.synchronized_dataset import SynchronizedSceneDataset

import rclpy
from rclpy.time import Time
from rclpy.node import Node
from sensor_msgs.msg import PointCloud2, PointField, Image
from std_msgs.msg import Header

from cv_bridge import CvBridge

image_01 = None
depth_01 = None

class Publisher(Node):
    def __init__(self):
        super().__init__('ddad_publisher')
        self.lidar_publisher = self.create_publisher(PointCloud2, "luminar")
        self.image_publisher = self.create_publisher(Image, "image01")
        self.depth_publisher = self.create_publisher(Image, "depth01")
        self.image_bridge = CvBridge()
        self.depth_bridge = CvBridge()

    def array_to_pointcloud2(self, points, stamp=None, parent_frame="luminar"):
        ros_dtype = PointField.FLOAT32
        dtype = np.float32
        itemsize = np.dtype(dtype).itemsize
        data = points.astype(dtype).tobytes()

        fields = [PointField(
            name=n, offset=i*itemsize, datatype=ros_dtype, count=1)
            for i, n in enumerate('xyz')]

        header = Header(frame_id=parent_frame)

        return PointCloud2(
            header=header,
            height=1,
            width=points.shape[0],
            is_dense=False,
            is_bigendian=False,
            fields=fields,
            point_step=(itemsize * 3),
            row_step=(itemsize * 3 * points.shape[0]),
            data=data
        )

    def pil2cv(self, image):
        ''' PIL型 -> OpenCV型 '''
        new_image = np.array(image, dtype=np.uint8)
        if new_image.ndim == 2:  # モノクロ
            pass
        elif new_image.shape[2] == 3:  # カラー
            new_image = cv2.cvtColor(new_image, cv2.COLOR_RGB2BGR)
        elif new_image.shape[2] == 4:  # 透過
            new_image = cv2.cvtColor(new_image, cv2.COLOR_RGBA2BGRA)
        return new_image

    def dataloader(self):
        st = time.time()
        dataset = SynchronizedSceneDataset(
            'path/ddad.json',
            datum_names=('lidar', 'CAMERA_01', 'CAMERA_05'),
            generate_depth_from_datum='lidar',
            split='train'
            )
        print('Loading dataset took {:.2f} s'.format(time.time() - st))
        # Iterate through the dataset.
        for sample in tqdm(dataset, desc='Loading samples from the dataset'):
            lidar, camera_01, camera_05 =sample[0:3]
            point_cloud = lidar['point_cloud']
            image_01 = camera_01['rgb']
            depth_01 = camera_01['depth']

            print(depth_01.shape)

            pc2 = self.array_to_pointcloud2(point_cloud)
            self.lidar_publisher.publish(pc2)

            image_cv_01 = self.pil2cv(image_01)
            depth_cv_01 = self.pil2cv(depth_01)

            image_msg_01 = self.image_bridge.cv2_to_imgmsg(image_cv_01, encoding="bgr8")
            depth_msg_01 = self.depth_bridge.cv2_to_imgmsg(depth_cv_01, encoding="8UC1")

            self.image_publisher.publish(image_msg_01)
            self.depth_publisher.publish(depth_msg_01)

def main(args=None):
    rclpy.init(args=args)
    ddad_publisher = Publisher()
    ddad_publisher.dataloader()

if __name__ == '__main__':
    main()