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

# %% auto 0
__all__ = ['TSCategorize', 'TSRegression', 'TSForecasting', 'get_tsimage_dls', 'NumpyTensor', 'ToNumpyTensor', 'TSTensor',
           'ToTSTensor', 'show_tuple', 'TSLabelTensor', 'TSMaskTensor', 'ToFloat', 'ToInt', 'TSClassification',
           'TSMultiLabelClassification', 'NumpyTensorBlock', 'TSTensorBlock', 'TorchDataset', 'NumpyDataset',
           'TSDataset', 'NoTfmLists', 'TSTfmdLists', 'NumpyDatasets', 'tscoll_repr', 'TSDatasets', 'add_ds',
           'NumpyDataLoader', 'TSDataLoader', 'NumpyDataLoaders', 'TSDataLoaders', 'StratifiedSampler', 'get_c',
           'get_best_dl_params', 'get_best_dls_params', 'get_ts_dls', 'get_ts_dl', 'get_subset_dl',
           'get_time_per_batch', 'get_dl_percent_per_epoch']

# %% ../../nbs/006_data.core.ipynb 3
import warnings
from types import MethodType

import matplotlib.colors as mcolors
from fastai.data.block import CategoryBlock, DataBlock
from fastai.data.core import DataLoaders, Datasets, TfmdDL, TfmdLists
from fastai.data.load import DataLoader
from fastai.data.transforms import (
    Categorize,
    Category,
    CategoryMap,
    ItemGetter,
    MultiCategorize,
    MultiCategory,
    RandomSplitter,
)
from fastai.losses import BCEWithLogitsLossFlat, CrossEntropyLossFlat, MSELossFlat
from fastai.vision.data import get_grid
from fasttransform import DisplayedTransform, Pipeline, Transform
from matplotlib import rcParams
from matplotlib.ticker import PercentFormatter
from sklearn.model_selection import StratifiedKFold

from ..imports import *
from ..utils import *

warnings.filterwarnings("ignore", category=UserWarning)

# %% ../../nbs/006_data.core.ipynb 6
class NumpyTensor(TensorBase):
    "Returns a `tensor` with subclass `NumpyTensor` that has a show method"

    def __new__(cls, o, dtype=None, device=None, copy=None, requires_grad=False, **kwargs):
        o = torch.asarray(o, dtype=dtype, device=device, copy=copy, requires_grad=requires_grad)
        res = cast(o, cls) # if the tensor results in a dtype torch.float64 a copy is made as dtype torch.float32
        for k,v in kwargs.items(): setattr(res, k, v)
        return res

    @property
    def data(self): return cast(self, Tensor)

    def __repr__(self):
        if self.ndim > 0: return f'NumpyTensor(shape:{tuple(self.shape)}, device={self.device}, dtype={self.dtype})'
        else: return f'NumpyTensor([{self.data}], device={self.device}, dtype={self.dtype})'


    def show(self, ax=None, ctx=None, title=None, **kwargs):
        if self.ndim == 0: return str(self.data)
        elif self.ndim != 2: self = type(self)(to2d(self))
        if not isinstance(self,np.ndarray): self = self.detach().cpu().numpy()
        ax = ifnone(ax, ctx)
        if ax is None: _, ax = plt.subplots(**kwargs)
        ax.plot(self.T)
        ax.axis(xmin=0, xmax=self.shape[-1] - 1)
        title_color = kwargs['title_color'] if 'title_color' in kwargs else rcParams['axes.labelcolor']
        if title is not None:
            if is_listy(title): title = str(title)[1:-1]
            ax.set_title(title, weight='bold', color=title_color)
        plt.tight_layout()
        return ax


class ToNumpyTensor(Transform):
    "Transforms an object into NumpyTensor"
    def encodes(self, o): return NumpyTensor(o)
    def decodes(self, o): return o.detach().cpu().numpy()

# %% ../../nbs/006_data.core.ipynb 7
class TSTensor(TensorBase):
    '''Returns a `tensor` with subclass `TSTensor` that has a show method'''

    def __new__(cls, o, dtype=None, device=None, copy=None, requires_grad=False, **kwargs):
        o = torch.asarray(o, dtype=dtype, device=device, copy=copy, requires_grad=requires_grad)
        res = cast(o, cls) # if the tensor results in a dtype torch.float64 a copy is made as dtype torch.float32
        for k,v in kwargs.items(): setattr(res, k, v)
        return res

    @property
    def data(self): return cast(self, Tensor)

    def show(self, ax=None, ctx=None, title=None, **kwargs):
        if self.ndim == 0: return str(self.data)
        elif self.ndim != 2: self = type(self)(to2d(self))
        if not isinstance(self,np.ndarray): self = self.detach().cpu().numpy()
        ax = ifnone(ax, ctx)
        if ax is None: _, ax = plt.subplots(**kwargs)
        ax.plot(self.T)
        ax.axis(xmin=0, xmax=self.shape[-1] - 1)
        title_color = kwargs['title_color'] if 'title_color' in kwargs else rcParams['axes.labelcolor']
        if title is not None:
            if is_listy(title): title = str(title)[1:-1]
            ax.set_title(title, weight='bold', color=title_color)
        plt.tight_layout()
        return ax

    @property
    def vars(self):
        return self.shape[-2]

    @property
    def len(self): return self.shape[-1]

    def __repr__(self):
        if self.ndim > 3:
            return f'TSTensor(shape:{tuple(self.shape)}, device={self.device}, dtype={self.dtype})'
        elif self.ndim == 3:
            return f'TSTensor(samples:{self.shape[-3]}, vars:{self.shape[-2]}, len:{self.shape[-1]}, device={self.device}, dtype={self.dtype})'
        elif self.ndim == 2:
            return f'TSTensor(vars:{self.shape[-2]}, len:{self.shape[-1]}, device={self.device}, dtype={self.dtype})'
        elif self.ndim == 1:
            return f'TSTensor(len:{self.shape[-1]}, device={self.device}, dtype={self.dtype})'
        else: return f'TSTensor([{self.data}], device={self.device}, dtype={self.dtype})'

# %% ../../nbs/006_data.core.ipynb 8
class ToTSTensor(Transform):
    "Transforms an object into TSTensor"
    def encodes(self, o): return TSTensor(o)
    def decodes(self, o): return o.data


@delegates(plt.subplots)
def show_tuple(tup, **kwargs):
    "Display a timeseries plot from a decoded tuple"
    if len(tup) == 1: title = 'unlabeled'
    elif is_listy(tup[1]): title = str(tup[1])[1:-1]
    else: title = str(tup[1])
    tup[0].show(title=title, **kwargs)

# %% ../../nbs/006_data.core.ipynb 27
class TSLabelTensor(NumpyTensor):
    def __repr__(self):
        if self.ndim == 0: return f'{self.data}'
        else: return f'TSLabelTensor(shape:{tuple(self.shape)}, device={self.device}, dtype={self.dtype})'

