import os
import math
import copy
import argparse
from os import path as osp
from collections import OrderedDict
from typing import List, Tuple, Union
import torch
import cv2
from scipy.optimize import linear_sum_assignment
import numpy as np
from pyquaternion import Quaternion
from shapely.geometry import MultiPoint, box, LineString
from numpy.typing import NDArray
from shapely import affinity, ops
from typing import List, Tuple, Union, Dict
from projects.mmdet3d_plugin.datasets.visualize.renderer import Renderer
from PIL import Image, ImageDraw

import imageio

import mmcv

from nuscenes.nuscenes import NuScenes
from nuscenes.can_bus.can_bus_api import NuScenesCanBus
from nuscenes.utils.geometry_utils import transform_matrix
from nuscenes.utils.data_classes import Box
from nuscenes.utils.geometry_utils import view_points
from nuscenes.prediction import PredictHelper, convert_local_coords_to_global

from projects.mmdet3d_plugin.datasets.map_utils.nuscmap_extractor import NuscMapExtractor

from nusc_split import TRAIN_SCENES, VAL_SCENES

NameMapping = {
    "movable_object.barrier": "barrier",
    "vehicle.bicycle": "bicycle",
    "vehicle.bus.bendy": "bus",
    "vehicle.bus.rigid": "bus",
    "vehicle.car": "car",
    "vehicle.construction": "construction_vehicle",
    "vehicle.motorcycle": "motorcycle",
    "human.pedestrian.adult": "pedestrian",
    "human.pedestrian.child": "pedestrian",
    "human.pedestrian.construction_worker": "pedestrian",
    "human.pedestrian.police_officer": "pedestrian",
    "movable_object.trafficcone": "traffic_cone",
    "vehicle.trailer": "trailer",
    "vehicle.truck": "truck",
}

def quart_to_rpy(qua):
    x, y, z, w = qua
    roll = math.atan2(2 * (w * x + y * z), 1 - 2 * (x * x + y * y))
    pitch = math.asin(2 * (w * y - x * z))
    yaw = math.atan2(2 * (w * z + x * y), 1 - 2 * (z * z + y * y))
    return roll, pitch, yaw

def locate_message(utimes, utime):
    i = np.searchsorted(utimes, utime)
    if i == len(utimes) or (i > 0 and utime - utimes[i-1] < utimes[i] - utime):
        i -= 1
    return i

def geom2anno(map_geoms):
    MAP_CLASSES = (
        'ped_crossing',
        'divider',
        'boundary',
    )
    vectors = {}
    for cls, geom_list in map_geoms.items():
        if cls in MAP_CLASSES:
            label = MAP_CLASSES.index(cls)
            vectors[label] = []
            for geom in geom_list:
                line = np.array(geom.coords)
                vectors[label].append(line)
    return vectors

def create_nuscenes_infos(root_path,
                          out_path,
                          can_bus_root_path,
                          info_prefix,
                          version='v1.0-trainval',
                          max_sweeps=10,
                          roi_size=(30, 60),):
    """Create info file of nuscene dataset.

    Given the raw data, generate its related info file in pkl format.

    Args:
        root_path (str): Path of the data root.
        info_prefix (str): Prefix of the info file to be generated.
        version (str): Version of the data.
            Default: 'v1.0-trainval'
        max_sweeps (int): Max number of sweeps.
            Default: 10
    """
    print(version, root_path)
    nusc = NuScenes(version=version, dataroot=root_path, verbose=True)
    nusc_map_extractor = NuscMapExtractor(root_path, roi_size)
    nusc_can_bus = NuScenesCanBus(dataroot=can_bus_root_path)
    from nuscenes.utils import splits
    available_vers = ['v1.0-trainval', 'v1.0-test', 'v1.0-mini']
    assert version in available_vers
    if version == 'v1.0-trainval':
        train_scenes = splits.train
        val_scenes = splits.val
    elif version == 'v1.0-test':
        train_scenes = splits.test
        val_scenes = []
    elif version == 'v1.0-mini':
        train_scenes = splits.mini_train
        val_scenes = splits.mini_val
        out_path = osp.join(out_path, 'mini')
    else:
        raise ValueError('unknown')

    if "infos_newsplit" in out_path:
        import pdb; pdb.set_trace()
        train_scenes = TRAIN_SCENES
        val_scenes = VAL_SCENES
        
        print(" val_scenes========", val_scenes)

    os.makedirs(out_path, exist_ok=True)

    # filter existing scenes.
    available_scenes = get_available_scenes(nusc)
    available_scene_names = [s['name'] for s in available_scenes]
    train_scenes = list(
        filter(lambda x: x in available_scene_names, train_scenes))
    val_scenes = list(filter(lambda x: x in available_scene_names, val_scenes))
    print("++++++++++++++++++++++++++")
    print("train_scenes =", train_scenes)
    print("val_scenes ==", val_scenes)
    train_scenes = set([
        available_scenes[available_scene_names.index(s)]['token']
        for s in train_scenes
    ])
    val_scenes = set([
        available_scenes[available_scene_names.index(s)]['token']
        for s in val_scenes
    ])

    test = 'test' in version
    if test:
        print('test scene: {}'.format(len(train_scenes)))
    else:
        print('train scene: {}, val scene: {}'.format(
            len(train_scenes), len(val_scenes)))

    train_nusc_infos, val_nusc_infos = _fill_trainval_infos(
        nusc, nusc_map_extractor, nusc_can_bus, train_scenes, val_scenes, test, max_sweeps=max_sweeps)

    metadata = dict(version=version)
    if test:
        print('test sample: {}'.format(len(train_nusc_infos)))
        data = dict(infos=train_nusc_infos, metadata=metadata)
        info_path = osp.join(out_path,
                             '{}_infos_test.pkl'.format(info_prefix))
        mmcv.dump(data, info_path)
    else:
        print('train sample: {}, val sample: {}'.format(
            len(train_nusc_infos), len(val_nusc_infos)))
        data = dict(infos=train_nusc_infos, metadata=metadata)
        info_path = osp.join(out_path,
                             '{}_infos_train.pkl'.format(info_prefix))
        mmcv.dump(data, info_path)
        data['infos'] = val_nusc_infos
        # import pdb; pdb.set_trace()
        info_val_path = osp.join(out_path,
                                 '{}_infos_val.pkl'.format(info_prefix))
        mmcv.dump(data, info_val_path)

