import mmcv
import numpy as np
import os
import pandaset as ps
from os import path as osp
from pyquaternion import Quaternion
import warnings
from mmdet3d.datasets import PandasetDataset

pandaset_categories = (
    '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 create_pandaset_infos(root_path,
                          info_prefix,
                          split='train',
                          version=None):
    """Create info file of pandaset 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.01-train'
        max_sweeps (int): Max number of sweeps.
            Default: 10
    """
    dataset = ps.DataSet(os.path.join(root_path, 'dataset'))
    available_sequences = dataset.sequences()

    split_sequences = PandasetDataset.SplitSequences[split]

    sequences = list(
        filter(lambda x: x in available_sequences, split_sequences))

    infos = []
    next_sample_idx = 0
    for seq in sequences:
        info, next_sample_idx = _get_sequence_infos(dataset, seq, root_path,
                                                    next_sample_idx)
        infos.extend(info)

    print(f'{split} split: {len(infos)}')
    data = dict(infos=infos, metadata={})
    info_name = f'{info_prefix}_infos_{split}'
    info_path = osp.join(root_path, f'{info_name}.pkl')
    mmcv.dump(data, info_path)
    print(f'Dumping to {info_path}')


def _get_zrot_from_pose(pose):
    """Get rotation around vertical axis from pose
    """
    yaxis_points_from_pose = ps.geometry.lidar_points_to_ego(
                np.array([[0, 0, 0], [0, 1., 0]]), pose)
    yaxis_from_pose = yaxis_points_from_pose[1, :] - yaxis_points_from_pose[0, :]

    if yaxis_from_pose[-1] >= 10**-1:
        warnings.warn("The car's pitch is supposed to be negligible " +
                      "sin(pitch) is >= 10**-1 ({})".format(yaxis_from_pose[-1]))

    # rotation angle in rads of the y axis around thz z axis
    zrot_world_to_ego = np.arctan2(-yaxis_from_pose[0], yaxis_from_pose[1])

    return zrot_world_to_ego


def _get_cuboids_infos(cuboids, pose, zrot_world_to_ego):
        """
        Get box informations in the unified normative coordinate system for a
        given frame
        """
        # get bboxes
        xs = cuboids['position.x'].to_numpy()
        ys = cuboids['position.y'].to_numpy()
        zs = cuboids['position.z'].to_numpy()
        dxs = cuboids['dimensions.x'].to_numpy()
        dys = cuboids['dimensions.y'].to_numpy()
        dzs = cuboids['dimensions.z'].to_numpy()
        yaws = cuboids['yaw'].to_numpy()
        labels = cuboids['label'].to_numpy()
        sensor_ids = cuboids["cuboids.sensor_id"].to_numpy()

        del cuboids  # There seem to be issues with the automatic deletion of pandas datasets sometimes

        labels = np.array([PandasetDataset.NameMapping.get(lab, lab)
                           for lab in labels])

        # Compute the center points coordinates in ego coordinates
        centers = np.vstack([xs, ys, zs]).T
        ego_centers = ps.geometry.lidar_points_to_ego(centers, pose)

        # Compute the yaw in ego coordinates
        # The following implementation supposes that the pitch of the car is
        # negligible compared to its yaw, in order to be able to express the
        # bbox coordinates in the ego coordinate system with an {axis aligned
        # box + yaw} only representation
        ego_yaws = yaws + zrot_world_to_ego

        # Pandaset ego coordinates are:
        # - x pointing to the right
        # - y pointing to the front
        # - z pointing up
        # Normative coordinates are:
        # - x pointing foreward
        # - y pointings to the left
        # - z pointing to the top
        # So a transformation is required to the match the normative coordinates
        ego_xs = ego_centers[:, 1]
        ego_ys = -ego_centers[:, 0]
        ego_zs = ego_centers[:, 2]
        ego_dxs = dys
        ego_dys = dxs  # stays >= 0
        ego_dzs = dzs

        ego_boxes = np.vstack([ego_xs, ego_ys, ego_zs, ego_dxs, ego_dys, ego_dzs, ego_yaws]).T

        infos = dict(gt_boxes=ego_boxes.astype(np.float32),
                     gt_names=labels,
                     sensor_ids=sensor_ids)

        return infos


def _get_sequence_infos(dataset, seq, root_path, next_sample_idx=0):
        """
        Generate the dataset infos dict for each sample of the dataset.
        For each sample, this dict contains:
            - the sequence index
            - the frame index
            - the path to the lidar data
            - the path to the bounding box annotations
        TODO: Complete with images and calibration informations
        """
        s = dataset[seq]
        s.load_lidar()
        s.lidar._load_poses()
        s.load_cuboids()
        if len(s.lidar.data) > 100:
            raise ValueError(
                "The implementation for this dataset assumes that each " +
                "sequence is no longer than 100 frames. " +
                "The current sequence has {}".format(len(s.lidar.data))
            )
        infos = []
        for ii in range(len(s.lidar.data)):

            pose = dataset[seq].lidar.poses[ii]
            cuboids = dataset[seq].cuboids[ii]
            zrot_world_to_ego = _get_zrot_from_pose(pose)
            infos.append({
                'sample_idx': next_sample_idx,
                'lidar_pose': pose,
                'sequence': seq,
                'frame_idx': ii,
                'lidar_path': os.path.join(root_path, 'dataset', seq, 'lidar',
                                           ("{:02d}.pkl.gz".format(ii))),
                'cuboids': _get_cuboids_infos(cuboids, pose, zrot_world_to_ego),
                'zrot_world_to_ego': zrot_world_to_ego
            })
            next_sample_idx += 1
        del dataset._sequences[seq]

        return infos, next_sample_idx

