import argparse
import glob
from pathlib import Path

import rospy
from sensor_msgs.msg import PointCloud2, PointCloud
from sensor_msgs import point_cloud2
from visualization_msgs.msg import MarkerArray
from visualization_msgs.msg import Marker

from scipy.spatial.transform import Rotation as R

import time

try:
    import open3d
    from visual_utils import open3d_vis_utils as V
    OPEN3D_FLAG = True
except:
    import mayavi.mlab as mlab
    from visual_utils import visualize_utils as V
    OPEN3D_FLAG = False

import numpy as np
import torch

from pcdet.config import cfg, cfg_from_yaml_file
from pcdet.datasets import DatasetTemplate
from pcdet.models import build_network, load_data_to_gpu
from pcdet.utils import common_utils

color_list = {
    'Car': (0, 255, 0),
    'Pedestrian': (255, 0, 255),
    'Truck': (255, 0, 0),
    'Bicycle': (255, 165, 0),
    'Motorcycle': (255, 255, 0), 
    'Tricycle': (0, 0, 255), 
    'Other_vehicle': (0, 255, 255)
}


class DemoDataset(DatasetTemplate):
    def __init__(self, dataset_cfg, class_names, training=True, root_path=None, logger=None):
        """
        Args:
            root_path:
            dataset_cfg:
            class_names:
            training:
            logger:
        """
        super().__init__(
            dataset_cfg=dataset_cfg, class_names=class_names, training=training, root_path=root_path, logger=logger
        )
        self.root_path = root_path

    def __getitem__(self, index):
        pass

    def __len__(self):
        pass

    def prepare_data(self, points, index):

        input_dict = {
            'points': points,
            'frame_id': index,
        }

        return super().prepare_data(data_dict=input_dict)

def parse_config():
    parser = argparse.ArgumentParser(description='arg parser')
    parser.add_argument('--cfg_file', type=str, default='cfgs/kitti_models/second.yaml',
                        help='specify the config for demo')
    parser.add_argument('--data_path', type=str, default='demo_data',
                        help='specify the point cloud data file or directory')
    parser.add_argument('--ckpt', type=str, default=None, help='specify the pretrained model')
    parser.add_argument('--ext', type=str, default='.bin', help='specify the extension of your point cloud data file')

    args = parser.parse_args()

    cfg_from_yaml_file(args.cfg_file, cfg)

    return args, cfg

def get_cur_time_ms():
    time_stamp = time.time()
    return int(round(time_stamp * 1000))

args, cfg = parse_config()
logger = common_utils.create_logger()
logger.info('-----------------Quick Demo of OpenPCDet-------------------------')
demo_dataset = DemoDataset(
    dataset_cfg=cfg.DATA_CONFIG, class_names=cfg.CLASS_NAMES, training=False,
    root_path=Path(args.data_path), logger=logger)

model = build_network(model_cfg=cfg.MODEL, num_class=len(cfg.CLASS_NAMES), dataset=demo_dataset)
model.load_params_from_file(filename=args.ckpt, logger=logger, to_cpu=True)
model.cuda()
model.eval()

frame_idx = 0
object_publisher = None


