# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/005_data.tabular.ipynb (unless otherwise specified).

__all__ = ['TabularDataset', 'TabularDataLoader']

# Cell
from ..imports import *
from fastai2.tabular.all import *

# Cell
class TabularDataset():
    "A `Numpy` dataset from a `TabularPandas` object"
    def __init__(self, to):
        self.cats = to.cats.to_numpy().astype(np.long)
        self.conts = to.conts.to_numpy().astype(np.float32)
        self.ys = to.ys.to_numpy()
    def __getitem__(self, idx): return self.cats[idx], self.conts[idx], self.ys[idx]
    def __len__(self): return len(self.cats)
    @property
    def c(self): return 0 if self.ys is None else 1 if isinstance(self.ys[0], float) else len(np.unique(self.ys))

class TabularDataLoader(DataLoader):
    def __init__(self, dataset, bs=1, num_workers=0, device=None, train=False, **kwargs):
        device = ifnone(device, default_device())
        super().__init__(dataset, bs=min(bs, len(dataset)), num_workers=num_workers, shuffle=train, device=device, drop_last=train, **kwargs)
        self.device, self.shuffle = device, train
    def create_item(self, s): return s
    def get_idxs(self):
        idxs = Inf.count if self.indexed else Inf.nones
        if self.n is not None: idxs = list(range(len(self.dataset)))
        if self.shuffle: self.shuffle_fn()
        return idxs
    def create_batch(self, b):
        return self.dataset[b[0]:b[0]+self.bs]
    def shuffle_fn(self):
        "Shuffle dataset after each epoch"
        rng = np.random.permutation(len(self.dataset))
        self.dataset.cats = self.dataset.cats[rng]
        self.dataset.conts = self.dataset.conts[rng]
        self.dataset.ys = self.dataset.ys[rng]
    def to(self, device):
        self.device = device
#     def ds_to(self, device=None):
        self.dataset.cats = tensor(self.dataset.cats).to(device=self.device)
        self.dataset.conts = tensor(self.dataset.conts).to(device=self.device)
        self.dataset.ys = tensor(self.dataset.ys).to(device=self.device)