import os
from PIL import Image

from nuscenes.nuscenes import NuScenes as NuScenesDevKit
from torch.utils.data import Dataset

from carbrain.datasets.util import AnnotationType


class NuScenes(Dataset):
    """The NuScenes dataset

    Args:
        root (str): The location of the root directory of the NuScenes
            dataset.
        version (str, optional): The version of the NuScenes dataset to be
            used. Possible: 'v1.0-mini', 'v1.0-trainval', 'v1.0-test' (depends
            on which datasets you have downloaded).
        split (str, optional): The split from the dataset to be used.
            Possible: 'train', 'val', 'test'
        verbose (bool, optional): Whether to print additional statistics while
            loading the dataset.
        annotation_type (DatasetType, optional):
            It is possible to generate different types of annotations for the
            NuScenes dataset. This argument allows for setting the desired
            annotation type. Possible: OBJECT_DETECTION_2D, OBJECT_DETECTION_3D
    """
    available_annotation_types = [
        AnnotationType.NONE,
        AnnotationType.OBJECT_DETECTION_2D,
        AnnotationType.OBJECT_DETECTION_3D,
    ]

    def __init__(self, root, split='v1.0-mini', transforms=None, verbose=False, annotation_type=AnnotationType.OBJECT_DETECTION_2D):
        self._nusc = NuScenesDevKit(
            version=split,
            dataroot=os.path.expanduser(root),
            verbose=verbose,
        )
        # check if the desired annotation type is possible
        if annotation_type not in NuScenes.available_annotation_types:
            raise ValueError('The annotation type {} is not possible for NuScenes, possible types: {}'.format(
                annotation_type.name, ', '.join([elem.name for elem in NuScenes.available_annotation_types])))
        self._annotation_type = annotation_type
        self._transforms = transforms

    def __getitem__(self, idx):
        sample = self._nusc.sample[idx]
        cam_front_data = self._nusc.get('sample_data', sample['data']['CAM_FRONT'])
        assert cam_front_data['is_key_frame']

        if self._annotation_type == AnnotationType.NONE:
            target = None
        else:
            for ann_id in sample['anns']:
                annotation = self._nusc.get('sample_annotations', ann_id)
            if self._annotation_type == AnnotationType.OBJECT_DETECTION_2D:
                raise NotImplementedError()
            elif self._annotation_type == AnnotationType.OBJECT_DETECTION_3D:
                raise NotImplementedError()

        img = Image.open(os.path.join(self._nusc.dataroot, cam_front_data['filename'])).convert('RGB')
        if self._transforms is not None:
            if target is None:
                img = self._transforms(img)
            else:
                img, target = self._transforms(img, target)

        return img, target

    def __len__(self):
        return len(self._nusc.sample)

    @property
    def mean(self):
        if self._nusc.version == 'v1.0-mini':
            return [0.0, 0.0, 0.0]
        elif self._nusc.version == 'v1.0-trainval':
            return [0.0, 0.0, 0.0]
        elif self._nusc.version == 'v1.0-test':
            return [0.0, 0.0, 0.0]
        else:
            # v1.0-private-test
            return [0.0, 0.0, 0.0]

    @property
    def std(self):
        if self._nusc.version == 'v1.0-mini':
            return [1.0, 1.0, 1.0]
        elif self._nusc.version == 'v1.0-trainval':
            return [1.0, 1.0, 1.0]
        elif self._nusc.version == 'v1.0-test':
            return [1.0, 1.0, 1.0]
        else:
            # v1.0-private-test
            return [1.0, 1.0, 1.0]



if __name__ == '__main__':
    nu = NuScenes(root=r'~/Downloads/nuscenes', verbose=True)
    img, boxes = nu[0]