def callback_pc(msg):
    assert isinstance(msg, PointCloud2)

    t1 = get_cur_time_ms()
    points_list = point_cloud2.read_points_list(msg, field_names=('x', 'y', 'z', 'intensity'), skip_nans=True)  # 700+ms

    points = np.zeros([len(points_list), 4], dtype=np.float32)
    for idx, point in enumerate(points_list):
        points[idx] = [point.x, point.y, point.z, point.intensity]  # 100+ms

    # normalize intensity
    i_min = min(points[:,3])
    i_max = max(points[:,3])
    points[:,3] = (points[:,3] - i_min) / (i_max - i_min + 1e-3)    # 10+ms

    t2 = get_cur_time_ms()
    print('get point time: %dms.' % (t2-t1))    # 800+ms

    with torch.no_grad():
        global frame_idx
        data_dict = demo_dataset.prepare_data(points, frame_idx)
        t3 = get_cur_time_ms()
        print('prepare_data time: %dms.' % (t3-t2)) # 9ms

        data_dict = demo_dataset.collate_batch([data_dict])
        t4 = get_cur_time_ms()
        print('collate_batch: %dms.' % (t4-t3))     # 1ms

        load_data_to_gpu(data_dict)
        t5 = get_cur_time_ms()
        print('load_data_to_gpu: %dms.' % (t5-t4))  # 1ms

        pred_dicts, _ = model.forward(data_dict)
        t6 = get_cur_time_ms()
        print('inference+post_processing time: %dms.' % (t6-t5))    # 35ms
        print('total time: %dms.' % (t6-t1))
        print('*'*20)

        pred_boxes  = pred_dicts[0]['pred_boxes'].cpu().numpy()
        pred_scores = pred_dicts[0]['pred_scores'].cpu().numpy()
        pred_labels = pred_dicts[0]['pred_labels'].cpu().numpy()

        # detection threshold
        obj_mark    = pred_scores > 0.3
        pred_boxes  = pred_boxes[obj_mark]
        pred_labels = pred_labels[obj_mark]
        pred_scores = pred_scores[obj_mark]
        obj_num     = len(pred_boxes)

        print(pred_boxes)
        print(pred_scores)
        print(pred_labels)

        global object_publisher

        empty_markers           = MarkerArray()
        clear_obj_marker        = Marker()
        clear_obj_marker.header = msg.header
        clear_obj_marker.ns     = 'objects'
        clear_obj_marker.id     = 0
        clear_obj_marker.action = Marker.DELETEALL
        empty_markers.markers.append(clear_obj_marker)
        clear_txt_marker        = Marker()
        clear_txt_marker.header = msg.header
        clear_txt_marker.ns     = 'txt'
        clear_txt_marker.id     = 1
        clear_txt_marker.action = Marker.DELETEALL
        empty_markers.markers.append(clear_txt_marker)
        object_publisher.publish(empty_markers)

        object_markers = MarkerArray()
        object_markers.markers = []
        for idx, box, score, label in zip(range(obj_num), pred_boxes, pred_scores, pred_labels):

            obj_color = color_list[demo_dataset.class_names[label-1]]

            box_marker                 = Marker()
            box_marker.header          = msg.header
            box_marker.header.frame_id = msg.header.frame_id
            box_marker.ns              = 'objects'
            box_marker.id              = idx + 2
            box_marker.type            = Marker.CUBE
            box_marker.action          = Marker.ADD
            box_marker.color.r         = obj_color[0]
            box_marker.color.g         = obj_color[1]
            box_marker.color.b         = obj_color[2]
            box_marker.color.a         = 0.5
            box_marker.scale.x         = box[3]
            box_marker.scale.y         = box[4]
            box_marker.scale.z         = box[5]
            box_marker.pose.position.x = box[0]
            box_marker.pose.position.y = box[1]
            box_marker.pose.position.z = box[2]

            r   = R.from_euler('zxy', [box[6], 0., 0.])
            qua = r.as_quat()
            box_marker.pose.orientation.x = qua[0]
            box_marker.pose.orientation.y = qua[1]
            box_marker.pose.orientation.z = qua[2]
            box_marker.pose.orientation.w = qua[3]

            object_markers.markers.append(box_marker)

            txt_marker                    = Marker()
            txt_marker.header             = msg.header
            txt_marker.header.frame_id    = msg.header.frame_id
            txt_marker.ns                 = 'txt'
            txt_marker.id                 = idx + obj_num + 2
            txt_marker.type               = Marker.TEXT_VIEW_FACING
            txt_marker.action             = Marker.ADD
            txt_marker.color.r            = 255
            txt_marker.color.g            = 255
            txt_marker.color.b            = 0
            txt_marker.color.a            = 1.0
            txt_marker.scale.z            = 1.0
            txt_marker.pose.orientation.w = 1.0
            txt_marker.pose.position.x    = box[0]
            txt_marker.pose.position.y    = box[1]
            txt_marker.pose.position.z    = box[2] + box[5]/2
            txt_marker.text               = '%s: %.2f' % (demo_dataset.class_names[label-1], score)

            object_markers.markers.append(txt_marker)


        object_publisher.publish(object_markers)
        pc_publisher.publish(msg)

    frame_idx = frame_idx + 1

if __name__ == '__main__':

    rospy.init_node('demo', anonymous=True)
    rospy.Subscriber('/golfcart/all_combined_pc', PointCloud2, callback_pc)
    object_publisher = rospy.Publisher('/golfcart/objects', MarkerArray, queue_size=1)
    pc_publisher = rospy.Publisher('/golfcart/input_pc', PointCloud2, queue_size=1)
    rospy.spin()

import argparse
import glob
from pathlib import Path

