from typing import Iterator, List, Union

import numpy as np

from ...classes import DatasetInterface, UniformPointCloudData
from ...functions import read_uniform_point_cloud_data


def _copy(array: np.ndarray):
    if array is None:
        return None
    return np.copy(array)


class UniformPointCloudDataset(DatasetInterface):
    def __init__(self, path_list: List[str], memory_caching=False):
        self.path_list = path_list
        self.memory_caching = memory_caching
        self.cache: List[Union[UniformPointCloudData, None]] = [
            None for k in range(len(path_list))
        ]

    def read_file(self, path_index: int) -> UniformPointCloudData:
        path = self.path_list[path_index]
        if self.memory_caching:
            cached_data = self.cache[path_index]
            if cached_data is not None:
                return cached_data
        data = read_uniform_point_cloud_data(path)
        if self.memory_caching:
            self.cache[path_index] = data

        return data

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

    def __iter__(self) -> Iterator[UniformPointCloudData]:
        for path_index in range(len(self.path_list)):
            yield self.read_file(path_index)

    def __getitem__(self, indices) -> np.ndarray:
        if isinstance(indices, int):
            return self.read_file(indices)
        ret = []
        for index in indices:
            ret.append(self.read_file(index))
        return ret

    def shuffle(self) -> Iterator[np.ndarray]:
        indices = np.random.permutation(len(self.path_list))
        for index in indices:
            yield self[int(index)]