class TSMaskTensor(NumpyTensor):
    def __repr__(self):
        if self.ndim == 0: return f'{self.data}'
        else: return f'TSMaskTensor(shape:{tuple(self.shape)}, device={self.device}, dtype={self.dtype})'

# %% ../../nbs/006_data.core.ipynb 30
class ToFloat(Transform):
    "Transforms an object dtype to float (vectorized)"
    vectorized=True
    loss_func=MSELossFlat()
    def encodes(self, o:torch.Tensor): return o.float()
    def encodes(self, o): return np.asarray(o, dtype=np.float32)
    def decodes(self, o):
        if o.ndim==0: return TitledFloat(o)
        else:
            return TitledTuple(o.cpu().numpy().tolist())


class ToInt(Transform):
    "Transforms an object dtype to int"
    def encodes(self, o:torch.Tensor): return o.long()
    def encodes(self, o): return np.asarray(o).astype(np.float32).astype(np.int64)
    def decodes(self, o):
        if o.ndim==0: return TitledFloat(o)
        else:
            return TitledTuple(o.cpu().numpy().tolist())


class TSClassification(DisplayedTransform):
    "Vectorized, reversible transform of category string to `vocab` id"
    loss_func,order,vectorized=CrossEntropyLossFlat(),1,True

    def __init__(self, vocab=None, sort=True):
        if vocab is not None: vocab = CategoryMap(vocab, sort=sort, add_na=False)
        store_attr()

    def setups(self, dset):
        if self.vocab is None and dset is not None:
            dset = np.asarray(dset).flatten()
            self.vocab = CategoryMap(dset, sort=self.sort, add_na=False)
        self.c = len(self.vocab)
        self.vocab_keys = np.array(list(self.vocab.o2i.keys()))[:, None]

    def encodes(self, o:torch.Tensor): return o
    def encodes(self, o):
        try:
            if not hasattr(o, "ndim") or o.ndim <= 1:
                return TensorCategory((self.vocab_keys == o).argmax(axis=0))
            else:
                return TensorCategory((self.vocab_keys == o.flatten()).argmax(axis=0).reshape(*o.shape))
        except KeyError as e:
            raise KeyError(f"Label '{o}' was not included in the training dataset") from e
    def decodes(self, o):
        if not is_iter(o):
            return Category(self.vocab[o])
        elif o.ndim <= 1:
            return stack([Category(self.vocab[oi]) for oi in o])
        else:
            return stack(MultiCategory(self.vocab[o.flatten()])).reshape(*o.shape)


TSCategorize = TSClassification
TSRegression = ToFloat
TSForecasting = ToFloat

# %% ../../nbs/006_data.core.ipynb 35
class TSMultiLabelClassification(Categorize):
    "Reversible combined transform of multi-category strings to one-hot encoded `vocab` id"
    loss_func,order=BCEWithLogitsLossFlat(),1
    def __init__(self, c=None, vocab=None, add_na=False, sort=True):
        super().__init__(vocab=vocab,add_na=add_na,sort=sort)
        self.c = c

    def setups(self, dsets):
        if not dsets: return
        if self.vocab is None:
            vals = set()
            for b in dsets: vals = vals.union(set(b))
            self.vocab = CategoryMap(list(vals), add_na=self.add_na)
            if self.c is None: self.c = len(self.vocab)
            if not self.c: warn("Couldn't infer the number of classes, please pass a value for `c` at init")

    def encodes(self, o:torch.Tensor): return o
    def encodes(self, o):
        if not all(elem in self.vocab.o2i.keys() for elem in o):
            diff = [elem for elem in o if elem not in self.vocab.o2i.keys()]
            diff_str = "', '".join(diff)
            raise KeyError(f"Labels '{diff_str}' were not included in the training dataset")
        return TensorMultiCategory(one_hot([self.vocab.o2i[o_] for o_ in o], self.c).float())
    def decodes(self, o):
        if o.ndim == 2:
            return MultiCategory([self.vocab[o_] for o_ in o])
        else:
            return MultiCategory(self.vocab[o])

# %% ../../nbs/006_data.core.ipynb 36
class NumpyTensorBlock():
    def __init__(self, type_tfms=None, item_tfms=None, batch_tfms=None, dl_type=None, dls_kwargs=None):
        self.type_tfms  =                 L(type_tfms)
        self.item_tfms  = ToNumpyTensor + L(item_tfms)
        self.batch_tfms =                 L(batch_tfms)
        self.dl_type,self.dls_kwargs = dl_type,({} if dls_kwargs is None else dls_kwargs)

class TSTensorBlock():
    def __init__(self, type_tfms=None, item_tfms=None, batch_tfms=None, dl_type=None, dls_kwargs=None):
        self.type_tfms  =              L(type_tfms)
        self.item_tfms  = ToTSTensor + L(item_tfms)
        self.batch_tfms =              L(batch_tfms)
        self.dl_type,self.dls_kwargs = dl_type,({} if dls_kwargs is None else dls_kwargs)

# %% ../../nbs/006_data.core.ipynb 38
class TorchDataset():
    def __init__(self, X, y=None): self.X, self.y = X, y
    def __getitem__(self, idx): return (self.X[idx],) if self.y is None else (self.X[idx], self.y[idx])
    def __len__(self): return len(self.X)


class NumpyDataset():
    def __init__(self, X, y=None, types=None): self.X, self.y, self.types = X, y, types
    def __getitem__(self, idx):
        if self.types is None: return (self.X[idx], self.y[idx]) if self.y is not None else (self.X[idx])
        else: return (self.types[0](self.X[idx]), self.types[1](self.y[idx])) if self.y is not None else (self.types[0](self.X[idx]))
    def __len__(self): return len(self.X)


class TSDataset():
    _types = TSTensor, TSLabelTensor
    def __init__(self, X, y=None, split=None, sel_vars=None, sel_steps=None, types=None, dtype=None, device=None):
        self.X, self.y, self.split = X, y, split
        self.sel_vars = sel_vars
        self.sel_steps = sel_steps
        self.multi_idx = sel_vars is not None or sel_steps is not None
        if types is not None: self._types = listify(types)
        self.dtype, self.device = dtype, device
    def __getitem__(self, idx):
        if self.split is not None:
            idx = self.split[idx]
        if hasattr(self.X, 'oindex'):
            X = self._types[0](self.X.oindex[idx, self.sel_vars, self.sel_steps] if self.multi_idx \
                               else self.X.oindex[idx], device=self.device, dtype=self.dtype)
        elif hasattr(self.X, 'compute'):
            X = self._types[0](self.X[idx, self.sel_vars, self.sel_steps].compute() if self.multi_idx \
                               else self.X[idx].compute(), device=self.device, dtype=self.dtype)
        else:
            X = self._types[0](self.X[idx, self.sel_vars, self.sel_steps] if self.multi_idx \
                               else self.X[idx], device=self.device, dtype=self.dtype)
        if self.y is None:
            return (X, )
        if hasattr(self.y, 'oindex'):
            y = self._types[1](self.y.oindex[idx], device=self.device, dtype=self.dtype)
        elif hasattr(self.X, 'compute'):
            y = self._types[1](self.y[idx].compute(), device=self.device, dtype=self.dtype)
        else:
            y = self._types[1](self.y[idx], device=self.device, dtype=self.dtype)
        return (X, y)
    def __len__(self): return len(self.X) if self.split is None else len(self.split)

