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

# %% ../../nbs/loader.jax.ipynb 3
from __future__ import print_function, division, annotations
from ..imports import *
from ..datasets import ArrayDataset, JAXDataset
from . import BaseDataLoader
from ..utils import get_config, asnumpy
from ..tests import *
import jax_dataloader as jdl
from threading import Thread, Event
from queue import Queue

# %% auto 0
__all__ = ['EpochIterator', 'to_jax_dataset', 'DataLoaderJAX']

# %% ../../nbs/loader.jax.ipynb 4
def EpochIterator(
    data,
    batch_size: int,
    indices: Sequence[int]
):
    for i in range(0, len(indices), batch_size):
        idx = indices[i:i+batch_size]
        yield data[idx]

# %% ../../nbs/loader.jax.ipynb 5
@dispatch
def to_jax_dataset(dataset: JAXDataset):
    if isinstance(dataset, ArrayDataset):
        dataset.asnumpy()
    return dataset

@dispatch
def to_jax_dataset(dataset: HFDataset):
    return dataset.with_format('numpy')

# %% ../../nbs/loader.jax.ipynb 6
class DataLoaderJAX(BaseDataLoader):

    @typecheck
    def __init__(
        self, 
        dataset: Union[JAXDataset, HFDataset], 
        batch_size: int = 1,  # batch size
        shuffle: bool = False,  # if true, dataloader shuffles before sampling each batch
        num_workers: int = 0,  # how many subprocesses to use for data loading. Ignored.
        drop_last: bool = False,
        **kwargs
    ):
        self.key = jrand.PRNGKey(get_config().global_seed)
        self.dataset = to_jax_dataset(dataset)
        
        self.indices = np.arange(len(dataset))
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.drop_last = drop_last
    
    def __iter__(self):
        # shuffle (permutation) indices every epoch        
        indices = jrand.permutation(self.next_key(), self.indices).__array__() if self.shuffle else self.indices
        
        if self.drop_last:
            indices = indices[:len(self.indices) - len(self.indices) % self.batch_size]
        return EpochIterator(self.dataset, self.batch_size, indices)
    
    def next_key(self):
        self.key, subkey = jrand.split(self.key)
        return subkey
    
    def __len__(self):
        return len(self.indices) // self.batch_size + int(not self.drop_last)