def get_available_scenes(nusc):
    """Get available scenes from the input nuscenes class.

    Given the raw data, get the information of available scenes for
    further info generation.

    Args:
        nusc (class): Dataset class in the nuScenes dataset.

    Returns:
        available_scenes (list[dict]): List of basic information for the
            available scenes.
    """
    available_scenes = []
    print('total scene num: {}'.format(len(nusc.scene)))
    for scene in nusc.scene:
        scene_token = scene['token']
        scene_rec = nusc.get('scene', scene_token)
        sample_rec = nusc.get('sample', scene_rec['first_sample_token'])
        sd_rec = nusc.get('sample_data', sample_rec['data']['LIDAR_TOP'])
        has_more_frames = True
        scene_not_exist = False
        while has_more_frames:
            lidar_path, boxes, _ = nusc.get_sample_data(sd_rec['token'])
            lidar_path = str(lidar_path)
            if os.getcwd() in lidar_path:
                # path from lyftdataset is absolute path
                lidar_path = lidar_path.split(f'{os.getcwd()}/')[-1]
                # relative path
            if not mmcv.is_filepath(lidar_path):
                scene_not_exist = True
                break
            else:
                break
        if scene_not_exist:
            continue
        available_scenes.append(scene)
    print('exist scene num: {}'.format(len(available_scenes)))
    return available_scenes