import rospy
from sensor_msgs.msg import PointCloud2, PointCloud
from sensor_msgs import point_cloud2
from visualization_msgs.msg import MarkerArray
from visualization_msgs.msg import Marker

from scipy.spatial.transform import Rotation as R

import time

try:
    import open3d
    from visual_utils import open3d_vis_utils as V
    OPEN3D_FLAG = True
except:
    import mayavi.mlab as mlab
    from visual_utils import visualize_utils as V
    OPEN3D_FLAG = False

import numpy as np
import torch

from pcdet.config import cfg, cfg_from_yaml_file
from pcdet.datasets import DatasetTemplate
from pcdet.models import build_network, load_data_to_gpu
from pcdet.utils import common_utils

color_list = {
    'Car': (0, 255, 0),
    'Pedestrian': (255, 0, 255),
    'Truck': (255, 0, 0),
    'Bicycle': (255, 165, 0),
    'Motorcycle': (255, 255, 0), 
    'Tricycle': (0, 0, 255), 
    'Other_vehicle': (0, 255, 255)
}


class DemoDataset(DatasetTemplate):
    def __init__(self, dataset_cfg, class_names, training=True, root_path=None, logger=None):
        """
        Args:
            root_path:
            dataset_cfg:
            class_names:
            training:
            logger:
        """
        super().__init__(
            dataset_cfg=dataset_cfg, class_names=class_names, training=training, root_path=root_path, logger=logger
        )
        self.root_path = root_path

    def __getitem__(self, index):
        pass

    def __len__(self):
        pass

    def prepare_data(self, points, index):

        input_dict = {
            'points': points,
            'frame_id': index,
        }

        return super().prepare_data(data_dict=input_dict)

def parse_config():
    parser = argparse.ArgumentParser(description='arg parser')
    parser.add_argument('--cfg_file', type=str, default='cfgs/kitti_models/second.yaml',
                        help='specify the config for demo')
    parser.add_argument('--data_path', type=str, default='demo_data',
                        help='specify the point cloud data file or directory')
    parser.add_argument('--ckpt', type=str, default=None, help='specify the pretrained model')
    parser.add_argument('--ext', type=str, default='.bin', help='specify the extension of your point cloud data file')

    args = parser.parse_args()

    cfg_from_yaml_file(args.cfg_file, cfg)

    return args, cfg

def get_cur_time_ms():
    time_stamp = time.time()
    return int(round(time_stamp * 1000))

args, cfg = parse_config()
logger = common_utils.create_logger()
logger.info('-----------------Quick Demo of OpenPCDet-------------------------')
demo_dataset = DemoDataset(
    dataset_cfg=cfg.DATA_CONFIG, class_names=cfg.CLASS_NAMES, training=False,
    root_path=Path(args.data_path), logger=logger)

model = build_network(model_cfg=cfg.MODEL, num_class=len(cfg.CLASS_NAMES), dataset=demo_dataset)
model.load_params_from_file(filename=args.ckpt, logger=logger, to_cpu=True)
model.cuda()
model.eval()

frame_idx = 0
object_publisher = None


