import numpy as np
import torch
from nuplan.common.actor_state.tracked_objects_types import (
    TrackedObjectType,
)

OBJECT_TYPE_DICT = {
    "vehicle": TrackedObjectType.VEHICLE,
    "pedestrian": TrackedObjectType.PEDESTRIAN,
    "bicycle": TrackedObjectType.BICYCLE,
    "traffic_cone": TrackedObjectType.TRAFFIC_CONE,
    "barrier": TrackedObjectType.BARRIER,
    "czone_sign": TrackedObjectType.CZONE_SIGN,
    "generic_object": TrackedObjectType.GENERIC_OBJECT,
}


def limit_period(val, offset=0.5, period=2 * np.pi):
    """Limit the value into a period for periodic function.

    Args:
        val (torch.Tensor | np.ndarray): The value to be converted.
        offset (float, optional): Offset to set the value range.
            Defaults to 0.5.
        period ([type], optional): Period of the value. Defaults to np.pi.

    Returns:
        (torch.Tensor | np.ndarray): Value in the range of
            [-offset * period, (1-offset) * period]
    """
    limited_val = val - torch.floor(val / period + offset) * period
    return limited_val


class LiDARAug(object):
    def __init__(self,
                 bda_aug_conf, is_train,
                 x_range='(-50.0, 50.0)',
                 y_range='(-50.0, 50.0)',
                 z_range='(-5, 20)',
                 ):
        for k in ['rot_lim', 'scale_lim', 'tran_lim']:
            bda_aug_conf[k] = eval(bda_aug_conf[k])
        self.bda_aug_conf = bda_aug_conf
        self.is_train = False
        self.x_range = eval(x_range)
        self.y_range = eval(y_range)
        self.z_range = eval(z_range)

    def sample_bda_augmentation(self):
        """Generate bda augmentation values based on bda_config."""
        if self.is_train:
            rotate_bda = np.random.uniform(*self.bda_aug_conf['rot_lim'])
            scale_bda = np.random.uniform(*self.bda_aug_conf['scale_lim'])
            flip_dx = np.random.uniform() < self.bda_aug_conf['flip_dx_ratio']
            flip_dy = np.random.uniform() < self.bda_aug_conf['flip_dy_ratio']
            translation_std = self.bda_aug_conf.get('tran_lim', [0.0, 0.0, 0.0])
            tran_bda = np.random.normal(scale=translation_std, size=3).T
        else:
            rotate_bda = 0
            scale_bda = 1.0
            flip_dx = False
            flip_dy = False
            tran_bda = np.zeros((1, 3), dtype=np.float32)
        return rotate_bda, scale_bda, flip_dx, flip_dy, tran_bda

    def bev_transform(self, gt_boxes, rotate_angle, scale_ratio, flip_dx,
                      flip_dy, tran_bda, rot_mat):
        if gt_boxes.shape[0] > 0:
            gt_boxes[:, :3] = (
                rot_mat @ gt_boxes[:, :3].unsqueeze(-1)).squeeze(-1)
            gt_boxes[:, 3:6] *= scale_ratio
            gt_boxes[:, 6] += rotate_angle
            if flip_dx:
                gt_boxes[:,
                         6] = 2 * torch.asin(torch.tensor(1.0)) - gt_boxes[:,
                                                                           6]
            if flip_dy:
                gt_boxes[:, 6] = -gt_boxes[:, 6]
            gt_boxes[:, 7:] = (
                rot_mat[:2, :2] @ gt_boxes[:, 7:].unsqueeze(-1)).squeeze(-1)
            gt_boxes[:, :3] = gt_boxes[:, :3] + tran_bda
        return gt_boxes

    def __call__(self, features, targets):
        # 1. filter box based on ranges
        # 2. filter label based on classes
        if 'dets' in targets and 'labels' in targets:
            boxes = targets['dets']
            labels = targets['labels']

            for t, (box, label) in enumerate(zip(boxes, labels)):
                label_mask = np.array([n in OBJECT_TYPE_DICT for n in label], dtype=np.bool_)
                label_mask = torch.from_numpy(label_mask)
                range_mask = ((box[:, 0] > self.x_range[0]) &
                              (box[:, 0] < self.x_range[1]) &
                              (box[:, 1] > self.y_range[0]) &
                              (box[:, 1] < self.y_range[1]))
                mask = range_mask & label_mask
                box_of_interest = box[mask]
                box_of_interest[:, 6] = limit_period(box_of_interest[:, 6])
                boxes[t] = box_of_interest.float()

                labels[t] = torch.from_numpy(np.array([OBJECT_TYPE_DICT[x].value for
                                                       x in label], dtype=np.int64))[mask]
            targets['dets'] = boxes
            targets['labels'] = labels

        rotate_bda, scale_bda, flip_dx, flip_dy, tran_bda = \
            self.sample_bda_augmentation()
        bda_mat = torch.zeros(4, 4)
        bda_mat[3, 3] = 1
        rotate_angle = torch.tensor(rotate_bda / 180 * np.pi)
        rot_sin = torch.sin(rotate_angle)
        rot_cos = torch.cos(rotate_angle)
        rot_mat = torch.Tensor([[rot_cos, -rot_sin, 0], [rot_sin, rot_cos, 0],
                                [0, 0, 1]])
        scale_mat = torch.Tensor([[scale_bda, 0, 0], [0, scale_bda, 0],
                                  [0, 0, scale_bda]])
        flip_mat = torch.Tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
        if flip_dx:
            flip_mat = flip_mat @ torch.Tensor([[-1, 0, 0], [0, 1, 0],
                                                [0, 0, 1]])
        if flip_dy:
            flip_mat = flip_mat @ torch.Tensor([[1, 0, 0], [0, -1, 0],
                                                [0, 0, 1]])
        bda_rot = flip_mat @ (scale_mat @ rot_mat)

        if 'dets' in targets:
            for idx, boxes in enumerate(targets['dets']):
                targets['dets'][idx] = self.bev_transform(boxes, rotate_bda, scale_bda,
                                                       flip_dx, flip_dy, tran_bda, bda_rot)
        # print('before bda')
        # print(features['lidars_warped'][-1][:, 0].max())
        # print(features['lidars_warped'][-1][:, 0].min())
        # print(features['lidars_warped'][-1][:, 1].max())
        # print(features['lidars_warped'][-1][:, 1].min())
        for idx, points in enumerate(features['lidars_warped']):
            points_aug = (bda_rot @ points[:, :3].unsqueeze(-1)).squeeze(-1)
            points[:, :3] = points_aug + tran_bda
            features['lidars_warped'][idx] = points

        # print('after bda')
        # print(features['lidars_warped'][-1][:, 0].max())
        # print(features['lidars_warped'][-1][:, 0].min())
        # print(features['lidars_warped'][-1][:, 1].max())
        # print(features['lidars_warped'][-1][:, 1].min())
        bda_mat[:3, :3] = bda_rot
        bda_mat[:3, 3] = torch.from_numpy(tran_bda)
        features['bda'] = bda_mat
        return features, targets