def _fill_trainval_infos(nusc,
                         nusc_map_extractor,
                         nusc_can_bus,
                         train_scenes,
                         val_scenes,
                         test=False,
                         max_sweeps=10,
                         fut_ts=12,
                         ego_fut_ts=6):
    """Generate the train/val infos from the raw data.

    Args:
        nusc (:obj:`NuScenes`): Dataset class in the nuScenes dataset.
        train_scenes (list[str]): Basic information of training scenes.
        val_scenes (list[str]): Basic information of validation scenes.
        test (bool): Whether use the test mode. In the test mode, no
            annotations can be accessed. Default: False.
        max_sweeps (int): Max number of sweeps. Default: 10.

    Returns:
        tuple[list[dict]]: Information of training set and validation set
            that will be saved to the info file.
    """
    train_nusc_infos = []
    val_nusc_infos = []
    cat2idx = {}
    for idx, dic in enumerate(nusc.category):
        cat2idx[dic['name']] = idx

    predict_helper = PredictHelper(nusc)
    # classify 
    scene_name2idx = {}
    tmp_idx = 0
    scene_name2idx_val = {}
    tmp_idx_val = 0

    scene_vectors_seq = {}
    scene_vectors_seq_val = {}
    scene_matchings_seq = {}
    scene_matchings_seq_val = {}
    # local_idx = 0
    for sample_tmp in mmcv.track_iter_progress(nusc.sample):
        scene_record_ = nusc.get('scene', sample_tmp['scene_token'])
        scene_n = scene_record_['name']
        # scene_n = sample_tmp['scene_name']
        if sample_tmp['scene_token'] in train_scenes:

            if scene_n not in scene_name2idx:
                scene_name2idx[scene_n] = []
                scene_vectors_seq[scene_n] = []
                scene_matchings_seq[scene_n] = []
                # local_idx = 0
            scene_name2idx[scene_n].append(tmp_idx)
            tmp_idx += 1
            # local_idx += 1
        elif sample_tmp['scene_token'] in val_scenes:

            # if scene_n not in scene_name2idx_val:
            #     scene_name2idx_val[scene_n] = []
            #     scene_vectors_seq_val[scene_n] = []
            #     scene_matchings_seq_val[scene_n] = []
            # scene_name2idx_val[scene_n].append(tmp_idx_val)
            # tmp_idx_val += 1
            if scene_n not in scene_name2idx:
                scene_name2idx[scene_n] = []
                scene_vectors_seq_val[scene_n] = []
                scene_matchings_seq_val[scene_n] = []
                # local_idx = 0
            scene_name2idx[scene_n].append(tmp_idx_val)
            tmp_idx_val += 1
            # local_idx += 1

    train_sample_idx = 0
    val_sample_idx = 0
    train_prev = None
    val_prev = None
    local_idx = 0
    scene_local_idx = {}
    for sample in mmcv.track_iter_progress(nusc.sample):
        map_location = nusc.get('log', nusc.get('scene', sample['scene_token'])['log_token'])['location']
        lidar_token = sample['data']['LIDAR_TOP']
        sd_rec = nusc.get('sample_data', lidar_token)
        cs_record = nusc.get('calibrated_sensor',
                             sd_rec['calibrated_sensor_token'])
        pose_record = nusc.get('ego_pose', sd_rec['ego_pose_token'])
        lidar_path, boxes, _ = nusc.get_sample_data(lidar_token)
        mmcv.check_file_exist(lidar_path)
        scene_record = nusc.get('scene', sample['scene_token'])
        scene_name = scene_record['name']
        # if scene_name in ['scene-0061','scene-0499', 'scene-0502', 'scene-0515', 'scene-0517', 'scene-0655', 'scene-0757',  'scene-0796', 'scene-1077', 'scene-1094','scene-1100'] :
        #     continue
        # get local idx
        if scene_name not in scene_local_idx:
            scene_local_idx[scene_name] = []
            local_idx = 0
            val_prev = None
            train_prev = None

        info = {
            'lidar_path': lidar_path,
            'token': sample['token'],
            'sweeps': [],
            'cams': dict(),
            'scene_token': sample['scene_token'],
            'lidar2ego_translation': cs_record['translation'],
            'lidar2ego_rotation': cs_record['rotation'],
            'ego2global_translation': pose_record['translation'],
            'ego2global_rotation': pose_record['rotation'],
            'timestamp': sample['timestamp'],
            'map_location': map_location,
            'scene_name':scene_name,
            'scene_idx_list': scene_name2idx[scene_name],
            'local_idx': local_idx,
        }
        local_idx += 1

        l2e_r = info['lidar2ego_rotation']
        l2e_t = info['lidar2ego_translation']
        e2g_r = info['ego2global_rotation']
        e2g_t = info['ego2global_translation']
        l2e_r_mat = Quaternion(l2e_r).rotation_matrix
        e2g_r_mat = Quaternion(e2g_r).rotation_matrix

        # extract map annos
        lidar2ego = np.eye(4)
        lidar2ego[:3, :3] = Quaternion(
            info["lidar2ego_rotation"]
        ).rotation_matrix
        lidar2ego[:3, 3] = np.array(info["lidar2ego_translation"])
        ego2global = np.eye(4)
        ego2global[:3, :3] = Quaternion(
            info["ego2global_rotation"]
        ).rotation_matrix
        ego2global[:3, 3] = np.array(info["ego2global_translation"])
        lidar2global = ego2global @ lidar2ego

        translation = list(lidar2global[:3, 3])
        rotation = list(Quaternion(matrix=lidar2global).q)
        map_geoms = nusc_map_extractor.get_map_geom(map_location, translation, rotation)
        map_annos = geom2anno(map_geoms)
        info['map_annos'] = map_annos
        info['map_geoms'] = map_geoms

        # obtain 6 image's information per frame
        camera_types = [
            'CAM_FRONT',
            'CAM_FRONT_RIGHT',
            'CAM_FRONT_LEFT',
            'CAM_BACK',
            'CAM_BACK_LEFT',
            'CAM_BACK_RIGHT',
        ]
        for cam in camera_types:
            cam_token = sample['data'][cam]
            cam_path, _, cam_intrinsic = nusc.get_sample_data(cam_token)
            cam_info = obtain_sensor2top(nusc, cam_token, l2e_t, l2e_r_mat,
                                         e2g_t, e2g_r_mat, cam)
            cam_info.update(cam_intrinsic=cam_intrinsic)
            info['cams'].update({cam: cam_info})

        # obtain sweeps for a single key-frame
        sd_rec = nusc.get('sample_data', sample['data']['LIDAR_TOP'])
        sweeps = []
        while len(sweeps) < max_sweeps:
            if not sd_rec['prev'] == '':
                sweep = obtain_sensor2top(nusc, sd_rec['prev'], l2e_t,
                                          l2e_r_mat, e2g_t, e2g_r_mat, 'lidar')
                sweeps.append(sweep)
                sd_rec = nusc.get('sample_data', sd_rec['prev'])
            else:
                break
        info['sweeps'] = sweeps
        # obtain annotation
        if not test:
            # object detection annos: boxes (locs, dims, yaw, velocity), names and valid flags
            annotations = [
                nusc.get('sample_annotation', token)
                for token in sample['anns']
            ]
            locs = np.array([b.center for b in boxes]).reshape(-1, 3)
            dims = np.array([b.wlh for b in boxes]).reshape(-1, 3)
            rots = np.array([b.orientation.yaw_pitch_roll[0]
                             for b in boxes]).reshape(-1, 1)
            velocity = np.array(
                [nusc.box_velocity(token)[:2] for token in sample['anns']])
            # convert velo from global to lidar
            for i in range(len(boxes)):
                velo = np.array([*velocity[i], 0.0])
                velo = velo @ np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(
                    l2e_r_mat).T
                velocity[i] = velo[:2]
            names = [b.name for b in boxes]
            for i in range(len(names)):
                if names[i] in NameMapping:
                    names[i] = NameMapping[names[i]]
            names = np.array(names)
            valid_flag = np.array(
                [(anno['num_lidar_pts'] + anno['num_radar_pts']) > 0
                 for anno in annotations],
                dtype=bool).reshape(-1)  ## TODO update valid flag for tracking
            # we need to convert box size to
            # the format of our lidar coordinate system
            # which is x_size, y_size, z_size (corresponding to l, w, h)
            gt_boxes = np.concatenate([locs, dims[:, [1, 0, 2]], rots], axis=1)
            assert len(gt_boxes) == len(
                annotations), f'{len(gt_boxes)}, {len(annotations)}'
            
            # object tracking annos: instance_ids
            instance_inds = [nusc.getind('instance', anno['instance_token'])
                             for anno in annotations]

            # motion prediction annos: future trajectories offset in lidar frame and valid mask
            num_box = len(boxes)
            gt_fut_trajs = np.zeros((num_box, fut_ts, 2))
            gt_fut_masks = np.zeros((num_box, fut_ts))
            for i, anno in enumerate(annotations):
                instance_token = anno['instance_token']
                fut_traj_local = predict_helper.get_future_for_agent(
                    instance_token, 
                    sample['token'], 
                    seconds=fut_ts/2, 
                    in_agent_frame=True
                )
                if fut_traj_local.shape[0] > 0:
                    box = boxes[i]
                    trans = box.center
                    rot = Quaternion(matrix=box.rotation_matrix)
                    fut_traj_scene = convert_local_coords_to_global(fut_traj_local, trans, rot)
                    valid_step = fut_traj_scene.shape[0]
                    gt_fut_trajs[i, 0] = fut_traj_scene[0] - box.center[:2]
                    gt_fut_trajs[i, 1:valid_step] = fut_traj_scene[1:] - fut_traj_scene[:-1]
                    gt_fut_masks[i, :valid_step] = 1

            # motion planning annos: future trajectories offset in lidar frame and valid mask
            ego_fut_trajs = np.zeros((ego_fut_ts + 1, 3))
            ego_fut_masks = np.zeros((ego_fut_ts + 1))
            sample_cur = sample
            ego_status = get_ego_status(nusc, nusc_can_bus, sample_cur)
            for i in range(ego_fut_ts + 1):
                pose_mat = get_global_sensor_pose(sample_cur, nusc)
                ego_fut_trajs[i] = pose_mat[:3, 3]
                ego_fut_masks[i] = 1
                if sample_cur['next'] == '':
                    ego_fut_trajs[i+1:] = ego_fut_trajs[i]
                    break
                else:
                    sample_cur = nusc.get('sample', sample_cur['next'])
            # global to ego
            ego_fut_trajs = ego_fut_trajs - np.array(pose_record['translation'])
            rot_mat = Quaternion(pose_record['rotation']).inverse.rotation_matrix
            ego_fut_trajs = np.dot(rot_mat, ego_fut_trajs.T).T
            # ego to lidar
            ego_fut_trajs = ego_fut_trajs - np.array(cs_record['translation'])
            rot_mat = Quaternion(cs_record['rotation']).inverse.rotation_matrix
            ego_fut_trajs = np.dot(rot_mat, ego_fut_trajs.T).T
            # drive command according to final fut step offset
            if ego_fut_trajs[-1][0] >= 2:
                command = np.array([1, 0, 0])  # Turn Right
            elif ego_fut_trajs[-1][0] <= -2:
                command = np.array([0, 1, 0])  # Turn Left
            else:
                command = np.array([0, 0, 1])  # Go Straight
            # get offset
            ego_fut_trajs = ego_fut_trajs[1:] - ego_fut_trajs[:-1]      

            info['gt_boxes'] = gt_boxes
            info['gt_names'] = names
            info['gt_velocity'] = velocity.reshape(-1, 2)
            info['num_lidar_pts'] = np.array(
                [a['num_lidar_pts'] for a in annotations])
            info['num_radar_pts'] = np.array(
                [a['num_radar_pts'] for a in annotations])
            info['valid_flag'] = valid_flag
            info['instance_inds'] = instance_inds
            info['gt_agent_fut_trajs'] = gt_fut_trajs.astype(np.float32)
            info['gt_agent_fut_masks'] = gt_fut_masks.astype(np.float32)
            info['gt_ego_fut_trajs'] = ego_fut_trajs[:, :2].astype(np.float32)
            info['gt_ego_fut_masks'] = ego_fut_masks[1:].astype(np.float32)
            info['gt_ego_fut_cmd'] = command.astype(np.float32)
            info['ego_status'] = ego_status
        # map track id
        # import pdb; pdb.set_trace()
        # if scene_name in train_scenes:
        if sample['scene_token'] in train_scenes:
            info.update({
                'sample_idx': train_sample_idx,
                'prev': train_sample_idx -1,
                'next': train_sample_idx+1,
            })
            if sample['prev'] == '':
                info['prev'] = -1
            if sample['next'] == '':
                info['next'] = -1
            # train_samples.append(info)
            # train_sample_idx += 1
        # elif scene_name in val_scenes:
        elif sample['scene_token'] in val_scenes:
            info.update({
                'sample_idx': val_sample_idx,
                'prev': val_sample_idx - 1,
                'next': val_sample_idx + 1,
            })
            if sample['prev'] == '':
                info['prev'] = -1
            if sample['next'] == '':
                info['next'] = -1
            # val_samples.append(info)
            # val_sample_idx += 1

        # import pdb; pdb.set_trace()
        roi_size = torch.tensor([60, 30])
        origin = torch.tensor([-30, -15])
        cfg = [200, 100]
        if sample['scene_token'] in train_scenes:
            if sample['scene_token'] == '325cef682f064c55a255f2625c533b75' or info['scene_name'] == 'scene-0103':
                print(scene_name)
                # import pdb; pdb.set_trace()

            # import pdb; pdb.set_trace()
            train_nusc_infos.append(info)
            # TODO Match prev and curr
            # print("==================",info['local_idx'])
            if train_prev ==None:
                print('------train------init-----',info['scene_name'], info['sample_idx'],"local idx ==", info['local_idx'])
                train_prev = info
            elif train_prev['scene_name'] == info['scene_name'] \
                and train_prev['sample_idx']+1 == info['sample_idx'] \
                and info['prev'] == train_prev['sample_idx'] \
                and info['sample_idx'] in info['scene_idx_list']:
                matchings, prev_vectors = match_two_consecutive_frames(train_prev, info, roi_size, origin, cfg)
                print("------------train--------------",info['scene_name'], info['sample_idx'],"local idx ==",info['local_idx'], info['timestamp'])
                # TODO ASSIGN_GLOBAL_IDS
                # import pdb; pdb.set_trace()
                scene_vectors_seq[scene_name].append(prev_vectors)
                scene_matchings_seq[scene_name].append(matchings)
                # print('== works fine ==')
                if info['next'] == -1: #info['sample_idx'] == 37:
                    # import pdb; pdb.set_trace()
                    print('------train------not in sequence-----', info['scene_name'], info['sample_idx'],"local idx ==",info['local_idx'] )
                    matchings, prev_vectors = match_two_consecutive_frames(info, info, roi_size, origin, cfg)
                    scene_vectors_seq[scene_name].append(prev_vectors)
                    # scene_matchings_seq[scene_name].append(matchings)
                    train_prev = info
            else:
                break
                print('------train-=========-----', info['scene_name'], info['sample_idx'],"local idx ==",info['local_idx'] )
                matchings, prev_vectors = match_two_consecutive_frames(train_prev, train_prev, roi_size, origin, cfg)
                scene_vectors_seq[scene_name].append(prev_vectors)
                # scene_matchings_seq[scene_name].append(matchings)
                train_prev = info
            train_prev = info
            train_sample_idx += 1
        else:
            # import pdb; pdb.set_trace()
            val_nusc_infos.append(info)
            # TODO Match prev and curr
            # TODO Match prev and curr
            if val_prev ==None:
                val_prev = info
            elif val_prev['scene_name'] == info['scene_name'] \
                and val_prev['sample_idx']+1 == info['sample_idx'] \
                and info['prev'] == val_prev['sample_idx'] \
                and info['sample_idx'] in info['scene_idx_list']:
                roi_size = torch.tensor([60, 30])
                origin = torch.tensor([-30, -15])
                cfg = [200, 100]
                # import pdb; pdb.set_trace()
                
                matchings,prev_vectors = match_two_consecutive_frames(val_prev, info, roi_size, origin, cfg)
                print("-------------val-------------",info['scene_name'], info['sample_idx'], info['timestamp'])
                # TODO ASSIGN_GLOBAL_IDS
                # matchings_seq.append
                scene_vectors_seq_val[scene_name].append(prev_vectors)
                scene_matchings_seq_val[scene_name].append(matchings)
                # print('== works fine ==')
                if info['next'] == -1: #info['sample_idx'] == 37:
                    # import pdb; pdb.set_trace()
                    print('------val------not in sequence-----', info['scene_name'], info['sample_idx'],"local idx ==",info['local_idx'] )
                    matchings, prev_vectors = match_two_consecutive_frames(info, info, roi_size, origin, cfg)
                    scene_vectors_seq_val[scene_name].append(prev_vectors)
                    # scene_matchings_seq_val[scene_name].append(matchings)
                    val_prev = info
            else:
                print('------val------not in sequence-----', val_prev['scene_name'], info['scene_name'])
                
                matchings, prev_vectors = match_two_consecutive_frames(val_prev, val_prev, roi_size, origin, cfg)
                # break
                scene_vectors_seq_val[scene_name].append(prev_vectors)
                # scene_matchings_seq_val[scene_name].append(matchings)
                val_prev = info
            val_prev = info
            val_sample_idx += 1
    # import pdb; pdb.set_trace()
    ids_info_dict = {}
    ids_info_dict_val = {}
    for (scenename, matchings_seq), (scenename_key, vectors_seq) in zip(scene_matchings_seq.items(), scene_vectors_seq.items()):
        # import pdb; pdb.set_trace()

        assert scenename == scenename_key
        print(" scene name ==", scenename)
        scene_id = int(scenename[6:])
        ids_info = assign_global_ids(matchings_seq, vectors_seq, scene_id)
        if scenename not in ids_info_dict:

            ids_info_dict[scenename] = ids_info
            vis(scenename, ids_info, vectors_seq, "track_vis")
        # ids_info_dict[scenename].append(ids_info)
        
        else:
            print("-----error 1-----")
            break
        
    for (scenename, matchings_seq_val), (scenename_key,vectors_seq_val) in zip(scene_matchings_seq_val.items(), scene_vectors_seq_val.items()):
        assert scenename == scenename_key
        scene_id = int(scenename[6:])
        # import pdb; pdb.set_trace()
        print(" scene name val --- ==", scenename)
        ids_info_val = assign_global_ids(matchings_seq_val, vectors_seq_val, scene_id)
        if scenename not in ids_info_dict_val:
            # import pdb; pdb.set_trace()
            ids_info_dict_val[scenename] = ids_info_val
            vis(scenename, ids_info_val, vectors_seq_val, "track_vis")
        else:
            print("-----error 2-----")
            break

    for train_nusc_infos_it in train_nusc_infos:
        # print(len(ids_info_dict[train_nusc_infos_it['scene_name']]),"  --  ", train_nusc_infos_it['local_idx'], )
        # print(train_nusc_infos_it['scene_name'], "==train=",ids_info_dict.keys())

        train_nusc_infos_it.update({
                'track_globle_id': ids_info_dict[train_nusc_infos_it['scene_name']][train_nusc_infos_it['local_idx']] 
            })

    for val_nusc_infos_it in val_nusc_infos:
        # print(len(ids_info_dict_val[val_nusc_infos_it['scene_name']]),"  =val=  ", val_nusc_infos_it['local_idx'])
        val_nusc_infos_it.update({
                'track_globle_id': ids_info_dict_val[val_nusc_infos_it['scene_name']][val_nusc_infos_it['local_idx']] 
            })
    # import pdb; pdb.set_trace()
    return train_nusc_infos, val_nusc_infos

