import mmcv
import numpy as np
import pandas as pd
import tempfile
from os import path as osp
from pyquaternion import Quaternion
import pandaset as ps

from mmdet.datasets import DATASETS
from ..core import show_result
from ..core.bbox import Box3DMode, LiDARInstance3DBoxes
from .custom_3d import Custom3DDataset


@DATASETS.register_module()
class PandasetDataset(Custom3DDataset):
    r"""Pandaset Dataset.

    This class serves as the API for experiments on the Pandaset Dataset.

    Please refer to
    `https://github.com/scaleapi/pandaset-devkit`
    for data downloading.

    Args:
        ann_file (str): Path of annotation file.
        pipeline (list[dict], optional): Pipeline used for data processing.
            Defaults to None.
        data_root (str): Path of dataset root.
        classes (tuple[str], optional): Classes used in the dataset.
            Defaults to None.
        load_interval (int, optional): Interval of loading the dataset. It is
            used to uniformly sample the dataset. Defaults to 1.
        modality (dict, optional): Modality to specify the sensor data used
            as input. Defaults to None.
        box_type_3d (str, optional): Type of 3D box of this dataset.
            Based on the `box_type_3d`, the dataset will encapsulate the box
            to its original format then converted them to `box_type_3d`.
            Defaults to 'LiDAR' in this dataset. Available options includes

            - 'LiDAR': Box in LiDAR coordinates.
            - 'Depth': Box in depth coordinates, usually for indoor dataset.
            - 'Camera': Box in camera coordinates.
        filter_empty_gt (bool, optional): Whether to filter empty GT.
            Defaults to True.
        test_mode (bool, optional): Whether the dataset is in test mode.
            Defaults to False.
    """
    NameMapping = {
        'Animals - Other': 'animal',
        'Bicycle': 'bicycle',
        'Bus': 'bus',
        'Car': 'car',
        'Cones': 'cone',
        'Construction Signs': 'sign',
        'Emergency Vehicle': 'emergency_vehicle',
        'Medium-sized Truck': 'truck',
        'Motorcycle': 'motorcycle',
        'Motorized Scooter': 'mobility_device',
        'Other Vehicle - Construction Vehicle': 'other_vehicle',
        'Other Vehicle - Pedicab': 'other_vehicle',
        'Other Vehicle - Uncommon': 'other_vehicle',
        'Pedestrian': 'pedestrian',
        'Pedestrian with Object': 'pedestrian',
        'Personal Mobility Device': 'mobility_device',
        'Pickup Truck': 'other_vehicle',
        'Pylons': 'pylon',
        'Road Barriers': 'barrier',
        'Rolling Containers': 'container',
        'Semi-truck': 'truck',
        'Signs': 'sign',
        'Temporary Construction Barriers': 'barrier',
        'Towed Object': 'towed_object',
        'Train': 'train',
        'Tram / Subway': 'train'}

    SplitSequences = {
        # ~60% of the sequences, randomly chosen
        'train': ['014', '050', '079', '048', '093', '091', '063', '104',
                  '100', '092', '012', '047', '018', '006', '099', '085',
                  '035', '041', '052', '105', '030', '113', '002', '084',
                  '028', '119', '044', '005', '102', '034', '077', '064',
                  '067', '058', '019', '015', '037', '095', '120', '066',
                  '023', '071', '117', '098', '139', '038', '116', '046',
                  '088', '089', '040', '033', '016', '024', '122', '039',
                  '158', '069', '124', '123', '106'],
        # ~20% of the sequences, randomly chosen
        'val': ['045', '059', '055', '051', '020', '097', '073', '043', '003',
                '101', '027', '056', '011', '078', '080', '109', '042', '021',
                '094', '057'],
        # ~20% of the sequences, randomly chosen
        'test': ['074', '004', '086', '062', '068', '008', '001', '110', '053',
                 '115', '054', '065', '017', '103', '072', '013', '029', '090',
                 '112', '149', '070', '032']
    }

    CLASSES = ('car', 'bus', 'bicycle', 'motorcycle', 'truck', 'pickup',
               'emergency_vehicle', 'other_vehicle', 'towed_object',
               'pedestrian', 'mobility_device', 'train', 'cone', 'pylon',
               'sign', 'construction_sign', 'road_barrier',
               'construction_barrier', 'container', 'animal')

    def __init__(self,
                 ann_file,
                 pipeline=None,
                 data_root=None,
                 classes=None,
                 load_interval=1,
                 modality=None,
                 box_type_3d='LiDAR',
                 filter_empty_gt=True,
                 test_mode=False,
                 lidar_device=0):
        self.load_interval = load_interval
        super().__init__(
            data_root=data_root,
            ann_file=ann_file,
            pipeline=pipeline,
            classes=classes,
            modality=modality,
            box_type_3d=box_type_3d,
            filter_empty_gt=filter_empty_gt,
            test_mode=test_mode)

        if self.modality is None:
            self.modality = dict(
                use_camera=False,
                use_lidar=True,
                use_radar=False,
                use_map=False,
                use_external=False,
            )
            # To support other modalities, update dataset infos

        self.dataset = ps.DataSet(data_root)
        self.lidar_device = lidar_device

    def get_cat_ids(self, idx):
        """Get category distribution of single scene.

        Args:
            idx (int): Index of the data_info.

        Returns:
            dict[list]: for each category, if the current scene
                contains such boxes, store a list containing idx,
                otherwise, store empty list.
        """
        info = self.data_infos[idx]
        gt_names = info['cuboids']['gt_names']
        if self.lidar_device != -1:
            boxes_sensor_id = info["cuboids"]["sensor_ids"]
            # keep cuboids that are seen by a given device
            valid_indices = boxes_sensor_id != 1 - self.lidar_device
            gt_names = gt_names[valid_indices]
        gt_names = set(gt_names)

        cat_ids = []
        for name in gt_names:
            if name in self.CLASSES:
                cat_ids.append(self.cat2id[name])
        return cat_ids

    def load_annotations(self, ann_file):
        """Load annotations from ann_file.

        Args:
            ann_file (str): Path of the annotation file.

        Returns:
            list[dict]: List of annotations sorted by timestamps.
        """
        data = mmcv.load(ann_file)
        data_infos = data['infos']

        return data_infos

    def get_data_info(self, index):
        """Get data info according to the given index.

        Args:
            index (int): Index of the sample data to get.

        Returns:
            dict: Data information that will be passed to the data \
                preprocessing pipelines. It includes the following keys:

                - sample_idx (str): sample index
                - pts_filename (str): filename of point clouds
                - sweeps (list[dict]): infos of sweeps
                - timestamp (float): sample timestamp
                - img_filename (str, optional): image filename
                - lidar2img (list[np.ndarray], optional): transformations \
                    from lidar to different cameras
                - ann_info (dict): annotation info
        """
        info = self.data_infos[index]

        # standard protocal modified from SECOND.Pytorch
        input_dict = dict(
            pts_filename=info['lidar_path'],
            lidar_device=self.lidar_device,
            pose=info['lidar_pose'],
            sample_idx=info['sample_idx']
        )

        if self.modality['use_camera']:
            raise NotImplementedError("Camera data is not yet supported for Pandaset")
            # image_paths = []
            # lidar2img_rts = []
            # for cam_type, cam_info in info['cams'].items():
            #     image_paths.append(cam_info['data_path'])
            #     # obtain lidar to image transformation matrix
            #     lidar2cam_r = np.linalg.inv(cam_info['sensor2lidar_rotation'])
            #     lidar2cam_t = cam_info[
            #         'sensor2lidar_translation'] @ lidar2cam_r.T
            #     lidar2cam_rt = np.eye(4)
            #     lidar2cam_rt[:3, :3] = lidar2cam_r.T
            #     lidar2cam_rt[3, :3] = -lidar2cam_t
            #     intrinsic = cam_info['cam_intrinsic']
            #     viewpad = np.eye(4)
            #     viewpad[:intrinsic.shape[0], :intrinsic.shape[1]] = intrinsic
            #     lidar2img_rt = (viewpad @ lidar2cam_rt.T)
            #     lidar2img_rts.append(lidar2img_rt)

            # input_dict.update(
            #     dict(
            #         img_filename=image_paths,
            #         lidar2img=lidar2img_rts,
            #     ))

        if not self.test_mode:
            annos = self.get_ann_info(index)
            input_dict['ann_info'] = annos

        return input_dict

    def get_ann_info(self, index):
        """Get annotation info according to the given index.

        Args:
            index (int): Index of the annotation data to get.

        Returns:
            dict: Annotation information consists of the following keys:

                - gt_bboxes_3d (:obj:`LiDARInstance3DBoxes`): \
                    3D ground truth bboxes.
                - gt_labels_3d (np.ndarray): Labels of ground truths.
                - gt_names (list[str]): Class names of ground truths.
        """
        info = self.data_infos[index]
        gt_bboxes_3d = info["cuboids"]["gt_boxes"]
        # Add speed
        gt_bboxes_3d = np.pad(gt_bboxes_3d, ((0, 0), (0, 2)), 'constant',
                              constant_values=(0))
        gt_names_3d = info["cuboids"]["gt_names"]

        if self.lidar_device != -1:
            boxes_sensor_id = info["cuboids"]["sensor_ids"]
            # keep cuboids that are seen by a given device
            valid_indices = boxes_sensor_id != 1 - self.lidar_device
            gt_bboxes_3d = gt_bboxes_3d[valid_indices]
            gt_names_3d = gt_names_3d[valid_indices]

        gt_labels_3d = np.array([self.CLASSES.index(cat) if cat in self.CLASSES
                                 else -1 for cat in gt_names_3d])
        gt_labels_3d = np.array(gt_labels_3d)

        if 'gt_shape' in info:
            gt_shape = info['gt_shape']
            gt_bboxes_3d = np.concatenate([gt_bboxes_3d, gt_shape], axis=-1)

        # the pandaset box center is [0.5, 0.5, 0.5], we change it to be
        # the same as KITTI (0.5, 0.5, 0)
        gt_bboxes_3d = LiDARInstance3DBoxes(
            gt_bboxes_3d,
            box_dim=gt_bboxes_3d.shape[-1],
            origin=(0.5, 0.5, 0.5)).convert_to(self.box_mode_3d)

        anns_results = dict(
            gt_bboxes_3d=gt_bboxes_3d,
            gt_labels_3d=gt_labels_3d,
            gt_names=gt_names_3d
        )
        return anns_results

    def _format_bbox(self, results, jsonfile_prefix=None):
        """Convert the results to the standard format.

        Args:
            results (list[dict]): Testing results of the dataset.
            jsonfile_prefix (str): The prefix of the output jsonfile.
                You can specify the output directory/filename by
                modifying the jsonfile_prefix. Default: None.

        Returns:
            str: Path of the output json file.
        """
        all_annos = {}
        mapped_class_names = self.CLASSES

        print('Start to convert detection format...')
        all_annos = []
        for sample_id, det in enumerate(mmcv.track_iter_progress(results)):
            boxes = output_to_pandaset_frame(det, self.data_infos[sample_id],
                                             self.CLASSES)

            frame_id = str(self.data_infos[sample_id]["frame_idx"]).zfill(2)
            seq_id = str(self.data_infos[sample_id]["sequence"]).zfill(3)
            cur_det_file = osp.join(jsonfile_prefix, seq_id, 'predictions',
                                    'cuboids', ("{}.pkl.gz".format(frame_id)))
            mmcv.mkdir_or_exist(osp.dirname(cur_det_file))
            pd.DataFrame(boxes).to_pickle(cur_det_file)
            all_annos.append(
                {'preds': boxes,
                 'frame_id': frame_id,
                 'sequence': seq_id}
            )
            # seq_idx was converted to int in self.__getitem__` because strings

        res_path = osp.join(jsonfile_prefix, 'results_pandaset.json')
        print('Results written to', res_path)
        mmcv.dump(all_annos, res_path)

        return res_path

    def _evaluate_single(self,
                         result_path,
                         logger=None,
                         metric='bbox',
                         result_name='pts_bbox'):
        """Evaluation for a single model in Pandaset protocol.

        Args:
            result_path (str): Path of the result file.
            logger (logging.Logger | str | None): Logger used for printing
                related information during evaluation. Default: None.
            metric (str): Metric name used for evaluation. Default: 'bbox'.
            result_name (str): Result name in the metric prefix.
                Default: 'pts_bbox'.

        Returns:
            dict: Dictionary of evaluation details.
        """
        print("Evaluation not implemented yet")

        return dict()

    def format_results(self, results, jsonfile_prefix=None, csv_savepath=None):
        """Format the results to json (standard format for COCO evaluation).

        Args:
            results (list[dict]): Testing results of the dataset.
            jsonfile_prefix (str | None): The prefix of json files. It includes
                the file path and the prefix of filename, e.g., "a/b/prefix".
                If not specified, a temp file will be created. Default: None.
            csv_savepath (str | None): The path for saving csv files.
                It includes the file path and the csv filename,
                e.g., "a/b/filename.csv". If not specified,
                the result will not be converted to csv file.

        Returns:
            tuple: Returns (result_files, tmp_dir), where `result_files` is a \
                dict containing the json filepaths, `tmp_dir` is the temporal \
                directory created for saving json files when \
                `jsonfile_prefix` is not specified.
        """
        assert isinstance(results, list), 'results must be a list'
        assert len(results) == len(self), (
            'The length of results is not equal to the dataset len: {} != {}'.
            format(len(results), len(self)))

        if jsonfile_prefix is None:
            tmp_dir = tempfile.TemporaryDirectory()
            jsonfile_prefix = osp.join(tmp_dir.name, 'results')
        else:
            tmp_dir = None

        if not isinstance(results[0], dict):
            result_files = self._format_bbox(results, jsonfile_prefix)
        else:
            result_files = dict()
            for name in results[0]:
                print(f'\nFormating bboxes of {name}')
                results_ = [out[name] for out in results]
                tmp_file_ = osp.join(jsonfile_prefix, name)
                result_files.update(
                    {name: self._format_bbox(results_, tmp_file_)})

        if csv_savepath is not None:
            print("Conversion to CSV not implemented for pandaset. Skipping.")

        return result_files, tmp_dir

    def evaluate(self,
                 results,
                 metric='bbox',
                 logger=None,
                 jsonfile_prefix=None,
                 csv_savepath=None,
                 result_names=['pts_bbox'],
                 show=False,
                 out_dir=None):
        """Evaluation in Pandaset protocol.

        Args:
            results (list[dict]): Testing results of the dataset.
            metric (str | list[str]): Metrics to be evaluated.
            logger (logging.Logger | str | None): Logger used for printing
                related information during evaluation. Default: None.
            jsonfile_prefix (str | None): The prefix of json files. It includes
                the file path and the prefix of filename, e.g., "a/b/prefix".
                If not specified, a temp file will be created. Default: None.
            csv_savepath (str | None): The path for saving csv files.
                It includes the file path and the csv filename,
                e.g., "a/b/filename.csv". If not specified,
                the result will not be converted to csv file.
            show (bool): Whether to visualize.
                Default: False.
            out_dir (str): Path to save the visualization results.
                Default: None.

        Returns:
            dict[str, float]: Evaluation results.
        """
        result_files, tmp_dir = self.format_results(results, jsonfile_prefix,
                                                    csv_savepath)

        if isinstance(result_files, dict):
            results_dict = dict()
            for name in result_names:
                print(f'Evaluating bboxes of {name}')
                ret_dict = self._evaluate_single(result_files[name])
            results_dict.update(ret_dict)
        elif isinstance(result_files, str):
            results_dict = self._evaluate_single(result_files)

        if tmp_dir is not None:
            tmp_dir.cleanup()

        if show:
            self.show(results, out_dir)
        return results_dict

    def show(self, results, out_dir):
        """Results visualization.

        Args:
            results (list[dict]): List of bounding boxes results.
            out_dir (str): Output directory of visualization result.
        """
        for i, result in enumerate(results):
            example = self.prepare_test_data(i)
            points = example['points'][0]._data.numpy()
            data_info = self.data_infos[i]
            pts_path = data_info['lidar_path']
            file_name = osp.split(pts_path)[-1].split('.')[0]
            # for now we convert points into depth mode
            points = points[..., [1, 0, 2]]
            points[..., 0] *= -1
            inds = result['pts_bbox']['scores_3d'] > 0.1
            gt_bboxes = self.get_ann_info(i)['gt_bboxes_3d'].tensor
            gt_bboxes = Box3DMode.convert(gt_bboxes, Box3DMode.LIDAR,
                                          Box3DMode.DEPTH)
            gt_bboxes[..., 2] += gt_bboxes[..., 5] / 2
            pred_bboxes = result['pts_bbox']['boxes_3d'][inds].tensor.numpy()
            pred_bboxes = Box3DMode.convert(pred_bboxes, Box3DMode.LIDAR,
                                            Box3DMode.DEPTH)
            pred_bboxes[..., 2] += pred_bboxes[..., 5] / 2
            show_result(points, gt_bboxes, pred_bboxes, out_dir, file_name)