# %% ../../nbs/006_data.core.ipynb 40
def _flatten_list(lst):
    "Flattens a list of lists with splits"

    def __flatten_list(lst):
        if lst is None:
            return L([])
        if not hasattr(lst, "__iter__"):
            lst = [lst]

        # clean_up_list
        if len(lst) > 10:
            return lst
        else:
            lst = [l for l in lst if l is not None or (hasattr(l, "__len__") and len(l) == 0)]

        # count lists
        n_lists = sum([hasattr(l, "__iter__") for l in lst])
        if n_lists == 0:
            return lst
        elif len(lst) == n_lists == 1:
            return lst[0]
        else:
            output = np.concatenate([l if hasattr(l, "__iter__") else [l] for l in lst])
            if len(output) == 0:
                return L([])
            return output.astype(smallest_dtype(int(max(output))))

    output = __flatten_list(lst)
    if len(output) == 0: return output
    dtype = smallest_dtype(np.max(output))
    return np.asarray(output, dtype=dtype)

def _remove_brackets(l):
    return [li if (not li or not is_listy(li) or len(li) > 1) else li[0] for li in l]

class NoTfmLists(TfmdLists):
    def __init__(self, items, tfms=None, splits=None, split_idx=None, types=None, do_setup=False, **kwargs):
        self.splits = ifnone(splits, L(np.arange(len(items)).tolist(),[]))
        self._splits = _flatten_list(self.splits)
        store_attr('items,types,split_idx')
        self.tfms = Pipeline(split_idx=split_idx)
    def subset(self, i, **kwargs): return type(self)(self.items, splits=self.splits[i], split_idx=i, do_setup=False, types=self.types,
                                                     **kwargs)
    def __getitem__(self, it):
        if hasattr(self.items, 'oindex'): return self.items.oindex[self._splits[it]]
        else: return self.items[self._splits[it]]
    def __len__(self): return len(self._splits)
    def __repr__(self):
        if hasattr(self.items, "shape"):
            return f"{self.__class__.__name__}: {self.items.__class__.__name__}{(len(self), *self.items.shape[1:])}"
        else:
            return f"{self.__class__.__name__}: {self.items.__class__.__name__}(0)"
    def _new(self, items, split_idx=None, **kwargs):
        split_idx = ifnone(split_idx, self.split_idx)
        return type(self)(items, split_idx=split_idx, do_setup=False, types=self.types, **kwargs)
    def decode(self, o, **kwargs): return o
    def new_empty(self): return self._new([])

NoTfmLists.train, NoTfmLists.valid = add_props(lambda i,x: x.subset(i))

class TSTfmdLists(TfmdLists):
    def __getitem__(self, it):
        # res = self._get(it)
        if hasattr(self.items, 'oindex'): res = self.items.oindex[it]
        else: res = self.items[it]
        if self._after_item is None: return res
        else: return self._after_item(res)

# %% ../../nbs/006_data.core.ipynb 47
@delegates(Datasets.__init__)
class NumpyDatasets(Datasets):
    "A dataset that creates tuples from X (and y) and applies `tfms` of type item_tfms"
    typs = NumpyTensor,tensor
    def __init__(self, X=None, y=None, items=None, tfms=None, tls=None, n_inp=None, dl_type=None, inplace=True, **kwargs):

        self.tfms, self.inplace = tfms, inplace

        if X is not None:
            if not hasattr(X, '__array__'): X = np.asarray(X)
            elif hasattr(X, "iloc"): X = to3d(X)
        if y is not None:
            if not hasattr(y, '__array__'):  y = np.asarray(y)
            elif hasattr(y, "iloc"): y = toarray(y)

        if tls is None:
            items = tuple((X,)) if y is None else tuple((X, y))

            if tfms is None:
                self.tfms, lts = [None] * len(items), [NoTfmLists] * len(items)
            else:
                self.tfms = _remove_brackets(tfms)
                lts = [NoTfmLists if t is None else TSTfmdLists if getattr(t, 'vectorized', None) else TfmdLists for t in self.tfms]

            self.tls = L(lt(item, t, **kwargs) for lt,item,t in zip(lts, items, self.tfms))
            # if len(self.tls) > 0 and len(self.tls[0]) > 0:
            #     self.typs = [type(tl[0]) if isinstance(tl[0], torch.Tensor) else self.typs[i] for i,tl in enumerate(self.tls)]
            self.ptls = L([typ(stack(tl[:])) for i,(tl,typ) in enumerate(zip(self.tls,self.typs))]) if inplace else self.tls
        else:
            self.tls = tls
            # if len(self.tls) > 0 and len(self.tls[0]) > 0:
            #     self.typs = [type(tl[0]) if isinstance(tl[0], torch.Tensor) else self.typs[i] for i,tl in enumerate(self.tls)]
            self.ptls = L([typ(stack(tl[:])) for i,(tl,typ) in enumerate(zip(self.tls,self.typs))]) if inplace and len(tls[0]) != 0 else tls

        self.n_inp = 1
        if kwargs.get('splits', None) is not None:
            split_idxs = _flatten_list(kwargs['splits'])
        else:
            split_idxs = _flatten_list(np.arange(len(self)))
        self.split_idxs = split_idxs

    def __getitem__(self, it):
        if self.inplace:
            return tuple([ptl[it] for ptl in self.ptls])
        else:
            return tuple([typ(stack(ptl[it])) for i,(ptl,typ) in enumerate(zip(self.ptls,self.typs))])

    def subset(self, i):
        if is_indexer(i):
            return type(self)(tls=L([tl.subset(i) for tl in self.tls]), inplace=self.inplace, tfms=self.tfms,
                              splits=None if self.splits is None else self.splits[i], split_idx=i)
        else:
            splits = None if self.splits is None else L(np.arange(len(i)).tolist())
            return type(self)(*self[i], inplace=True, tfms=None, splits=splits, split_idx=ifnone(self.split_idx, 1))

    def __len__(self): return len(self.tls[0])

    def _new(self, X, y=None, **kwargs):
        return type(self)(X, y=y, tfms=self.tfms, inplace=self.inplace, do_setup=False, **kwargs)

    def new_empty(self): return type(self)(tls=[tl.new_empty() for tl in self.tls], n_inp=self.n_inp, inplace=self.inplace)

    def show_at(self, idx, **kwargs):
        self.show(self[idx], **kwargs)
        plt.show()

    def __repr__(self): return tscoll_repr(self)