# map tracking func
# vis
def vis(scene_name, ids_info, vectors_seq, out_dir):
    # vis
    font                   = cv2.FONT_HERSHEY_SIMPLEX
    location               = (200,60)
    fontScale              = 2
    fontColor              = (255,0,0)
    thickness              = 2
    lineType               = 2
    cat2id_ = {'ped_crossing':0,
        'divider':1,
        'boundary':2}
    renderer = Renderer(cat2id_, (60, 30), 'nusc')
    if False: #True:
        print('Visualize gt tracks for scene {}'.format(scene_name))
        scene_dir = os.path.join(out_dir, scene_name)
        os.makedirs(scene_dir, exist_ok=True)
        # visualize with matched track ids
        imgs = []
        for idx, (id_info, vectors) in enumerate(zip(ids_info, vectors_seq)):
            # vectors = _denorm(vectors, roi_size.numpy(), origin.numpy())
            # import pdb; pdb.set_trace()
            save_path = os.path.join(scene_dir, f'{idx}_with_id.png')
            renderer.render_bev_from_vectors(vectors, out_dir=None, specified_path=save_path, id_info=id_info)
            viz_img = np.ascontiguousarray(cv2.imread(save_path)[:, :, ::-1], dtype=np.uint8)
            if idx == 0:
                img_shape = (viz_img.shape[1], viz_img.shape[0])
            else:
                viz_img = cv2.resize(viz_img, img_shape)
            cv2.putText(viz_img, 't={}'.format(idx), location, font, fontScale, fontColor,
            thickness, lineType)
            imgs.append(viz_img)
        gif_path = os.path.join(scene_dir, 'matching.gif')
        imageio.mimsave(gif_path, imgs, duration=500)

