import pickle
import numpy as np

from ._point_dataset import _Point_Dataset


class FullCloudsDataset(_Point_Dataset):

    setting = 'prim'

    def __init__(self, *args, **kwargs):
        super(FullCloudsDataset, self).__init__(*args, **kwargs)

    def __len__(self):
        return len(self._files)

    def __getitem__(self, item):
        processed_file_path = self._read_or_load(item)

        with processed_file_path.open('rb') as processed_file:
            pointcloud = pickle.load(processed_file)
        points = np.stack((pointcloud['x'], pointcloud['y'], pointcloud['z'],
                           pointcloud['xn'], pointcloud['yn'], pointcloud['zn']
                           ),
                          axis=-1)

        # When yopu want to return points and normal seperately
        # normal = np.stack((pointcloud['xn'], pointcloud['yn'], pointcloud['zn']), axis=-1)
        label = pointcloud['cl_idx']
        sample_idxs = self.sampling(points)

        return points[sample_idxs], label[sample_idxs]