def tscoll_repr(c, max_n=10):
    "String repr of up to `max_n` items of (possibly lazy) collection `c`"
    _len = len(c)
    if _len == 0: return coll_repr(c)
    return f'(#{_len}) {L(c[i] for i in range(min(len(c), max_n)))} ...]'

# %% ../../nbs/006_data.core.ipynb 48
@delegates(Datasets.__init__)
class TSDatasets(Datasets):
    """A dataset that creates tuples from X (and optionally y) and applies `item_tfms`"""
    typs = TSTensor, torch.as_tensor
    def __init__(self, X=None, y=None, items=None, sel_vars=None, sel_steps=None, tfms=None, tls=None, n_inp=None, dl_type=None,
                 inplace=True, **kwargs):

        # Prepare X (and y)
        if X is not None:
            if not hasattr(X, '__array__'):
                X = np.asarray(X)
            X = to3d(X)
        if y is not None:
            if not hasattr(y, '__array__'):
                y = np.asarray(y)
            elif hasattr(y, "iloc"):
                y = toarray(y)

        # Prepare sel_vars and sel_steps
        self.multi_index = False
        if sel_vars is None or (type(sel_vars) == slice and sel_vars == slice(None)):
            self.sel_vars = slice(None)
        elif type(sel_vars) == slice:
            self.sel_vars = sel_vars
            self.multi_index = True
        else:
            self.sel_vars = np.asarray(sel_vars)
            if sel_steps is not None and type(sel_steps) != slice: self.sel_vars = sel_vars[:, None]
            self.multi_index = True
        if sel_steps is None or (type(sel_steps) == slice and sel_steps == slice(None)):
            self.sel_steps = slice(None)
        elif type(sel_steps) == slice:
            self.sel_steps = sel_steps
            self.multi_index = True
        else:
            self.sel_steps = np.asarray(sel_steps)
            self.multi_index = True
        self.tfms, self.inplace = tfms, inplace

        # Prepare tls (transform lists) and ptls (preprocessed transform lists)
        if tls is None:
            items = tuple((X,)) if y is None else tuple((X, y))
            if tfms is None:
                self.tfms, lts = [None] * len(items), [NoTfmLists] * len(items)
            else:
                self.tfms = _remove_brackets(tfms)
                lts = [NoTfmLists if t is None else TSTfmdLists if getattr(t, 'vectorized', None) else TfmdLists for t in self.tfms]
            self.tls = L(lt(item, t, **kwargs) for lt,item,t in zip(lts, items, self.tfms))
            # if len(self.tls) > 0 and len(self.tls[0]) > 0:
            #     self.typs = [type(tl.items[0]) if isinstance(tl.items[0], torch.Tensor) else self.typs[i] for i,tl in enumerate(self.tls)]
            if self.inplace and (tfms is None or tfms == [None] * len(self.tls)):
                for tl,typ in zip(self.tls, self.typs):
                    tl.items = typ(tl.items)
                self.ptls = self.tls
                self.no_tfm = True
            else:
                self.ptls = L([typ(stack(tl[:]))[...,self.sel_vars, self.sel_steps] if (i==0 and self.multi_index) else typ(stack(tl[:])) \
                    for i,(tl,typ) in enumerate(zip(self.tls,self.typs))]) if inplace else self.tls
                self.no_tfm = False
        else:
            self.tls = tls
            # if len(self.tls) > 0 and len(self.tls[0]) > 0:
            #     self.typs = [type(tl[0]) if isinstance(tl[0], torch.Tensor) else self.typs[i] for i,tl in enumerate(self.tls)]
            self.ptls = L([typ(stack(tl[:]))[...,self.sel_vars, self.sel_steps] if (i==0 and self.multi_index) else typ(stack(tl[:])) \
                for i,(tl,typ) in enumerate(zip(self.tls,self.typs))]) if inplace and len(tls[0]) != 0 else tls
            self.no_tfm = False

        self.n_inp = 1
        if kwargs.get('splits', None) is not None:
            split_idxs = _flatten_list(kwargs.get('splits'))
        else:
            split_idxs = np.arange(len(self), dtype=smallest_dtype(len(self)))
        self.split_idxs = split_idxs

    def __getitem__(self, it):
        if self.inplace or self.no_tfm:
            return tuple([ptl[it] for ptl in self.ptls])
        else:
            return tuple([typ(stack(ptl[it]))[...,self.sel_vars, self.sel_steps] if (i==0 and self.multi_index) else typ(stack(ptl[it])) \
                          for i,(ptl,typ) in enumerate(zip(self.ptls,self.typs))])

    def subset(self, i):
        if is_indexer(i):
            return type(self)(tls=L([tl.subset(i) for tl in self.tls]), inplace=self.inplace, tfms=self.tfms,
                              sel_vars=self.sel_vars, sel_steps=self.sel_steps, splits=None if self.splits is None else self.splits[i],
                              split_idx=i)
        else:
            if self.splits is None:
                splits = None
            else:
                min_dtype = np.min_scalar_type(len(i))
                splits = np.arange(len(i), dtype=min_dtype)
            return type(self)(*self[i], inplace=True, tfms=None,
                              sel_vars=self.sel_vars, sel_steps=self.sel_steps, splits=splits, split_idx=ifnone(self.split_idx, 1))

    def _new(self, X, y=None, **kwargs):
        return type(self)(X, y=y, sel_vars=self.sel_vars, sel_steps=self.sel_steps, tfms=self.tfms, inplace=self.inplace,
                          do_setup=False, **kwargs)

    def new_empty(self): return type(self)(tls=[tl.new_empty() for tl in self.tls], sel_vars=self.sel_vars, sel_steps=self.sel_steps,
                                           n_inp=self.n_inp, inplace=self.inplace)

    def __len__(self): return len(self.tls[0])

    def show_at(self, idx, **kwargs):
        self.show(self[idx], **kwargs)
        plt.show()

    def __repr__(self): return tscoll_repr(self)