def assign_global_ids(matchings_seq, vectors_seq, scene_id):
    ids_seq = []
    # global_map_index = {
    #     0: 0,
    #     1: 0,
    #     2: 0,
    # }
    global_map_index = {
        0: scene_id*1000,
        1: scene_id*1000+330,
        2: scene_id*1000+660,
    }
    
    ids_0 = dict()
    for label, vectors in vectors_seq[0].items():
        id_mapping = dict()
        for i, _ in enumerate(vectors):
            id_mapping[i] = global_map_index[label]
            global_map_index[label] += 1
        ids_0[label] = id_mapping
    ids_seq.append(ids_0)

    # Trace all frames following the consecutive matching
    for t, vectors_t in enumerate(vectors_seq[1:]):
        ids_t = dict()
        for label, vectors in vectors_t.items():
            reverse_matching = matchings_seq[t][label][1]
            id_mapping = dict()
            for i, _ in enumerate(vectors):
                if reverse_matching[i] != -1:
                    prev_id = reverse_matching[i]
                    global_id = ids_seq[-1][label][prev_id]
                else:
                    global_id = global_map_index[label]
                    global_map_index[label] += 1
                id_mapping[i] = global_id
            ids_t[label] = id_mapping
        ids_seq.append(ids_t)
    return ids_seq
    
def interp_fixed_num( line: LineString) -> NDArray:
    ''' Interpolate a line to fixed number of points.
    
    Args:
        line (LineString): line
    
    Returns:
        points (array): interpolated points, shape (N, 2)
    '''

    distances = np.linspace(0, line.length, 20)
    sampled_points = np.array([list(line.interpolate(distance).coords) 
        for distance in distances]).squeeze()

    return sampled_points
# def get_vectorized_lines( map_geoms: Dict) -> Dict:
#     ''' Vectorize map elements. Iterate over the input dict and apply the 
#     specified sample funcion.
    
#     Args:
#         line (LineString): line
    
#     Returns:
#         vectors (array): dict of vectorized map elements.
#     '''

#     vectors = {}
#     for label, geom_list in map_geoms.items():
#         vectors[label] = []
#         for geom in geom_list:
#             if geom.geom_type == 'LineString':

#                 line = interp_fixed_num(geom)
#                 coords_dim = 2
#                 line = line[:, :coords_dim]

#                 vectors[label].append(line)

#             elif geom.geom_type == 'Polygon':
#                 # polygon objects will not be vectorized
#                 continue
            
