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

# %% ../nbs/core.ipynb 3
from __future__ import print_function, division, annotations
from .imports import *
from .utils import *
from .datasets import *
from .loaders import *

# %% auto 0
__all__ = ['SUPPORTED_DATASETS', 'DataloaderBackends', 'get_backend_compatibilities', 'DataLoader']

# %% ../nbs/core.ipynb 4
SUPPORTED_DATASETS = [
    JAXDataset,
    TorchDataset,
    TFDataset,
    HFDataset
]

# %% ../nbs/core.ipynb 5
@dataclass(frozen=True)
class DataloaderBackends:
    jax = DataLoaderJAX
    pytorch: BaseDataLoader = DataLoaderPytorch
    tensorflow: BaseDataLoader = DataLoaderTensorflow
    merlin: BaseDataLoader = None

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

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

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

# %% ../nbs/core.ipynb 6
def _get_backends() -> List[str]:
    """Return list of supported dataloader backends"""
    return DataloaderBackends().__all__.keys()


def _dispatch_dataloader(
    backend: str # dataloader backend
) -> BaseDataLoader:
    """Return Dataloader class based on given `backend`"""
    backends = DataloaderBackends()
    if not backend in backends.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/core.ipynb 7
def _check_backend_compatibility(ds, backend: str):
    return DataLoader(ds, backend=backend)

# %% ../nbs/core.ipynb 8
def get_backend_compatibilities() -> dict[str, list[type]]: # { backend: [supported datasets] }
    """Return list of supported dataloader backends for each dataset type"""

    ds = {
        JAXDataset: ArrayDataset(np.array([1,2,3])),
        TorchDataset: torch_data.Dataset(),
        TFDataset: tf.data.Dataset.from_tensor_slices(np.array([1,2,3])),
        HFDataset: hf_datasets.Dataset.from_dict({'a': [1,2,3]})
    }
    assert len(ds) == len(SUPPORTED_DATASETS)
    backends = {b: [] for b in _get_backends()}
    for b in _get_backends():
        for name, dataset in ds.items():
            try:
                _check_backend_compatibility(dataset, b)
                backends[b].append(name)
            except:
                pass

    return backends

# %% ../nbs/core.ipynb 9
class DataLoader:
    """Main Dataloader class to load Numpy data batches"""

    def __init__(
        self,
        dataset, # Dataset from which to load the data
        backend: Literal['jax', 'pytorch', 'tensorflow'], # Dataloader backend to load the dataset
        batch_size: int = 1,  # How many samples per batch to load
        shuffle: bool = False,  # If true, dataloader reshuffles every epoch
        drop_last: bool = False, # If true, drop the last incomplete batch
        **kwargs
    ):
        dl_cls = _dispatch_dataloader(backend)
        self.dataloader = dl_cls(
            dataset=dataset, 
            batch_size=batch_size, 
            shuffle=shuffle, 
            drop_last=drop_last,
            **kwargs
        )

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

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

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