# %% ../../nbs/006_data.core.ipynb 51
def add_ds(dsets, X, y=None, inplace=True):
    "Create test datasets from X (and y) using validation transforms of `dsets`"
    items = tuple((X,)) if y is None else tuple((X, y))
    with_labels = False if y is None else True
    if isinstance(dsets, TSDatasets):
        tls = dsets.tls if with_labels else dsets.tls[:dsets.n_inp]
        new_tls = L([tl._new(item, split_idx=1) for tl,item in zip(tls, items)])
        return type(dsets)(tls=new_tls, sel_vars=dsets.sel_vars, sel_steps=dsets.sel_steps, inplace=dsets.inplace)
    elif isinstance(dsets, NumpyDatasets):
        tls = dsets.tls if with_labels else dsets.tls[:dsets.n_inp]
        new_tls = L([tl._new(item, split_idx=1) for tl,item in zip(tls, items)])
        return type(dsets)(tls=new_tls, inplace=dsets.inplace)
    elif isinstance(dsets, Datasets):
        tls = dsets.tls if with_labels else dsets.tls[:dsets.n_inp]
        new_tls = L([tl._new(item, split_idx=1) for tl,item in zip(tls, items)])
        return type(dsets)(tls=new_tls)
    elif isinstance(dsets, TfmdLists):
        new_tl = dsets._new(items, split_idx=1)
        return new_tl
    else:
        raise Exception(f"Expected a `Datasets` or a `TfmdLists` but got {dsets.__class__.__name__}")

@patch
def add_dataset(self:NumpyDatasets, X, y=None, inplace=True):
    return add_ds(self, X, y=y, inplace=inplace)

@patch
def add_test(self:NumpyDatasets, X, y=None, inplace=True):
    return add_ds(self, X, y=y, inplace=inplace)

@patch
def add_unlabeled(self:NumpyDatasets, X, inplace=True):
    return add_ds(self, X, y=None, inplace=inplace)

@patch
def add_dataset(self:TSDatasets, X, y=None, inplace=True):
    return add_ds(self, X, y=y, inplace=inplace)

@patch
def add_test(self:TSDatasets, X, y=None, inplace=True):
    return add_ds(self, X, y=y, inplace=inplace)

@patch
def add_unlabeled(self:TSDatasets, X, inplace=True):
    return add_ds(self, X, y=None, inplace=inplace)

# %% ../../nbs/006_data.core.ipynb 68
@patch
def _one_pass(self:TfmdDL):
    b = self.do_batch([self.do_item(0)])
    if self.device is not None: b = to_device(b, self.device)
    its = self.after_batch(b)
    self._n_inp = 1 if not isinstance(its, (list,tuple)) or len(its)==1 else len(its)-1
    self._types = explode_types(its)

# %% ../../nbs/006_data.core.ipynb 69
_batch_tfms = ('after_item','before_batch','after_batch')