def callback_pc(msg):
    assert isinstance(msg, PointCloud2)

    t1 = get_cur_time_ms()
    points_list = point_cloud2.read_points_list(msg, field_names=('x', 'y', 'z', 'intensity'), skip_nans=True)  # 700+ms

    points = np.zeros([len(points_list), 4], dtype=np.float32)
    for idx, point in enumerate(points_list):
        points[idx] = [point.x, point.y, point.z, point.intensity]  # 100+ms

    # normalize intensity
    i_min = min(points[:,3])
    i_max = max(points[:,3])
    points[:,3] = (points[:,3] - i_min) / (i_max - i_min + 1e-3)    # 10+ms

    t2 = get_cur_time_ms()
    print('get point time: %dms.' % (t2-t1))    # 800+ms

    with torch.no_grad():
        global frame_idx
        data_dict = demo_dataset.prepare_data(points, frame_idx)
        t3 = get_cur_time_ms()
        print('prepare_data time: %dms.' % (t3-t2)) # 9ms

        data_dict = demo_dataset.collate_batch([data_dict])
        t4 = get_cur_time_ms()
        print('collate_batch: %dms.' % (t4-t3))     # 1ms

        load_data_to_gpu(data_dict)
        t5 = get_cur_time_ms()
        print('load_data_to_gpu: %dms.' % (t5-t4))  # 1ms

        pred_dicts, _ = model.forward(data_dict)
        t6 = get_cur_time_ms()
        print('inference+post_processing time: %dms.' % (t6-t5))    # 35ms
        print('total time: %dms.' % (t6-t1))
        print('*'*20)

        pred_boxes  = pred_dicts[0]['pred_boxes'].cpu().numpy()
        pred_scores = pred_dicts[0]['pred_scores'].cpu().numpy()
        pred_labels = pred_dicts[0]['pred_labels'].cpu().numpy()

        # detection threshold
        obj_mark    = pred_scores > 0.3
        pred_boxes  = pred_boxes[obj_mark]
        pred_labels = pred_labels[obj_mark]
        pred_scores = pred_scores[obj_mark]
        obj_num     = len(pred_boxes)

        print(pred_boxes)
        print(pred_scores)
        print(pred_labels)

        global object_publisher

        empty_markers           = MarkerArray()
        clear_obj_marker        = Marker()
        clear_obj_marker.header = msg.header
        clear_obj_marker.ns     = 'objects'
        clear_obj_marker.id     = 0
        clear_obj_marker.action = Marker.DELETEALL
        empty_markers.markers.append(clear_obj_marker)
        clear_txt_marker        = Marker()
        clear_txt_marker.header = msg.header
        clear_txt_marker.ns     = 'txt'
        clear_txt_marker.id     = 1
        clear_txt_marker.action = Marker.DELETEALL
        empty_markers.markers.append(clear_txt_marker)
        object_publisher.publish(empty_markers)

        object_markers = MarkerArray()
        object_markers.markers = []
        for idx, box, score, label in zip(range(obj_num), pred_boxes, pred_scores, pred_labels):

            obj_color = color_list[demo_dataset.class_names[label-1]]

            box_marker                 = Marker()
            box_marker.header          = msg.header
            box_marker.header.frame_id = msg.header.frame_id
            box_marker.ns              = 'objects'
            box_marker.id              = idx + 2
            box_marker.type            = Marker.CUBE
            box_marker.action          = Marker.ADD
            box_marker.color.r         = obj_color[0]
            box_marker.color.g         = obj_color[1]
            box_marker.color.b         = obj_color[2]
            box_marker.color.a         = 0.5
            box_marker.scale.x         = box[3]
            box_marker.scale.y         = box[4]
            box_marker.scale.z         = box[5]
            box_marker.pose.position.x = box[0]
            box_marker.pose.position.y = box[1]
            box_marker.pose.position.z = box[2]

            r   = R.from_euler('zxy', [box[6], 0., 0.])
            qua = r.as_quat()
            box_marker.pose.orientation.x = qua[0]
            box_marker.pose.orientation.y = qua[1]
            box_marker.pose.orientation.z = qua[2]
            box_marker.pose.orientation.w = qua[3]

            object_markers.markers.append(box_marker)

            txt_marker                    = Marker()
            txt_marker.header             = msg.header
            txt_marker.header.frame_id    = msg.header.frame_id
            txt_marker.ns                 = 'txt'
            txt_marker.id                 = idx + obj_num + 2
            txt_marker.type               = Marker.TEXT_VIEW_FACING
            txt_marker.action             = Marker.ADD
            txt_marker.color.r            = 255
            txt_marker.color.g            = 255
            txt_marker.color.b            = 0
            txt_marker.color.a            = 1.0
            txt_marker.scale.z            = 1.0
            txt_marker.pose.orientation.w = 1.0
            txt_marker.pose.position.x    = box[0]
            txt_marker.pose.position.y    = box[1]
            txt_marker.pose.position.z    = box[2] + box[5]/2
            txt_marker.text               = '%s: %.2f' % (demo_dataset.class_names[label-1], score)

            object_markers.markers.append(txt_marker)


        object_publisher.publish(object_markers)
        pc_publisher.publish(msg)

    frame_idx = frame_idx + 1

if __name__ == '__main__':

    rospy.init_node('demo', anonymous=True)
    rospy.Subscriber('/golfcart/all_combined_pc', PointCloud2, callback_pc)
    object_publisher = rospy.Publisher('/golfcart/objects', MarkerArray, queue_size=1)
    pc_publisher = rospy.Publisher('/golfcart/input_pc', PointCloud2, queue_size=1)
    rospy.spin()

Args:             root_path:             dataset_cfg:             class_names:             training:             logger: # 700+ms # 100+ms # normalize intensity # 10+ms # 800+ms # 9ms # 1ms # 1ms # 35ms # detection threshold