# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/01_data.loader.ipynb.

# %% ../../nbs/01_data.loader.ipynb 3
from __future__ import annotations
from ..import_essentials import *
from ..utils import get_config

# %% auto 0
__all__ = ['Dataset', 'ArrayDataset', 'BaseDataLoader', 'JaxDataloader', 'TorchDataloader', 'DataloaderBackends', 'DataLoader']

# %% ../../nbs/01_data.loader.ipynb 4
try:
    import torch.utils.data as torch_data
except ModuleNotFoundError:
    torch_data = None

# %% ../../nbs/01_data.loader.ipynb 6
class Dataset(ABC):
    """A pytorch-like abstract Dataset class."""

    def __getitem__(self, index):
        raise NotImplementedError

# %% ../../nbs/01_data.loader.ipynb 7
class ArrayDataset(Dataset):
    """Dataset wrapping tensors."""

    def __init__(
        self, 
        *arrays: jnp.DeviceArray # Numpy array with same first dimension
    ):
        assert all(arrays[0].shape[0] == arr.shape[0] for arr in arrays), \
            "All arrays must have the same dimension."
        self.arrays = arrays

    def __len__(self):
        return self.arrays[0].shape[0]

    def __getitem__(self, index):
        return tuple(arr[index] for arr in self.arrays)

# %% ../../nbs/01_data.loader.ipynb 12
class BaseDataLoader(ABC):
    """Dataloader Interface"""
    def __init__(
        self, 
        dataset,
        backend: str,
        batch_size: int = 1,    # Batch size
        shuffle: bool = False,  # If true, dataloader shuffles before sampling each batch
        drop_last: bool = False, # Drop last batches or not
        **kwargs # Aux arguments
    ):
        pass 
    
    def __len__(self):
        raise NotImplementedError
    
    def __next__(self):
        raise NotImplementedError
    
    def __iter__(self):
        raise NotImplementedError

# %% ../../nbs/01_data.loader.ipynb 14
class JaxDataloader(BaseDataLoader):
    """Dataloder in vanilla Jax"""

    def __init__(
        self, 
        dataset: Dataset,
        backend: str = 'jax', # Position argument
        batch_size: int = 1,  # Batch size
        shuffle: bool = False,  # If true, dataloader shuffles before sampling each batch
        drop_last: bool = False, # Drop last batches or not
        **kwargs
    ):
        self.dataset = dataset
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.drop_last = drop_last

        self.keys = hk.PRNGSequence(get_config().global_seed)
        self.data_len = len(dataset)  # Length of the dataset
        self.indices = jnp.arange(self.data_len) # available indices in the dataset
        self.pose = 0  # record the current position in the dataset
        self._shuffle()

    def _shuffle(self):
        if self.shuffle:
            self.indices = jax.random.permutation(next(self.keys), self.indices)
        
    def _stop_iteration(self):
        self.pose = 0
        self._shuffle()
        raise StopIteration

    def __len__(self):
        if self.drop_last:
            batches = len(self.dataset) // self.batch_size  # get the floor of division
        else:
            batches = -(len(self.dataset) // -self.batch_size)  # get the ceil of division
        return batches

    def __next__(self):
        if self.pose + self.batch_size <= self.data_len:
            batch_indices = self.indices[self.pose: self.pose + self.batch_size]
            batch_data = self.dataset[batch_indices, ...]
            self.pose += self.batch_size
            return batch_data
        elif self.pose < self.data_len and not self.drop_last:
            batch_indices = self.indices[self.pose:]
            batch_data = self.dataset[batch_indices, ...]
            self.pose += self.batch_size
            return batch_data
        else:
            self._stop_iteration()

    def __iter__(self):
        return self

# %% ../../nbs/01_data.loader.ipynb 16
# copy from https://jax.readthedocs.io/en/latest/notebooks/Neural_Network_and_Data_Loading.html#data-loading-with-pytorch
def _numpy_collate(batch):
    if isinstance(batch[0], np.ndarray):
        return np.stack(batch)
    elif isinstance(batch[0], (tuple, list)):
        transposed = zip(*batch)
        return [_numpy_collate(samples) for samples in transposed]
    else:
        return np.array(batch)

def _convert_dataset_pytorch(dataset: Dataset):
    class _TorchDataset(torch_data.Dataset):
        def __init__(self, dataset: Dataset): self.dataset = dataset
        def __len__(self): return len(self.dataset)
        def __getitem__(self, idx): return self.dataset[idx]
    
    return _TorchDataset(dataset)

# %% ../../nbs/01_data.loader.ipynb 17
class TorchDataloader(BaseDataLoader):
    """Use `Pytorch` to load batches. It requires [pytorch](https://pytorch.org/get-started/) to be installed."""
    
    def __init__(
        self, 
        dataset: Dataset,
        backend: str = 'pytorch', # positional argument
        batch_size: int = 1,  # batch size
        shuffle: bool = False,  # if true, dataloader shuffles before sampling each batch
        num_workers: int = 0, # number of workers
        drop_last: bool = False, # drop last batch or not
        **kwargs
    ):
        if torch_data is None:
            raise ModuleNotFoundError("`pytorch` library needs to be installed. Try `pip install torch`."
            "Please refer to pytorch documentation for details: https://pytorch.org/get-started/.")
        
        dataset = _convert_dataset_pytorch(dataset)
        self.dataloader = torch_data.DataLoader(
            dataset, 
            batch_size=batch_size, 
            shuffle=shuffle, 
            num_workers=num_workers, 
            drop_last=drop_last,
            collate_fn=_numpy_collate,
            **kwargs
        ) 

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

    def __next__(self):
        return next(self.dataloader)

    def __iter__(self):
        return self.dataloader.__iter__()

# %% ../../nbs/01_data.loader.ipynb 20
@dataclass(frozen=True)
class DataloaderBackends:
    jax: BaseDataLoader = JaxDataloader
    pytorch: BaseDataLoader = TorchDataloader
    tensorflow: BaseDataLoader = None
    merlin: BaseDataLoader = None

    __all__ = dict(
        jax=jax, pytorch=pytorch, tensorflow=tensorflow, merlin=merlin
    )

    def __getitem__(self, key):
        return self.__all__[key]

    @classmethod
    def supported(cls) -> List[str]:
        return [
            backend for backend, dl_cls in cls.__all__.items() if dl_cls is not None
        ]

# %% ../../nbs/01_data.loader.ipynb 21
def _dispatch_dataloader(
    backend: str # dataloader backend
) -> BaseDataLoader:
    """Return Dataloader class based on given `backend`"""

    backends = DataloaderBackends()
    if not backend in DataloaderBackends.supported():
        raise ValueError(f"backend=`{backend}` is either an invalid backend or not supported yet. "
            f"Should be one of {backends.supported}.")
    
    dl_cls = backends[backend]
    return dl_cls

# %% ../../nbs/01_data.loader.ipynb 24
class DataLoader(BaseDataLoader):
    """Main Dataloader class to load Numpy data batches"""
    def __init__(
        self,
        dataset: Dataset,
        backend: str, # Dataloader backend; Currently supports `['jax', 'pytorch']`
        batch_size: int = 1,  # batch size
        shuffle: bool = False,  # if true, dataloader shuffles before sampling each batch
        num_workers: int = 0, # number of workers
        drop_last: bool = False, # drop last batches or not
        **kwargs
    ):
        self.__class__ = _dispatch_dataloader(backend)
        self.__init__(
            dataset=dataset, 
            backend=backend, 
            batch_size=batch_size, 
            shuffle=shuffle, 
            num_workers=num_workers,
            drop_last=drop_last,
            **kwargs
        )
