# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/00_data/03_block.ipynb.

# %% auto 0
__all__ = ['pad_sequence', 'SequenceBlock', 'ScalarNormalize', 'ScalarBlock']

# %% ../../nbs/00_data/03_block.ipynb 2
from fastai.data.all import *
from .core import *
from .split import ParentSplitter, ApplyToDict, PercentageSplitter
from .transforms import SeqNoiseInjection, Normalize,SeqSlice

# %% ../../nbs/00_data/03_block.ipynb 4
def pad_sequence(batch,sorting = False):
    '''collate_fn for padding of sequences of different lengths, use in before_batch of databunch, still quite slow'''
    #takes list of tuples as input, returns list of tuples
    sorted_batch = sorted(batch, key=lambda x: x[0].shape[0], reverse=True) if sorting else batch

    pad_func = partial(torch.nn.utils.rnn.pad_sequence,batch_first=True)
    padded_tensors = [pad_func([x[tup] for x in sorted_batch]) for tup in range(len(batch[0]))]
    padded_list = [retain_types(tuple([tup[entry] for tup in padded_tensors]),batch[0]) for entry in range(len(batch))]
    #retain types is important for decoding later back to source items
#     import pdb; pdb.set_trace()
    
    return padded_list

# %% ../../nbs/00_data/03_block.ipynb 12
class SequenceBlock(TransformBlock):
    def __init__(self, seq_extract,padding=False):
        return super().__init__(type_tfms=[seq_extract],
                                batch_tfms=[Normalize(axes=[0,1])],
                                dls_kwargs={} if not padding else {'before_batch': pad_sequence})

    @classmethod
    @delegates(HDF2Sequence, keep=True)
    def from_hdf(cls, clm_names, seq_cls=TensorSequencesInput,padding=False, **kwargs):
        return cls(HDF2Sequence(clm_names,to_cls=seq_cls,**kwargs), padding)

    @classmethod
    def from_numpy(cls, seq_cls=TensorSequencesInput,padding=False, **kwargs):
        return cls(ToTensor(enc=seq_cls), padding)

# %% ../../nbs/00_data/03_block.ipynb 15
class ScalarNormalize(DisplayedTransform):
    def __init__(self, mean=None, std=None, axes=(0,)): store_attr()
        
    @classmethod
    def from_stats(cls, mean, std, dim=1, ndim=4, cuda=True): return cls(*broadcast_vec(dim, ndim, mean, std, cuda=cuda))
    
    def setups(self, dl:DataLoader):
        if self.mean is None or self.std is None:
            b = dl.one_batch()
            for x in b:
                if isinstance(x,TensorScalarsInput):
                    self.mean,self.std = x.mean(self.axes, keepdim=True),x.std(self.axes, keepdim=True)+1e-7
                    return

    def encodes(self, x:TensorScalarsInput): 
        if x.device != self.mean.device:
            self.mean = self.mean.to(x.device)
            self.std = self.std.to(x.device)
        return (x-self.mean) / self.std
    
    def decodes(self, x:TensorScalarsInput):
        if x.device != self.mean.device:
            self.mean = self.mean.to(x.device)
            self.std = self.std.to(x.device)
        return (x*self.std + self.mean)

class ScalarBlock(TransformBlock):
    def __init__(self, scl_extract):
        return super().__init__(type_tfms=[scl_extract],
                                batch_tfms=[ScalarNormalize()])

    @classmethod
    @delegates(HDF2Scalars, keep=True)
    def from_hdf(cls, clm_names, scl_cls=TensorScalarsInput, **kwargs):
        return cls(HDF2Scalars(clm_names,to_cls=scl_cls,**kwargs))