#             else:
#                 raise ValueError('map geoms must be either LineString or Polygon!')
#     return vectors
def get_vectorized_lines( map_geoms: Dict) -> Dict:
    ''' Vectorize map elements. Iterate over the input dict and apply the 
    specified sample funcion.
    
    Args:
        line (LineString): line
    
    Returns:
        vectors (array): dict of vectorized map elements.
    '''

    vectors = {}
    cat2id = {'ped_crossing':0,
        'divider':1,
        'boundary':2}
    for label, geom_list in map_geoms.items():
        if label in ['ped_crossing', 'divider', 'boundary']:
            vectors[cat2id[label]] = []
            # import pdb; pdb.set_trace()
            for geom in geom_list:
                if geom.geom_type == 'LineString':
                    geom =  affinity.rotate(geom, -90, origin=(0,0), use_radians=False)

                    line = interp_fixed_num(geom)
                    coords_dim = 2
                    line = line[:, :coords_dim]
                    # import pdb; pdb.set_trace()
                    # line = line[:, ::-1]
                    vectors[cat2id[label]].append(line)

                elif geom.geom_type == 'Polygon':
                    # polygon objects will not be vectorized
                    continue
                
                else:
                    raise ValueError('map geoms must be either LineString or Polygon!')
    return vectors
def match_two_consecutive_frames(prev_data, curr_data, roi_size=torch.tensor([60,30]), origin=torch.tensor([-30,-15]), cfg=[200, 100]):
    # get relative pose
    '''
    l2e_r = info['lidar2ego_rotation']
        l2e_t = info['lidar2ego_translation']
        e2g_r = info['ego2global_rotation']
        e2g_t = info['ego2global_translation']
        l2e_r_mat = Quaternion(l2e_r).rotation_matrix
        e2g_r_mat = Quaternion(e2g_r).rotation_matrix
    '''
    # import pdb; pdb.set_trace()
    # prev_e2g_trans = torch.tensor(prev_data['img_metas'].data['ego2global_translation'], dtype=torch.float64)
    # prev_e2g_rot = torch.tensor(prev_data['img_metas'].data['ego2global_rotation'], dtype=torch.float64)
    # curr_e2g_trans  = torch.tensor(curr_data['img_metas'].data['ego2global_translation'], dtype=torch.float64)
    # curr_e2g_rot = torch.tensor(curr_data['img_metas'].data['ego2global_rotation'], dtype=torch.float64)
    # print("====", prev_data['ego2global_translation'])
    prev_e2g_trans = torch.tensor(prev_data['ego2global_translation'], dtype=torch.float64)
    # prev_e2g_rot = torch.tensor(prev_data['ego2global_rotation'], dtype=torch.float64)
    prev_e2g_rot = torch.tensor(Quaternion(prev_data['ego2global_rotation']).rotation_matrix, dtype=torch.float64)
    curr_e2g_trans  = torch.tensor(curr_data['ego2global_translation'], dtype=torch.float64)
    # curr_e2g_rot = torch.tensor(curr_data['ego2global_rotation'], dtype=torch.float64) 
    curr_e2g_rot = torch.tensor(Quaternion(curr_data['ego2global_rotation']).rotation_matrix, dtype=torch.float64) 

    prev_e2g_matrix = torch.eye(4, dtype=torch.float64)
    prev_e2g_matrix[:3, :3] = prev_e2g_rot
    prev_e2g_matrix[:3, 3] = prev_e2g_trans

    curr_g2e_matrix = torch.eye(4, dtype=torch.float64)
    curr_g2e_matrix[:3, :3] = curr_e2g_rot.T
    curr_g2e_matrix[:3, 3] = -(curr_e2g_rot.T @ curr_e2g_trans)

    prev2curr_matrix = curr_g2e_matrix @ prev_e2g_matrix

    # get vector data
    # prev_vectors = copy.deepcopy(prev_data['map_annos'])
    # curr_vectors_ = copy.deepcopy(curr_data['map_annos'])
    # curr_vectors = dict()
    # for label, vecs in curr_vectors_.items():
    #     if len(vecs) > 0:
    #         vecs_20 = []
    #         for vec_ in vecs:
    #             if vec_.shape[0] != 20:
    #                 vec_20 = np.tile(vec_, (20, 1))[:20, :]
    #                 vecs_20.append(vec_20)
    #         curr_vectors[label] = vecs_20
    #     else:
    #         curr_vectors[label] = vecs
    prev_vectors = copy.deepcopy(get_vectorized_lines(prev_data['map_geoms']))
    curr_vectors = copy.deepcopy(get_vectorized_lines(curr_data['map_geoms']))
    #meta_info = curr_data['img_metas'].data
    #imgs = [mmcv.imread(i) for i in meta_info['img_filenames']]
    #cam_extrinsics = meta_info['cam_extrinsics']
    #cam_intrinsics = meta_info['cam_intrinsics']
    #ego2cams = meta_info['ego2cam']
    
    # transform prev vectors
    prev2curr_vectors = dict()
    for label, vecs in prev_vectors.items():
        if len(vecs) > 0:
            # vecs_20 = []
            # for vec_ in vecs:
            #     if vec_.shape[0] != 20:
            #         vec_20 = np.tile(vec_, (20, 1))[:20, :]
            #         vecs_20.append(vec_20)

            # vecs = np.stack(vecs_20, 0)
            vecs = np.stack(vecs, 0)
            vecs = torch.tensor(vecs)
            N, num_points, _ = vecs.shape
            # denormed_vecs = vecs * roi_size + origin # (num_prop, num_pts, 2)
            
            denormed_vecs = vecs
            denormed_vecs = torch.cat([
                denormed_vecs,
                denormed_vecs.new_zeros((N, num_points, 1)), # z-axis
                denormed_vecs.new_ones((N, num_points, 1)) # 4-th dim
            ], dim=-1) # (num_prop, num_pts, 4)

            transformed_vecs = torch.einsum('lk,ijk->ijl', prev2curr_matrix, denormed_vecs.double()).float()
            normed_vecs = (transformed_vecs[..., :2] - origin) / roi_size # (num_prop, num_pts, 2)
            normed_vecs = torch.clip(normed_vecs, min=0., max=1.)
            prev2curr_vectors[label] = normed_vecs
        else:
            prev2curr_vectors[label] = vecs
    
    # convert to ego space for visualization
    for label in prev2curr_vectors:
        if len(prev2curr_vectors[label]) > 0:
            prev2curr_vectors[label] = prev2curr_vectors[label] * roi_size + origin
        if len(curr_vectors[label]) > 0:
            curr_vecs = torch.tensor(np.stack(curr_vectors[label]))
            # curr_vectors[label] = curr_vecs * roi_size + origin
            curr_vectors[label] = curr_vecs
    # import pdb; pdb.set_trace()
    prev2curr_masks = draw_instance_masks(prev2curr_vectors, roi_size, origin, cfg)
    curr_masks = draw_instance_masks(curr_vectors, roi_size, origin, cfg)
    
    prev2curr_matchings = find_matchings(prev2curr_masks, curr_masks, thresh=0.01)

    # ## For viz purpose, may display the maps in perspective images
    # meta_info = curr_data['img_metas'].data
    # imgs = [mmcv.imread(i) for i in meta_info['img_filenames']]
    # cam_extrinsics = meta_info['cam_extrinsics']
    # cam_intrinsics = meta_info['cam_intrinsics']
    # ego2cams = meta_info['ego2cam']
    # cat2id = {'ped_crossing':0,
    #     'divider':1,
    #     'boundary':2}
    # renderer = Renderer(cat2id, (60, 30), 'nusc')
    # viz_dir = os.path.join(scene_dir, '{}_viz_perspective'.format(local_idx))
    # if not os.path.exists(viz_dir):
    #    os.makedirs(viz_dir)
    # renderer.render_camera_views_from_vectors(curr_vectors, imgs, 
    #            cam_extrinsics, cam_intrinsics, ego2cams, 2, viz_dir)

    # renderer.render_bev_from_vectors(curr_vectors, out_dir=None, specified_path='cur.png')
    # renderer.render_bev_from_vectors(prev2curr_vectors, out_dir=None, specified_path='prev2cur.png')
    # from PIL import Image 
    # background = Image.open("cur.png")
    # overlay = Image.open("prev2cur.png")
    # background = background.convert("RGBA")
    # overlay = overlay.convert("RGBA")
    # new_img = Image.blend(background, overlay, 0.5)
    # new_img.save("cur_overlapped.png","PNG")
    # import pdb; pdb.set_trace()
    
    return prev2curr_matchings, prev_vectors