@delegates(TfmdDL.__init__)
class NumpyDataLoader(TfmdDL):
    idxs = None
    do_item = noops # create batch returns indices
    def __init__(self, dataset, bs=64, shuffle=False, drop_last=False, num_workers=0, verbose=False, do_setup=True, vocab=None,
                 sort=False, weights=None, partial_n=None, sampler=None, **kwargs):

        if num_workers is None: num_workers = min(16, defaults.cpus)
        if sampler is not None and shuffle:
            raise ValueError('sampler option is mutually exclusive with shuffle')

        for nm in _batch_tfms:
            if nm == 'after_batch' and kwargs.get('batch_tfms',None) is not None: kwargs[nm] = Pipeline(listify(kwargs.get('batch_tfms')))
            else: kwargs[nm] = Pipeline(kwargs.get(nm,None))
        bs = max(1, min(bs, len(dataset))) # bs cannot be 1
        if is_listy(partial_n): partial_n = partial_n[0]
        if isinstance(partial_n, float): partial_n = int(round(partial_n * len(dataset)))
        if partial_n is not None:
            partial_n = min(partial_n, len(dataset))
            bs = min(bs, partial_n)
        if weights is not None: weights = weights / weights.sum()
        self.weights, self.partial_n, self.sampler, self.sort, self.do_setup = weights, partial_n, sampler, sort, do_setup
        super().__init__(dataset, bs=bs, shuffle=shuffle, drop_last=drop_last, num_workers=num_workers, verbose=verbose, do_setup=do_setup, **kwargs)
        if vocab is not None:
            self.vocab = vocab

    def new_dl(self, X, y=None, bs=64):
        assert X.ndim == 3, "You must pass an X iterable with 3 dimensions [batch_size x n_vars x seq_len]"
        if y is not None:
            y = np.asarray(y)
            assert y.ndim > 0, "You must pass a y iterable with at least 1 dimension"
        ds = self.dataset.add_dataset(X, y=y)
        return self.new(ds, bs=min(bs, len(X)))

    def create_batch(self, b):
        if self.shuffle or self.sampler is not None:
            if self.sort and hasattr(b, 'sort'): b.sort()
            self.idxs = L(b)
        else:
            if self.n is not None:
                b = slice(b[0], min(self.n, b[0] + self.bs))
            else:
                b = slice(b[0], b[0] + self.bs)

            self.idxs = b
        if hasattr(self, "split_idxs"):
            self.input_idxs = self.split_idxs[b]
        else: self.input_idxs = self.idxs
        return self.dataset[b]

    def create_item(self, s):
        if self.indexed: return self.dataset[s or 0]
        elif s is None:  return next(self.it)
        else: raise IndexError("Cannot index an iterable dataset numerically - must use `None`.")

    def get_idxs(self):
        if self.n==0: return []
        if self.partial_n is not None: n = min(self.partial_n, self.n)
        else: n = self.n
        if self.sampler is not None:
            idxs = np.array(list(iter(self.sampler)))
            if self.partial_n is not None: idxs = idxs[:n]
            return idxs
        if self.weights is not None:
            return random_choice(self.n, n, p=self.weights)
        if self.n is not None:
            if self.partial_n is not None:
                return random_choice(self.n, n, False, dtype=smallest_dtype(self.n)) # already shuffled
            else:
                idxs = np.arange(self.n, dtype=smallest_dtype(self.n))
        else:
            idxs = Inf.count if self.indexed else Inf.nones
        if self.shuffle: idxs = self.shuffle_fn(idxs)
        return idxs

    def shuffle_fn(self, idxs):
        return np.random.permutation(idxs)

    def unique_batch(self, max_n=9):
        old_bs = self.bs
        self.bs = 1
        old_get_idxs = self.get_idxs
        self.get_idxs = lambda: Inf.zeros
        out_len = len(self.items)
        types = self.dataset.types
        x, y = [], []
        for _ in range(max_n):
            out = self.one_batch()
            if out_len == 2:
                x.extend(out[0])
                y.extend(out[1])
            else:
                x.extend(out)
        b = (types[0](stack(x)), types[1](stack(y))) if out_len == 2 else (types[0](stack(x)), )
        self.bs = old_bs
        self.get_idxs = old_get_idxs
        return b

    def __len__(self):
        if self.n == 0: return 0
        elif self.partial_n is None: return super().__len__()
        return self.partial_n//self.bs + (0 if self.drop_last or self.partial_n%self.bs==0 else 1)

    @delegates(plt.subplots)
    def show_batch(self, b=None, ctxs=None, max_n=9, nrows=3, ncols=3, figsize=None, unique=False, sharex=True, sharey=False, decode=False,
                   show_title=True, **kwargs):

        old_sort = self.sort
        self.sort = False # disable sorting when showing a batch to ensure varied samples

        if unique:
            b = self.unique_batch(max_n=max_n)
            sharex, sharey = True, True
        elif b is None: b = self.one_batch()
        if not decode:                                        # decode = False allows you to see the data as seen by the model
            after_batch = self.after_batch
            self.after_batch = Pipeline()
            db = self.decode_batch(b, max_n=max_n)
            self.after_batch = after_batch
        else:
            db = self.decode_batch(b, max_n=max_n)
        ncols = min(ncols, math.ceil(len(db) / ncols))
        nrows = min(nrows, math.ceil(len(db) / ncols))
        max_n = min(max_n, len(db), nrows*ncols)
        if figsize is None: figsize = (ncols*6, math.ceil(max_n/ncols)*4)
        if ctxs is None: ctxs = get_grid(max_n, nrows=nrows, ncols=ncols, figsize=figsize, sharex=sharex, sharey=sharey, **kwargs)
        if show_title:
            for i,ctx in enumerate(ctxs):
                show_tuple(db[i], ctx=ctx)
        else:
            db = [x for x,_ in db]
            for i,ctx in enumerate(ctxs):
                db[i].show(ctx=ctx)

        self.sort = old_sort

    @delegates(plt.subplots)
    def show_results(self, b, preds, ctxs=None, max_n=9, nrows=3, ncols=3, figsize=None, **kwargs):
        t = self.decode_batch(b, max_n=max_n)
        p = self.decode_batch((b[0],preds), max_n=max_n)
        if figsize is None: figsize = (ncols*6, max_n//ncols*4)
        if ctxs is None: ctxs = get_grid(min(len(t), nrows*ncols), nrows=None, ncols=ncols, figsize=figsize, **kwargs)
        for i,ctx in enumerate(ctxs):
            title = f'True: {t[i][1]}\nPred: {p[i][1]}'
            color = 'green' if t[i][1] == p[i][1] else 'red'
            t[i][0].show(ctx=ctx, title=title, title_color=color)

    @delegates(plt.subplots)
    def show_dist(self, figsize=None, color=None, **kwargs):
        if self.c == 0:
            print('\nunlabeled dataset.\n')
            return
        b = self.one_batch()
        i = getattr(self, 'n_inp', 1 if len(b)==1 else len(b)-1)
        yb = b[i:][0]
        if color == "random": color = random_shuffle(L(mcolors.CSS4_COLORS.keys()))
        elif color is None: color = ['m', 'orange', 'darkblue', 'lightgray']
        figsize = ifnone(figsize, (8, 6))
        plt.figure(figsize=figsize, **kwargs)
        ax = plt.axes()
        ax.set_axisbelow(True)
        plt.grid(color='gainsboro', linewidth=.1)
        plt.title('Target distribution in a single batch', fontweight='bold')
        if self.cat:
            if yb.ndim == 1:
                yb = yb.flatten().detach().cpu().numpy()
                data = np.unique(yb, return_counts=True)[1]
                data = data / np.sum(data)
            else:
                yb = yb.detach().cpu().numpy()
                data = yb.mean(0)
            plt.bar(self.vocab, data, color=color, edgecolor='black')
            plt.xticks(self.vocab)
        else:
            yb = yb.flatten().detach().cpu().numpy()
            weights=np.ones(len(yb)) / len(yb)
            plt.hist(yb, bins=min(len(yb) // 2, 100), weights=weights, color='violet', edgecolor='black')
        plt.gca().yaxis.set_major_formatter(PercentFormatter(1))
        plt.show()

    @property
    def c(self):
        if len(self.dataset) == 0:
            return 0
        if hasattr(self, "vocab"):
            return len(self.vocab)
        else:
            return 1
#             return self.d if not is_listy(self.d) else reduce(lambda x, y: x * y, self.d, 1)

    @property
    def d(self):
        if len(self.dataset) == 0: return None
        b = self.one_batch()
        if len(b) == 1: return None
        i = getattr(self, 'n_inp', len(b)-1)
        yb = b[i:]
        if len(yb[0].shape[1:]) == 0: return None
        elif len(yb[0].shape[1:]) == 1: return yb[0].shape[-1]
        else: return list(yb[0].shape[1:])

    @property
    def cat(self): return hasattr(self, "vocab")

    @property
    def cws(self):
        if self.cat:
            if self.ptls[-1].ndim == 2: # one hot encoded
                pos_per_class = self.ptls[-1].sum(0)
                neg_per_class = (len(self.ptls[-1]) - pos_per_class)
                pos_weights = neg_per_class / pos_per_class
                pos_weights[pos_weights == float('inf')] = 0
                return torch.Tensor(pos_weights).to(self.device)
            else:
                counts = torch.unique(self.ptls[-1].flatten(), return_counts=True, sorted=True)[-1]
                iw = (counts.sum() / counts)
                return (iw / iw.sum()).to(self.device)
        else: return None

    @property
    def class_priors(self):
        if self.cws is not None:
            cp = 1. / (self.cws + 1e-8)
            return (cp / cp.sum()).to(self.device)
        else: return None


class TSDataLoader(NumpyDataLoader):
    @property
    def vars(self):
        if len(self.dataset) == 0: return 0
        b = self.one_batch()
        i = getattr(self, 'n_inp', 1 if len(b)==1 else len(b)-1)
        xb = b[:i]
        if hasattr(xb[0], 'vars'): return xb[0].vars
        if xb[0].ndim >= 4: return xb[0].shape[-3]
        else: return xb[0].shape[-2]
    @property
    def len(self):
        if len(self.dataset) == 0: return 0
        b = self.one_batch()
        i = getattr(self, 'n_inp', 1 if len(b)==1 else len(b)-1)
        xb = b[:i]
        if hasattr(xb[0], 'len'): return xb[0].len
        if xb[0].ndim >= 4: return xb[0].shape[-2:]
        else: return xb[0].shape[-1]

# %% ../../nbs/006_data.core.ipynb 70
_batch_tfms = ('after_item','before_batch','after_batch')

class NumpyDataLoaders(DataLoaders):
    _xblock = NumpyTensorBlock
    _dl_type = NumpyDataLoader
    def __init__(self, *loaders, path='.', device=None):
        create_dir(path, verbose=False)
        self.loaders, self.path = list(loaders), Path(path)
        self.device = ifnone(device, default_device())

    def new_dl(self, X, y=None, bs=64):
        assert X.ndim == 3, "You must pass an X with 3 dimensions [batch_size x n_vars x seq_len]"
        if y is not None and not is_array(y) and not is_listy(y): y = [y]
        new_dloader = self.new(self.dataset.add_dataset(X, y=y), bs=min(bs, len(X)))
        return new_dloader

    @delegates(plt.subplots)
    def show_dist(self, figsize=None, **kwargs): self.loaders[0].show_dist(figsize=figsize, **kwargs)

    def decoder(self, o):
        if isinstance(o, tuple): return self.decode(o)
        if o.ndim <= 1: return self.decodes(o)
        else: return L([self.decodes(oi) for oi in o])

    def decode(self, b):
        return to_cpu(self.after_batch.decode(self._retain_dl(b)))


    @classmethod
    @delegates(DataLoaders.from_dblock)
    def from_numpy(cls, X, y=None, splitter=None, valid_pct=0.2, seed=0, item_tfms=None, batch_tfms=None, **kwargs):
        "Create timeseries dataloaders from arrays (X and y, unless unlabeled)"
        if splitter is None: splitter = RandomSplitter(valid_pct=valid_pct, seed=seed)
        getters = [ItemGetter(0), ItemGetter(1)] if y is not None else [ItemGetter(0)]
        dblock = DataBlock(blocks=(cls._xblock, CategoryBlock),
                           getters=getters,
                           splitter=splitter,
                           item_tfms=item_tfms,
                           batch_tfms=batch_tfms)

        source = itemify(X) if y is None else itemify(X,y)
        return cls.from_dblock(dblock, source, **kwargs)

    @classmethod
    def from_dsets(cls, *ds, path='.', bs=64, num_workers=0, batch_tfms=None, device=None, shuffle_train=True, drop_last=True,
                   weights=None, partial_n=None, sampler=None, sort=False, vocab=None, **kwargs):
        device = ifnone(device, default_device())
        if batch_tfms is not None and not isinstance(batch_tfms, list): batch_tfms = [batch_tfms]
        shuffle_default = (shuffle_train,) + (False,) * (len(ds)-1)
        drop_last_default = (drop_last,) + (False,) * (len(ds)-1)
        defaults = {'shuffle': shuffle_default, 'drop_last': drop_last_default}
        kwargs = merge(defaults, {k: tuplify(v, match=ds) for k,v in kwargs.items()})
        kwargs = [{k: v[i] for k,v in kwargs.items()} for i in range_of(ds)]
        if not is_listy(bs): bs = [bs]
        if len(bs) != len(ds): bs = bs * len(ds)
        if weights is None: weights = [None] * len(ds)
        if not is_listy(partial_n): partial_n = [partial_n]
        if len(partial_n) != len(ds) and len(ds) > 1: partial_n = partial_n + [None] * (len(ds) - len(partial_n))
        if not is_listy(sampler): sampler = [sampler]
        if len(sampler) != len(ds): sampler = sampler * len(ds)
        loaders = [cls._dl_type(d, bs=b, num_workers=num_workers, batch_tfms=batch_tfms, weights=w, partial_n=n, sampler=s, sort=sort, vocab=vocab, **k)\
                   for d,k,b,w,n,s in zip(ds, kwargs, bs, weights, partial_n, sampler)]
        return cls(*loaders, path=path, device=device)


class TSDataLoaders(NumpyDataLoaders):
    _xblock = TSTensorBlock
    _dl_type = TSDataLoader

# %% ../../nbs/006_data.core.ipynb 71
class StratifiedSampler:
    "Sampler where batches preserve the percentage of samples for each class"

    def __init__(self,
        y, # The target variable for supervised learning problems. Stratification is done based on the y labels.
        bs : int = 64, # Batch size
        shuffle : bool = False, # Flag to shuffle each class’s samples before splitting into batches.
        drop_last : bool = False # Flag to drop the last incomplete batch.
        ):
        self.n_splits = len(y) // bs if drop_last else int(np.ceil(len(y) / bs))
        self.skf = StratifiedKFold(n_splits=self.n_splits, shuffle=shuffle)
        self.y = y
        self.shuffle = shuffle
        self.n = self.n_splits * bs if drop_last else len(y)

    def __iter__(self):
        if self.shuffle:
            self.skf.random_state = np.random.randint(0, 1e8)
        chunklist = []
        for _,idx in self.skf.split(self.y, self.y):
            chunklist.extend(idx)
        yield chunklist

    def __len__(self):
        return self.n

# %% ../../nbs/006_data.core.ipynb 73
def get_c(dls):
    if getattr(dls, 'c', False): return dls.c
    if getattr(getattr(dls.train, 'after_item', None), 'c', False): return dls.train.after_item.c
    if getattr(getattr(dls.train, 'after_batch', None), 'c', False): return dls.train.after_batch.c
    vocab = getattr(dls, 'vocab', [])
    if len(vocab) > 0 and is_listy(vocab[-1]): vocab = vocab[-1]
    return len(vocab)

# %% ../../nbs/006_data.core.ipynb 74
def get_best_dl_params(dl, n_iters=10, num_workers=[0, 1, 2, 4, 8], pin_memory=[True, False], prefetch_factor=[2, 4, 8], return_best=True,
                       verbose=True):

    if not torch.cuda.is_available():
        num_workers = 0
    n_iters = min(n_iters, len(dl))
    if not return_best: verbose = True

    nw = dl.fake_l.num_workers
    pm = dl.fake_l.pin_memory
    pf = dl.fake_l.prefetch_factor

    try:
        best_nw = nw
        best_pm = pm
        best_pf = pf

        # num_workers
        if not num_workers: best_nw = nw
        elif isinstance(num_workers, Integral): best_nw = num_workers
        else:
            best_time = np.inf
            for _nw in num_workers:
                dl.fake_l.num_workers = _nw
                timer.start(False)
                for i, _ in enumerate(dl):
                    if i == n_iters - 1:
                        t = timer.stop() / (i + 1)
                        pv(f'   num_workers: {_nw:2}  pin_memory: {pm!s:^5}  prefetch_factor: {pf:2}  -  time: {1_000 * t/n_iters:8.3f} ms/iter', verbose)
                        if t < best_time:
                            best_nw = _nw
                            best_time = t
                        break
        dl.fake_l.num_workers = best_nw


        # pin_memory
        if not pin_memory: best_pm = pm
        elif isinstance(pin_memory, bool): best_pm = pin_memory
        else:
            best_time = np.inf
            if not pin_memory: pin_memory = [pm]
            for _pm in pin_memory:
                dl.fake_l.pin_memory = _pm
                timer.start(False)
                for i, _ in enumerate(dl):
                    if i == n_iters - 1:
                        t = timer.stop() / (i + 1)
                        pv(f'   num_workers: {best_nw:2}  pin_memory: {_pm!s:^5}  prefetch_factor: {pf:2}  -  time: {1_000 * t/n_iters:8.3f} ms/iter',
                           verbose)
                        if t < best_time:
                            best_pm = _pm
                            best_time = t
                        break
        dl.fake_l.pin_memory = best_pm

        # prefetch_factor
        if best_nw == 0: best_pf = 2
        elif not prefetch_factor: best_pf = pf
        elif isinstance(prefetch_factor, Integral): best_pf = prefetch_factor
        else:
            best_time = np.inf
            if not prefetch_factor: prefetch_factor = [pf]
            for _pf in prefetch_factor:
                dl.fake_l.prefetch_factor = _pf
                timer.start(False)
                for i, _ in enumerate(dl):
                    if i == n_iters - 1:
                        t = timer.stop() / (i + 1)
                        pv(f'   num_workers: {best_nw:2}  pin_memory: {best_pm!s:^5}  prefetch_factor: {_pf:2}  -  time: {1_000 * t/n_iters:8.3f} ms/iter',
                           verbose)
                        if t < best_time:
                            best_pf = _pf
                            best_time = t
                        break
        dl.fake_l.prefetch_factor = best_pf

    except KeyboardInterrupt:
        dl.fake_l.num_workers = best_nw if return_best else nw
        dl.fake_l.pin_memory = best_pm if return_best else pm
        dl.fake_l.prefetch_factor = best_pf if return_best else pf

    if not return_best:
        dl.fake_l.num_workers = nw
        dl.fake_l.pin_memory = pm
        dl.fake_l.prefetch_factor = pf

    if verbose:
        print('\n   best dl params:')
        print(f'       best num_workers    : {best_nw}')
        print(f'       best pin_memory     : {best_pm}')
        print(f'       best prefetch_factor: {best_pf}')
        print(f'       return_best         : {return_best}')
        print('\n')

    return dl

def get_best_dls_params(dls, n_iters=10, num_workers=[0, 1, 2, 4, 8], pin_memory=[True, False], prefetch_factor=[2, 4, 8],
                        return_best=True, verbose=True):

    for i in range(len(dls.loaders)):
        try:
            pv(f'\nDataloader {i}\n', verbose)
            dls.loaders[i] = get_best_dl_params(dls.loaders[i], n_iters=n_iters, num_workers=num_workers, pin_memory=pin_memory,
                                            prefetch_factor=prefetch_factor, return_best=return_best, verbose=verbose)
        except KeyboardInterrupt: pass
    return dls

# %% ../../nbs/006_data.core.ipynb 75
def _check_splits(X, splits):
    if splits is None:
        _dtype = smallest_dtype(len(X))
        if len(X) < 1e6:
            splits = (L(np.arange(len(X), dtype=_dtype).tolist()), L())
        else:
            _dtype = smallest_dtype(len(X))
            splits = (np.arange(len(X), dtype=_dtype), L())
    elif isinstance(splits, (tuple, list, L, np.ndarray)):
        if not isinstance(splits[0], (tuple, list, L, np.ndarray)):
            splits = (splits, L())
        elif len(splits) == 1:
            splits = (splits[0], L())
        elif len(splits) >= 2 and splits[1] is None:
            splits = (splits[0], L())
    assert len(splits) >= 2, 'splits must be a tuple or list of length >=2'
    assert len(splits[0]) > 0, 'splits[0] must be a non-empty list'
    return splits

def get_ts_dls(X, y=None, splits=None, sel_vars=None, sel_steps=None, tfms=None, inplace=True,
               path='.', bs=64, batch_tfms=None, num_workers=0, device=None, shuffle_train=True, drop_last=True,
               weights=None, partial_n=None, sampler=None, sort=False, **kwargs):
    splits = _check_splits(X, splits)
    create_dir(path, verbose=False)
    dsets = TSDatasets(X, y, splits=splits, sel_vars=sel_vars, sel_steps=sel_steps, tfms=tfms, inplace=inplace)
    dsets = [dsets.subset(i) for i in range(len(splits))]
    if weights is not None:
        assert len(X) == len(weights), 'len(X) != len(weights)'
        weights = [weights[split] if i == 0 else None for i,split in enumerate(splits)] # weights only applied to train set
    dls   = TSDataLoaders.from_dsets(*dsets, path=path, bs=bs, batch_tfms=batch_tfms, num_workers=num_workers,
                                     device=device, shuffle_train=shuffle_train, drop_last=drop_last, weights=weights,
                                     partial_n=partial_n, sampler=sampler, sort=sort, **kwargs)
    return dls

get_tsimage_dls = get_ts_dls

# %% ../../nbs/006_data.core.ipynb 77
def _check_split(X, split):
    if split is None:
        _dtype = smallest_dtype(len(X))
        if len(X) < 1e6:
            split = L(np.arange(len(X), dtype=_dtype).tolist())
        else:
            _dtype = smallest_dtype(len(X))
            split = np.arange(len(X), dtype=_dtype)
    return (split, L())

def get_ts_dl(X, y=None, split=None, sel_vars=None, sel_steps=None, tfms=None, inplace=True,
              path='.', bs=64, batch_tfms=None, num_workers=0, device=None, shuffle_train=True, drop_last=True, weights=None,
              partial_n=None, sampler=None, sort=False, **kwargs):
    splits = _check_split(X, split)
    create_dir(path, verbose=False)
    dsets = TSDatasets(X, y, splits=splits, sel_vars=sel_vars, sel_steps=sel_steps, tfms=tfms, inplace=inplace, **kwargs)
    if not is_listy(partial_n): partial_n = [partial_n]
    dls   = TSDataLoaders.from_dsets(dsets.train, path=path, bs=bs, batch_tfms=batch_tfms, num_workers=num_workers,
                                     device=device, shuffle_train=shuffle_train, drop_last=drop_last, weights=weights,
                                     partial_n=partial_n, sampler=sampler, sort=sort, **kwargs)
    return dls.train



def get_subset_dl(dl, idxs): return dl.new(dl.dataset.subset(idxs))

# %% ../../nbs/006_data.core.ipynb 116
def get_time_per_batch(dl, model=None, n_batches=None):
    try:
        timer.start(False)
        pbar = progress_bar(dl, leave=False)
        for i, (xb, _) in enumerate(pbar):
            if model is not None:
                _ = model(xb)
            if n_batches is not None and i >= n_batches - 1:
                t = timer.stop()
                pbar.on_interrupt()
                break
        if n_batches is None or i < n_batches - 1:
            t = timer.stop()

    except KeyboardInterrupt:
        t = timer.stop()
        pbar.on_interrupt()
    return t / (i+1)

def get_dl_percent_per_epoch(dl, model, n_batches=None):
    dl_time = get_time_per_batch(dl, model=None, n_batches=n_batches)
    model_time = get_time_per_batch(dl, model=model, n_batches=n_batches)
    return f'{min(1, dl_time/model_time):.2%}'