def output_to_pandaset_frame(detection, data_info, class_names):

    """Convert the output to a pandaset-like pandas dataframe.

    Args:
        detection (dict): Detection results.

    Returns:
        pd.DataFrame: All detection bboxes in pandaset format
                      (in world coordinates)
    """
    box3d = detection['boxes_3d']
    scores = detection['scores_3d'].numpy()
    labels = detection['labels_3d'].numpy()
    names = [class_names[ii] for ii in labels]

    zrot = data_info['zrot_world_to_ego']
    pose = data_info['lidar_pose']

    box_gravity_center = box3d.gravity_center.numpy()
    box_dims = box3d.dims.numpy()

    xs = box_gravity_center[:, 0]
    ys = box_gravity_center[:, 1]
    zs = box_gravity_center[:, 2]
    dxs = box_dims[:, 0]
    dys = box_dims[:, 1]
    dzs = box_dims[:, 2]
    yaws = box3d.yaw.numpy()

    # convert from normative coordinates to pandaset ego coordinates
    ego_xs = - ys
    ego_ys = xs
    ego_zs = zs
    ego_dxs = dys
    ego_dys = dxs
    ego_dzs = dzs
    ego_yaws = yaws

    # convert from pandaset ego coordinates to world coordinates
    # for the moment, an simplified estimation of the ego yaw is used
    # which sets ego_yaw = world_yaw + zrot_world_to_ego
    world_yaws = ego_yaws - zrot

    ego_centers = np.vstack([ego_xs, ego_ys, ego_zs]).T
    world_centers = ps.geometry.ego_to_lidar_points(ego_centers, pose)
    world_xs = world_centers[:, 0]
    world_ys = world_centers[:, 1]
    world_zs = world_centers[:, 2]
    # dx, dy, dz remain unchanged as the bbox orientation is handled by
    # the yaw information

    data_dict = {'position.x': world_xs,
                 'position.y': world_ys,
                 'position.z': world_zs,
                 'dimensions.x': ego_dxs,
                 'dimensions.y': ego_dys,
                 'dimensions.z': ego_dzs,
                 'yaw': world_yaws % (2 * np.pi),
                 'label': names,
                 'score': scores}

    return data_dict