def draw_polylines(vecs, roi_size, origin, cfg):
    results = []
    for line_coords in vecs:
        canvas = np.zeros((cfg[1], cfg[0]), dtype=np.uint8)
        coords = (line_coords - origin) / roi_size * torch.tensor(cfg)
        coords = coords.numpy()
        cv2.polylines(canvas, np.int32([coords]), False, color=1, thickness=3)
        result = np.flipud(canvas)
        if result.sum() < 20:
            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7,7))
            result = cv2.dilate(result, kernel, iterations=1)
        results.append(result)
    return results
        

def draw_polygons(vecs, roi_size, origin, cfg):
    results = []
    for poly_coords in vecs:
        mask = Image.new("L", size=(cfg[0], cfg[1]), color=0)
        coords = (poly_coords - origin) / roi_size * torch.tensor(cfg)
        coords = coords.numpy()
        vert_list = [(x, y) for x, y in coords]
        if not (coords[0] == coords[-1]).all():
            vert_list.append(vert_list[0])
        ImageDraw.Draw(mask).polygon(vert_list, outline=1, fill=1)
        result = np.flipud(np.array(mask))
        if result.sum() < 20:
            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7,7))
            result = cv2.dilate(result, kernel, iterations=1)
        results.append(result)
    return results
        

def draw_instance_masks(vectors, roi_size, origin, cfg):
    masks = {}
    for label, vecs in vectors.items():
        if label == 0:
            masks[label] = draw_polygons(vecs, roi_size, origin, cfg)
        else:
            masks[label] = draw_polylines(vecs, roi_size, origin, cfg)
    return masks

def find_matchings(src_masks, tgt_masks, thresh=0.1):
    """Find the matching of map elements between two temporally 
    connected frame

    Args:
        src_masks (_type_): instance masks of prev frame
        tgt_masks (_type_): instance masks of current frame
        thresh (float, optional): IOU threshold for matching. Defaults to 0.1.
    """
    matchings = {}
    for label, src_instances in src_masks.items():
        tgt_instances = tgt_masks[label]
        cost = np.zeros([len(src_instances), len(tgt_instances)])
        for i, src_ins in enumerate(src_instances):
            for j, tgt_ins in enumerate(tgt_instances):
                iou = _mask_iou(src_ins, tgt_ins)
                cost[i, j] = -iou
        row_ind, col_ind = linear_sum_assignment(cost)
        
        label_matching = [-1 for _ in range(len(src_instances))]
        label_matching_reverse = [-1 for _ in range(len(tgt_instances))]
        for i, j in zip(row_ind, col_ind):
            if -cost[i, j] > thresh:
                label_matching[i] = j
                label_matching_reverse[j] = i
        
        matchings[label] = (label_matching, label_matching_reverse)
    return matchings
def _mask_iou(mask1, mask2):
    intersection = (mask1 * mask2).sum()
    if intersection == 0:
        return 0.0
    union = np.logical_or(mask1, mask2).sum()
    return intersection / union
#map track end


def get_ego_status(nusc, nusc_can_bus, sample):
    ego_status = []
    ref_scene = nusc.get("scene", sample['scene_token'])
    try:
        pose_msgs = nusc_can_bus.get_messages(ref_scene['name'],'pose')
        steer_msgs = nusc_can_bus.get_messages(ref_scene['name'], 'steeranglefeedback')
        pose_uts = [msg['utime'] for msg in pose_msgs]
        steer_uts = [msg['utime'] for msg in steer_msgs]
        ref_utime = sample['timestamp']
        pose_index = locate_message(pose_uts, ref_utime)
        pose_data = pose_msgs[pose_index]
        steer_index = locate_message(steer_uts, ref_utime)
        steer_data = steer_msgs[steer_index]
        ego_status.extend(pose_data["accel"]) # acceleration in ego vehicle frame, m/s/s
        ego_status.extend(pose_data["rotation_rate"]) # angular velocity in ego vehicle frame, rad/s
        ego_status.extend(pose_data["vel"]) # velocity in ego vehicle frame, m/s
        ego_status.append(steer_data["value"]) # steering angle, positive: left turn, negative: right turn
    except:
        ego_status = [0] * 10
    
    return np.array(ego_status).astype(np.float32)

