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

__all__ = ['TSBlock', 'stack_train_valid', 'TSDataLoaders']

# Cell
from .imports import *
from .core import *
from fastai.basics import *
from fastai.torch_core import *
from fastai.vision.data import get_grid

# Cell
def TSBlock(cls=TSeries):
    "A TimeSeries Block to process one timeseries"
    return TransformBlock(type_tfms=cls.create)

# Cell
def stack_train_valid(df_train, df_valid):
    "Stack df_train and df_valid, adds `valid_col`=True/False for df_valid/df_train"
    return pd.concat([df_train.assign(valid_col=False), df_valid.assign(valid_col=True)]).reset_index(drop=True)

# Cell
class TSDataLoaders(DataLoaders):
    "A TimeSeries DataLoader"
    @classmethod
    @delegates(DataLoaders.from_dblock)
    def from_df(cls, df, path='.', valid_pct=0.2, seed=None, x_cols=None, label_col=None,
                y_block=None, valid_col=None, item_tfms=None, batch_tfms=None, **kwargs):
        "Create a DataLoader from a pandas DataFrame"
        y_block = ifnone(y_block, CategoryBlock)
        splitter = RandomSplitter(valid_pct, seed=seed) if valid_col is None else ColSplitter(valid_col)
        dblock = DataBlock(blocks=(TSBlock, y_block),
                           get_x=lambda o: o[x_cols].values.astype(np.float32),
                           get_y=ColReader(label_col),
                           splitter=splitter,
                           item_tfms=item_tfms,
                           batch_tfms=batch_tfms)
        return cls.from_dblock(dblock, df, path=path, **kwargs)
    @classmethod
    @delegates(DataLoaders.from_dblock)
    def from_dfs(cls, df_train, df_valid, path='.', x_cols=None, label_col=None,
                y_block=None, item_tfms=None, batch_tfms=None, **kwargs):
        "Create a DataLoader from a df_train and df_valid"
        df = stack_train_valid(df_train, df_valid)
        return cls.from_df(df, path, x_cols=x_cols, valid_col='valid_col', label_col=label_col,
                y_block=y_block, item_tfms=item_tfms, batch_tfms=batch_tfms,**kwargs)

# Cell
@typedispatch
def show_batch(x: TSeries, y, samples, ctxs=None, max_n=10,rows=None, cols=None, figsize=None, **kwargs):
    "Show batch for TSeries objects"
    if ctxs is None: ctxs = get_grid(min(len(samples), max_n), nrows=rows, ncols=cols, add_vert=1, figsize=figsize)
    ctxs = show_batch[object](x, y, samples=samples, ctxs=ctxs, max_n=max_n, **kwargs)
    return ctxs