def get_global_sensor_pose(rec, nusc):
    lidar_sample_data = nusc.get('sample_data', rec['data']['LIDAR_TOP'])

    pose_record = nusc.get("ego_pose", lidar_sample_data["ego_pose_token"])
    cs_record = nusc.get("calibrated_sensor", lidar_sample_data["calibrated_sensor_token"])

    ego2global = transform_matrix(pose_record["translation"], Quaternion(pose_record["rotation"]), inverse=False)
    sensor2ego = transform_matrix(cs_record["translation"], Quaternion(cs_record["rotation"]), inverse=False)
    pose = ego2global.dot(sensor2ego)

    return pose

def obtain_sensor2top(nusc,
                      sensor_token,
                      l2e_t,
                      l2e_r_mat,
                      e2g_t,
                      e2g_r_mat,
                      sensor_type='lidar'):
    """Obtain the info with RT matric from general sensor to Top LiDAR.

    Args:
        nusc (class): Dataset class in the nuScenes dataset.
        sensor_token (str): Sample data token corresponding to the
            specific sensor type.
        l2e_t (np.ndarray): Translation from lidar to ego in shape (1, 3).
        l2e_r_mat (np.ndarray): Rotation matrix from lidar to ego
            in shape (3, 3).
        e2g_t (np.ndarray): Translation from ego to global in shape (1, 3).
        e2g_r_mat (np.ndarray): Rotation matrix from ego to global
            in shape (3, 3).
        sensor_type (str): Sensor to calibrate. Default: 'lidar'.

    Returns:
        sweep (dict): Sweep information after transformation.
    """
    sd_rec = nusc.get('sample_data', sensor_token)
    cs_record = nusc.get('calibrated_sensor',
                         sd_rec['calibrated_sensor_token'])
    pose_record = nusc.get('ego_pose', sd_rec['ego_pose_token'])
    data_path = str(nusc.get_sample_data_path(sd_rec['token']))
    if os.getcwd() in data_path:  # path from lyftdataset is absolute path
        data_path = data_path.split(f'{os.getcwd()}/')[-1]  # relative path
    sweep = {
        'data_path': data_path,
        'type': sensor_type,
        'sample_data_token': sd_rec['token'],
        'sensor2ego_translation': cs_record['translation'],
        'sensor2ego_rotation': cs_record['rotation'],
        'ego2global_translation': pose_record['translation'],
        'ego2global_rotation': pose_record['rotation'],
        'timestamp': sd_rec['timestamp']
    }

    l2e_r_s = sweep['sensor2ego_rotation']
    l2e_t_s = sweep['sensor2ego_translation']
    e2g_r_s = sweep['ego2global_rotation']
    e2g_t_s = sweep['ego2global_translation']

    # obtain the RT from sensor to Top LiDAR
    # sweep->ego->global->ego'->lidar
    l2e_r_s_mat = Quaternion(l2e_r_s).rotation_matrix
    e2g_r_s_mat = Quaternion(e2g_r_s).rotation_matrix
    R = (l2e_r_s_mat.T @ e2g_r_s_mat.T) @ (
        np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T)
    T = (l2e_t_s @ e2g_r_s_mat.T + e2g_t_s) @ (
        np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T)
    T -= e2g_t @ (np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T
                  ) + l2e_t @ np.linalg.inv(l2e_r_mat).T
    sweep['sensor2lidar_rotation'] = R.T  # points @ R.T + T
    sweep['sensor2lidar_translation'] = T
    return sweep

def nuscenes_data_prep(root_path,
                       can_bus_root_path,
                       info_prefix,
                       version,
                       dataset_name,
                       out_dir,
                       max_sweeps=10):
    """Prepare data related to nuScenes dataset.

    Related data consists of '.pkl' files recording basic infos,
    2D annotations and groundtruth database.

    Args:
        root_path (str): Path of dataset root.
        info_prefix (str): The prefix of info filenames.
        version (str): Dataset version.
        dataset_name (str): The dataset class name.
        out_dir (str): Output directory of the groundtruth database info.
        max_sweeps (int): Number of input consecutive frames. Default: 10
    """
    create_nuscenes_infos(
        root_path, out_dir, can_bus_root_path, info_prefix, version=version, max_sweeps=max_sweeps)


parser = argparse.ArgumentParser(description='Data converter arg parser')
parser.add_argument('dataset', metavar='kitti', help='name of the dataset')
parser.add_argument(
    '--root-path',
    type=str,
    default='./data/kitti',
    help='specify the root path of dataset')
parser.add_argument(
    '--canbus',
    type=str,
    default='./data',
    help='specify the root path of nuScenes canbus')
parser.add_argument(
    '--version',
    type=str,
    default='v1.0',
    required=False,
    help='specify the dataset version, no need for kitti')
parser.add_argument(
    '--max-sweeps',
    type=int,
    default=10,
    required=False,
    help='specify sweeps of lidar per example')
parser.add_argument(
    '--out-dir',
    type=str,
    default='./data/kitti',
    required='False',
    help='name of info pkl')
parser.add_argument('--extra-tag', type=str, default='kitti')
parser.add_argument(
    '--workers', type=int, default=4, help='number of threads to be used')
args = parser.parse_args()

if __name__ == '__main__':
    if args.dataset == 'nuscenes' and args.version != 'v1.0-mini':
        train_version = f'{args.version}-trainval'
        nuscenes_data_prep(
            root_path=args.root_path,
            can_bus_root_path=args.canbus,
            info_prefix=args.extra_tag,
            version=train_version,
            dataset_name='NuScenesDataset',
            out_dir=args.out_dir,
            max_sweeps=args.max_sweeps)
        test_version = f'{args.version}-test'
        nuscenes_data_prep(
            root_path=args.root_path,
            can_bus_root_path=args.canbus,
            info_prefix=args.extra_tag,
            version=test_version,
            dataset_name='NuScenesDataset',
            out_dir=args.out_dir,
            max_sweeps=args.max_sweeps)
    elif args.dataset == 'nuscenes' and args.version == 'v1.0-mini':
        train_version = f'{args.version}'
        nuscenes_data_prep(
            root_path=args.root_path,
            can_bus_root_path=args.canbus,
            info_prefix=args.extra_tag,
            version=train_version,
            dataset_name='NuScenesDataset',
            out_dir=args.out_dir,
            max_sweeps=args.max_sweeps)
