# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/017_data.transforms.ipynb.

# %% auto 0
__all__ = ['TSMagScaleByVar', 'TSRandomZoomIn', 'TSSubsampleSteps', 'TSRandomRotate', 'all_TS_randaugs', 'TSIdentity',
           'TSShuffle_HLs', 'TSShuffleSteps', 'TSGaussianNoise', 'TSMagAddNoise', 'TSMagMulNoise',
           'random_curve_generator', 'random_cum_curve_generator', 'random_cum_noise_generator',
           'random_cum_linear_generator', 'TSTimeNoise', 'TSMagWarp', 'TSTimeWarp', 'TSWindowWarp', 'TSMagScale',
           'TSMagScalePerVar', 'TSRandomResizedCrop', 'TSWindowSlicing', 'TSRandomZoomOut', 'TSRandomTimeScale',
           'TSRandomTimeStep', 'TSResampleSteps', 'TSBlur', 'TSSmooth', 'maddest', 'TSFreqDenoise', 'TSRandomFreqNoise',
           'TSRandomResizedLookBack', 'TSRandomLookBackOut', 'TSVarOut', 'TSCutOut', 'TSTimeStepOut', 'TSRandomCropPad',
           'TSMaskOut', 'TSInputDropout', 'TSTranslateX', 'TSRandomShift', 'TSHorizontalFlip', 'TSRandomTrend',
           'TSVerticalFlip', 'TSResize', 'TSRandomSize', 'TSRandomLowRes', 'TSDownUpScale', 'TSRandomDownUpScale',
           'TSRandomConv', 'TSRandom2Value', 'TSMask2Value', 'RandAugment', 'TestTfm', 'get_tfm_name']

# %% ../../nbs/017_data.transforms.ipynb 3
from ..imports import *
from scipy.interpolate import CubicSpline
from scipy.ndimage import convolve1d
from fastcore.transform import compose_tfms
from fastai.vision.augment import RandTransform
from ..utils import *
from .core import *

# %% ../../nbs/017_data.transforms.ipynb 6
class TSIdentity(RandTransform):
    "Applies the identity tfm to a `TSTensor` batch"
    order = 90
    def __init__(self, magnitude=None, **kwargs): 
        self.magnitude = magnitude 
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor): return o

# %% ../../nbs/017_data.transforms.ipynb 8
# partial(TSShuffle_HLs, ex=0), 
class TSShuffle_HLs(RandTransform):
    "Randomly shuffles HIs/LOs of an OHLC `TSTensor` batch"
    order = 90
    def __init__(self, magnitude=1., ex=None, **kwargs): 
        self.magnitude, self.ex = magnitude, ex
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        timesteps = o.shape[-1] // 4
        pos_rand_list = np.random.choice(np.arange(timesteps),size=random.randint(1, timesteps),replace=False)
        rand_list = pos_rand_list * 4
        highs = rand_list + 1
        lows = highs + 1
        a = np.vstack([highs, lows]).flatten('F')
        b = np.vstack([lows, highs]).flatten('F')
        output = o.clone()
        output[...,a] = output[...,b]
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 10
# partial(TSShuffleSteps, ex=0), 
class TSShuffleSteps(RandTransform):
    "Randomly shuffles consecutive sequence datapoints in batch"
    order = 90
    def __init__(self, magnitude=1., ex=None, **kwargs): 
        self.magnitude, self.ex = magnitude, ex
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        odd = 1 - o.shape[-1]%2
        r = np.random.randint(2)
        timesteps = o.shape[-1] // 2
        pos_rand_list = np.random.choice(np.arange(0, timesteps - r * odd), size=random.randint(1, timesteps - r * odd),replace=False) * 2 + 1 + r
        a = np.vstack([pos_rand_list, pos_rand_list - 1]).flatten('F')
        b = np.vstack([pos_rand_list - 1, pos_rand_list]).flatten('F')
        output = o.clone()
        output[...,a] = output[...,b]
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 12
class TSGaussianNoise(RandTransform):
    "Applies additive or multiplicative gaussian noise"
    order = 90
    def __init__(self, magnitude=.5, additive=True, ex=None, **kwargs):
        self.magnitude, self.additive, self.ex = magnitude, additive, ex
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if self.magnitude <= 0: return o
        noise = self.magnitude * torch.randn_like(o)
        if self.ex is None:
            if self.additive: return o + noise
            else: return o * (1 + noise)
        else:
            if self.additive: output = o + noise
            else: output = o * (1 + noise)
        output[..., self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 14
class TSMagAddNoise(RandTransform):
    "Applies additive noise on the y-axis for each step of a `TSTensor` batch"
    order = 90
    def __init__(self, magnitude=1, ex=None, **kwargs):
        self.magnitude, self.ex = magnitude, ex
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        # output = o + torch.normal(0, o.std() * self.magnitude, o.shape, dtype=o.dtype, device=o.device)
        output = o + torch.normal(0, 1/3, o.shape, dtype=o.dtype, device=o.device) * (o[..., 1:] - o[..., :-1]).std(2, keepdims=True) * self.magnitude
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

class TSMagMulNoise(RandTransform): 
    "Applies multiplicative noise on the y-axis for each step of a `TSTensor` batch"
    order = 90
    def __init__(self, magnitude=1, ex=None, **kwargs): 
        self.magnitude, self.ex = magnitude, ex
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        noise = torch.normal(1, self.magnitude * .025, o.shape, dtype=o.dtype, device=o.device)
        output = o * noise
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 16
def random_curve_generator(o, magnitude=0.1, order=4, noise=None):
    seq_len = o.shape[-1]
    f = CubicSpline(np.linspace(-seq_len, 2 * seq_len - 1, 3 * (order - 1) + 1, dtype=int), 
                    np.random.normal(loc=1.0, scale=magnitude, size=3 * (order - 1) + 1), axis=-1)
    return f(np.arange(seq_len))

def random_cum_curve_generator(o, magnitude=0.1, order=4, noise=None):
    x = random_curve_generator(o, magnitude=magnitude, order=order, noise=noise).cumsum()
    x -= x[0]
    x /= x[-1]
    x = np.clip(x, 0, 1)
    return x * (o.shape[-1] - 1)

def random_cum_noise_generator(o, magnitude=0.1, noise=None):
    seq_len = o.shape[-1]
    x = np.clip(np.ones(seq_len) + np.random.normal(loc=0, scale=magnitude, size=seq_len), 0, 1000).cumsum()
    x -= x[0]
    x /= x[-1]
    return x * (o.shape[-1] - 1)

def random_cum_linear_generator(o, magnitude=0.1):
    seq_len = o.shape[-1]
    win_len = int(round(seq_len * np.random.rand() * magnitude))
    if win_len == seq_len: return np.arange(o.shape[-1])
    start = np.random.randint(0, seq_len - win_len)
    # mult between .5 and 2
    rand = np.random.rand()
    mult = 1 + rand
    if np.random.randint(2): mult = 1 - rand/2
    x = np.ones(seq_len)
    x[start : start + win_len] = mult
    x = x.cumsum()
    x -= x[0]
    x /= x[-1]
    return np.clip(x, 0, 1) * (seq_len - 1)

# %% ../../nbs/017_data.transforms.ipynb 17
class TSTimeNoise(RandTransform):
    "Applies noise to each step in the x-axis of a `TSTensor` batch based on smooth random curve"
    order = 90
    def __init__(self, magnitude=0.1, ex=None, **kwargs): 
        self.magnitude, self.ex = magnitude, ex
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        f = CubicSpline(np.arange(o.shape[-1]), o.cpu(), axis=-1)
        output = o.new(f(random_cum_noise_generator(o, magnitude=self.magnitude)))
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 19
class TSMagWarp(RandTransform):
    "Applies warping to the y-axis of a `TSTensor` batch based on a smooth random curve"
    order = 90
    def __init__(self, magnitude=0.02, ord=4, ex=None, **kwargs): 
        self.magnitude, self.ord, self.ex = magnitude, ord, ex
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if self.magnitude and self.magnitude <= 0: return o
        y_mult = random_curve_generator(o, magnitude=self.magnitude, order=self.ord)
        output = o * o.new(y_mult)
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 21
class TSTimeWarp(RandTransform):
    "Applies time warping to the x-axis of a `TSTensor` batch based on a smooth random curve"
    order = 90
    def __init__(self, magnitude=0.1, ord=6, ex=None, **kwargs):
        self.magnitude, self.ord, self.ex = magnitude, ord, ex
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        f = CubicSpline(np.arange(o.shape[-1]), o.cpu(), axis=-1)
        output = o.new(f(random_cum_curve_generator(o, magnitude=self.magnitude, order=self.ord)))
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 23
class TSWindowWarp(RandTransform):
    """Applies window slicing to the x-axis of a `TSTensor` batch based on a random linear curve based on
    https://halshs.archives-ouvertes.fr/halshs-01357973/document"""
    order = 90
    def __init__(self, magnitude=0.1, ex=None, **kwargs): 
        self.magnitude, self.ex = magnitude, ex
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0 or self.magnitude >= 1: return o
        f = CubicSpline(np.arange(o.shape[-1]), o.cpu(), axis=-1)
        output = o.new(f(random_cum_linear_generator(o, magnitude=self.magnitude)))
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 25
class TSMagScale(RandTransform):
    "Applies scaling to the y-axis of a `TSTensor` batch based on a scalar"
    order = 90
    def __init__(self, magnitude=0.5, ex=None, **kwargs): 
        self.magnitude, self.ex = magnitude, ex
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        rand = random_half_normal()
        scale = (1 - (rand  * self.magnitude)/2) if random.random() > 1/3 else (1 + (rand  * self.magnitude))
        output = o * scale
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output
    
class TSMagScalePerVar(RandTransform):
    "Applies per_var scaling to the y-axis of a `TSTensor` batch based on a scalar"
    order = 90
    def __init__(self, magnitude=0.5, ex=None, **kwargs): 
        self.magnitude, self.ex = magnitude, ex
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        s = [1] * o.ndim
        s[-2] = o.shape[-2]
        rand = random_half_normal_tensor(s, device=o.device)
        scale = (1 - (rand  * self.magnitude)/2) if random.random() > 1/3 else (1 + (rand  * self.magnitude))
        output = o * scale
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output
    
TSMagScaleByVar = TSMagScalePerVar

# %% ../../nbs/017_data.transforms.ipynb 27
class TSRandomResizedCrop(RandTransform):
    "Randomly amplifies a sequence focusing on a random section of the steps"
    order = 90
    def __init__(self, magnitude=0.1, size=None, scale=None, ex=None, mode='linear', **kwargs): 
        """
        Args:
            size: None, int or float
            scale: None or tuple of 2 floats 0 < float <= 1
            mode:  'nearest' | 'linear' | 'area'
        
        """
        self.magnitude, self.ex, self.mode = magnitude, ex, mode
        if scale is not None: 
            assert is_listy(scale) and len(scale) == 2 and min(scale) > 0 and min(scale) <= 1, "scale must be a tuple with 2 floats 0 < float <= 1"
        self.size,self.scale = size,scale
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        seq_len = o.shape[-1]
        if self.size is not None: 
            size = self.size if isinstance(self.size, Integral) else int(round(self.size * seq_len))
        else:
            size = seq_len
        if self.scale is not None: 
            lambd = np.random.uniform(self.scale[0], self.scale[1])
        else: 
            lambd = np.random.beta(self.magnitude, self.magnitude)
            lambd = max(lambd, 1 - lambd)
        win_len = int(round(seq_len * lambd))
        if win_len == seq_len: 
            if size == seq_len: return o
            _slice = slice(None) 
        else:
            start = np.random.randint(0, seq_len - win_len)
            _slice = slice(start, start + win_len)
        return F.interpolate(o[..., _slice], size=size, mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)
    
TSRandomZoomIn = TSRandomResizedCrop

# %% ../../nbs/017_data.transforms.ipynb 29
class TSWindowSlicing(RandTransform):
    "Randomly extracts an resize a ts slice based on https://halshs.archives-ouvertes.fr/halshs-01357973/document"
    order = 90
    def __init__(self, magnitude=0.1, ex=None, mode='linear', **kwargs): 
        "mode:  'nearest' | 'linear' | 'area'"
        self.magnitude, self.ex, self.mode = magnitude, ex, mode
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0 or self.magnitude >= 1: return o
        seq_len = o.shape[-1]
        win_len = int(round(seq_len * (1 - self.magnitude)))
        if win_len == seq_len: return o
        start = np.random.randint(0, seq_len - win_len)
        return F.interpolate(o[..., start : start + win_len], size=seq_len, mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)

# %% ../../nbs/017_data.transforms.ipynb 31
class TSRandomZoomOut(RandTransform):
    "Randomly compresses a sequence on the x-axis"
    order = 90
    def __init__(self, magnitude=0.1, ex=None, mode='linear', **kwargs): 
        "mode:  'nearest' | 'linear' | 'area'"
        self.magnitude, self.ex, self.mode = magnitude, ex, mode
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        seq_len = o.shape[-1]
        lambd = np.random.beta(self.magnitude, self.magnitude)
        lambd = max(lambd, 1 - lambd)
        win_len = int(round(seq_len * lambd))
        if win_len == seq_len: return o
        start = (seq_len - win_len) // 2
        output = torch.zeros_like(o, dtype=o.dtype, device=o.device)
        interp = F.interpolate(o, size=win_len, mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)
        output[..., start:start + win_len] = o.new(interp)
        return output

# %% ../../nbs/017_data.transforms.ipynb 33
class TSRandomTimeScale(RandTransform):
    "Randomly amplifies/ compresses a sequence on the x-axis keeping the same length"
    order = 90
    def __init__(self, magnitude=0.1, ex=None, mode='linear', **kwargs): 
        "mode:  'nearest' | 'linear' | 'area'"
        self.magnitude, self.ex, self.mode = magnitude, ex, mode
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        if np.random.rand() <= 0.5: return TSRandomZoomIn(magnitude=self.magnitude, ex=self.ex, mode=self.mode)(o, split_idx=0)
        else: return TSRandomZoomOut(magnitude=self.magnitude, ex=self.ex, mode=self.mode)(o, split_idx=0)

# %% ../../nbs/017_data.transforms.ipynb 35
class TSRandomTimeStep(RandTransform):
    "Compresses a sequence on the x-axis by randomly selecting sequence steps and interpolating to previous size"
    order = 90
    def __init__(self, magnitude=0.02, ex=None, mode='linear', **kwargs): 
        "mode:  'nearest' | 'linear' | 'area'"
        self.magnitude, self.ex, self.mode = magnitude, ex, mode
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        seq_len = o.shape[-1]
        new_seq_len = int(round(seq_len * max(.5, (1 - np.random.rand() * self.magnitude))))
        if  new_seq_len == seq_len: return o
        timesteps = np.sort(np.random.choice(np.arange(seq_len),new_seq_len, replace=False))
        output = F.interpolate(o[..., timesteps], size=seq_len, mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 37
class TSResampleSteps(RandTransform):
    "Transform that randomly selects and sorts sequence steps (with replacement) maintaining the sequence length"

    order = 90
    def __init__(self, step_pct=1., same_seq_len=True, magnitude=None, **kwargs):
        assert step_pct > 0, 'seq_len_pct must be subsample > 0'
        self.step_pct, self.same_seq_len = step_pct, same_seq_len
        super().__init__(**kwargs)

    def encodes(self, o: TSTensor):
        S = o.shape[-1]
        if isinstance(self.step_pct, tuple):
            step_pct = np.random.rand() * (self.step_pct[1] - self.step_pct[0]) + self.step_pct[0]
        else: 
            step_pct = self.step_pct
        if step_pct != 1 and self.same_seq_len:
            idxs = np.sort(np.tile(np.random.choice(S, round(S * step_pct), True), math.ceil(1 / step_pct))[:S])
        else:
            idxs = np.sort(np.random.choice(S, round(S * step_pct), True))
        return o[..., idxs]
    
TSSubsampleSteps = TSResampleSteps

# %% ../../nbs/017_data.transforms.ipynb 39
class TSBlur(RandTransform):
    "Blurs a sequence applying a filter of type [1, 0, 1]"
    order = 90
    def __init__(self, magnitude=1., ex=None, filt_len=None, **kwargs): 
        self.magnitude, self.ex = magnitude, ex
        if filt_len is None: 
            filterargs = [1, 0, 1]  
        else: 
            filterargs = ([1] * max(1, filt_len // 2) + [0] + [1] * max(1, filt_len // 2))
        self.filterargs = np.array(filterargs) 
        self.filterargs = self.filterargs/self.filterargs.sum()
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        output = o.new(convolve1d(o.cpu(), self.filterargs, mode='nearest'))
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 41
class TSSmooth(RandTransform):
    "Smoothens a sequence applying a filter of type [1, 5, 1]"
    order = 90
    def __init__(self, magnitude=1., ex=None, filt_len=None, **kwargs): 
        self.magnitude, self.ex = magnitude, ex
        self.filterargs = np.array([1, 5, 1])
        if filt_len is None: 
            filterargs = [1, 5, 1]  
        else: 
            filterargs = ([1] * max(1, filt_len // 2) + [5] + [1] * max(1, filt_len // 2))
        self.filterargs = np.array(filterargs) 
        self.filterargs = self.filterargs/self.filterargs.sum()
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        output = o.new(convolve1d(o.cpu(), self.filterargs, mode='nearest'))
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 43
def maddest(d, axis=None): 
    #Mean Absolute Deviation
    return np.mean(np.absolute(d - np.mean(d, axis=axis)), axis=axis)

class TSFreqDenoise(RandTransform):
    "Denoises a sequence applying a wavelet decomposition method"
    order = 90
    def __init__(self, magnitude=0.1, ex=None, wavelet='db4', level=2, thr=None, thr_mode='hard', pad_mode='per', **kwargs): 
        self.magnitude, self.ex = magnitude, ex
        self.wavelet, self.level, self.thr, self.thr_mode, self.pad_mode = wavelet, level, thr, thr_mode, pad_mode
        super().__init__(**kwargs)
        try: 
            import pywt
        except ImportError: 
            raise ImportError('You need to install pywt to run TSFreqDenoise')
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        """
        1. Adapted from waveletSmooth function found here:
        http://connor-johnson.com/2016/01/24/using-pywavelets-to-remove-high-frequency-noise/
        2. Threshold equation and using hard mode in threshold as mentioned
        in section '3.2 denoising based on optimized singular values' from paper by Tomas Vantuch:
        http://dspace.vsb.cz/bitstream/handle/10084/133114/VAN431_FEI_P1807_1801V001_2018.pdf
        """
        seq_len = o.shape[-1]
        # Decompose to get the wavelet coefficients
        coeff = pywt.wavedec(o.cpu(), self.wavelet, mode=self.pad_mode)
        # Calculate sigma for threshold as defined in http://dspace.vsb.cz/bitstream/handle/10084/133114/VAN431_FEI_P1807_1801V001_2018.pdf
        # As noted by @harshit92 MAD referred to in the paper is Mean Absolute Deviation not Median Absolute Deviation
        sigma = (1/0.6745) * maddest(coeff[-self.level])
        # Calculate the univeral threshold
        uthr = sigma * np.sqrt(2*np.log(seq_len)) * (1 if self.thr is None else self.magnitude)
        coeff[1:] = (pywt.threshold(c, value=uthr, mode=self.thr_mode) for c in coeff[1:])
        # Reconstruct the signal using the thresholded coefficients
        output = o.new(pywt.waverec(coeff, self.wavelet, mode=self.pad_mode)[..., :seq_len])
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 46
class TSRandomFreqNoise(RandTransform):
    "Applys random noise using a wavelet decomposition method"
    order = 90
    def __init__(self, magnitude=0.1, ex=None, wavelet='db4', level=2, mode='constant', **kwargs): 
        self.magnitude, self.ex = magnitude, ex
        self.wavelet, self.level, self.mode = wavelet, level, mode
        super().__init__(**kwargs)
        try: 
            import pywt
        except ImportError: 
            raise ImportError('You need to install pywt to run TSRandomFreqNoise')
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        self.level = 1 if self.level is None else self.level
        coeff = pywt.wavedec(o.cpu(), self.wavelet, mode=self.mode, level=self.level)
        coeff[1:] = [c * (1 + 2 * (np.random.rand() - 0.5) * self.magnitude) for c in coeff[1:]]
        output = o.new(pywt.waverec(coeff, self.wavelet, mode=self.mode)[..., :o.shape[-1]])
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 48
class TSRandomResizedLookBack(RandTransform):
    "Selects a random number of sequence steps starting from the end and return an output of the same shape"
    order = 90
    def __init__(self, magnitude=0.1, mode='linear', **kwargs): 
        "mode:  'nearest' | 'linear' | 'area'"
        self.magnitude, self.mode = magnitude, mode
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        seq_len = o.shape[-1]
        lambd = np.random.beta(self.magnitude, self.magnitude)
        lambd = min(lambd, 1 - lambd)
        output = o.clone()[..., int(round(lambd * seq_len)):]
        return F.interpolate(output, size=seq_len, mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)

# %% ../../nbs/017_data.transforms.ipynb 50
class TSRandomLookBackOut(RandTransform):
    "Selects a random number of sequence steps starting from the end and set them to zero"
    order = 90
    def __init__(self, magnitude=0.1, **kwargs): 
        self.magnitude = magnitude
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        seq_len = o.shape[-1]
        lambd = np.random.beta(self.magnitude, self.magnitude)
        lambd = min(lambd, 1 - lambd)
        output = o.clone()
        output[..., :int(round(lambd * seq_len))] = 0 
        return output

# %% ../../nbs/017_data.transforms.ipynb 52
class TSVarOut(RandTransform):
    "Set the value of a random number of variables to zero"
    order = 90
    def __init__(self, magnitude=0.05, ex=None, **kwargs): 
        self.magnitude, self.ex = magnitude, ex
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        in_vars = o.shape[-2]
        if in_vars == 1: return o
        lambd = np.random.beta(self.magnitude, self.magnitude)
        lambd = min(lambd, 1 - lambd)
        p = np.arange(in_vars).cumsum()
        p = p/p[-1]
        p = p / p.sum()
        p = p[::-1]
        out_vars = np.random.choice(np.arange(in_vars), int(round(lambd * in_vars)), p=p, replace=False)
        if len(out_vars) == 0:  return o
        output = o.clone()
        output[...,out_vars,:] = 0
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 54
class TSCutOut(RandTransform):
    "Sets a random section of the sequence to zero"
    order = 90
    def __init__(self, magnitude=0.05, ex=None, **kwargs): 
        self.magnitude, self.ex = magnitude, ex
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        seq_len = o.shape[-1]
        lambd = np.random.beta(self.magnitude, self.magnitude)
        lambd = min(lambd, 1 - lambd)
        win_len = int(round(seq_len * lambd))
        start = np.random.randint(-win_len + 1, seq_len)
        end = start + win_len
        start = max(0, start)
        end = min(end, seq_len)
        output = o.clone()
        output[..., start:end] = 0
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 56
class TSTimeStepOut(RandTransform):
    "Sets random sequence steps to zero"
    order = 90
    def __init__(self, magnitude=0.05, ex=None, **kwargs): 
        self.magnitude, self.ex = magnitude, ex
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        magnitude = min(.5, self.magnitude)
        seq_len = o.shape[-1]
        timesteps = np.sort(np.random.choice(np.arange(seq_len), int(round(seq_len * magnitude)), replace=False))
        output = o.clone()
        output[..., timesteps] = 0
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 58
class TSRandomCropPad(RandTransform):
    "Crops a section of the sequence of a random length"
    order = 90
    def __init__(self, magnitude=0.05, ex=None, **kwargs): 
        self.magnitude, self.ex = magnitude, ex
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        seq_len = o.shape[-1]
        lambd = np.random.beta(self.magnitude, self.magnitude)
        lambd = max(lambd, 1 - lambd)
        win_len = int(round(seq_len * lambd))
        if win_len == seq_len: return o
        start = np.random.randint(0, seq_len - win_len)
        output = torch.zeros_like(o, dtype=o.dtype, device=o.device)
        output[..., start : start + win_len] = o[..., start : start + win_len]
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 60
class TSMaskOut(RandTransform):
    """Applies a random mask"""
    order = 90
    def __init__(self, magnitude=0.1, compensate:bool=False, ex=None, **kwargs):
        store_attr()
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        mask = torch.rand_like(o) > (1 - self.magnitude)
        if self.compensate: # per sample and feature
            mean_per_seq = (torch.max(torch.ones(1, device=mask.device), torch.sum(mask, dim=-1).unsqueeze(-1)) / mask.shape[-1])
            output = o.masked_fill(mask, 0) / (1 - mean_per_seq)
        else:
            output = o.masked_fill(mask, 0)
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 62
class TSInputDropout(RandTransform):
    """Applies input dropout with required_grad=False"""
    order = 90
    def __init__(self, magnitude=0., ex=None, **kwargs):
        self.magnitude, self.ex = magnitude, ex
        self.dropout = nn.Dropout(magnitude)
        super().__init__(**kwargs)
    
    @torch.no_grad()
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        output = self.dropout(o)
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 64
class TSTranslateX(RandTransform):
    "Moves a selected sequence window a random number of steps"
    order = 90
    def __init__(self, magnitude=0.1, ex=None, **kwargs): 
        self.magnitude, self.ex = magnitude, ex
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        seq_len = o.shape[-1]
        lambd = np.random.beta(self.magnitude, self.magnitude)
        lambd = min(lambd, 1 - lambd)
        shift = int(round(seq_len * lambd))
        if shift == 0 or shift == seq_len: return o
        if np.random.rand() < 0.5: shift = -shift
        new_start = max(0, shift)
        new_end = min(seq_len + shift, seq_len)
        start = max(0, -shift)
        end = min(seq_len - shift, seq_len)
        output = torch.zeros_like(o, dtype=o.dtype, device=o.device)
        output[..., new_start : new_end] = o[..., start : end]
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 66
class TSRandomShift(RandTransform):
    "Shifts and splits a sequence"
    order = 90
    def __init__(self, magnitude=0.02, ex=None, **kwargs): 
        self.magnitude, self.ex = magnitude, ex
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        pos = int(round(np.random.randint(0, o.shape[-1]) * self.magnitude)) * (random.randint(0, 1)*2-1)
        output = torch.cat((o[..., pos:], o[..., :pos]), dim=-1)
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 68
class TSHorizontalFlip(RandTransform):
    "Flips the sequence along the x-axis"
    order = 90
    def __init__(self, magnitude=1., ex=None, **kwargs): 
        self.magnitude, self.ex = magnitude, ex
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        output = torch.flip(o, [-1])
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 70
class TSRandomTrend(RandTransform):
    "Randomly rotates the sequence along the z-axis"
    order = 90
    def __init__(self, magnitude=0.1, ex=None, **kwargs): 
        self.magnitude, self.ex = magnitude, ex
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        flat_x = o.reshape(o.shape[0], -1)
        ran = flat_x.max(dim=-1, keepdim=True)[0] - flat_x.min(dim=-1, keepdim=True)[0]
        trend = torch.linspace(0, 1, o.shape[-1], device=o.device) * ran
        t = (1 + self.magnitude * 2 * (np.random.rand() - 0.5) * trend)
        t -= t.mean(-1, keepdim=True)
        if o.ndim == 3: t = t.unsqueeze(1)
        output = o + t
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output
    
TSRandomRotate = TSRandomTrend

# %% ../../nbs/017_data.transforms.ipynb 72
class TSVerticalFlip(RandTransform):
    "Applies a negative value to the time sequence"
    order = 90
    def __init__(self, magnitude=1., ex=None, **kwargs): 
        self.magnitude, self.ex = magnitude, ex
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor): 
        if not self.magnitude or self.magnitude <= 0: return o
        return - o

# %% ../../nbs/017_data.transforms.ipynb 74
class TSResize(RandTransform):
    "Resizes the sequence length of a time series"
    order = 90
    def __init__(self, magnitude=-0.5, size=None, ex=None, mode='linear', **kwargs): 
        "mode:  'nearest' | 'linear' | 'area'"
        self.magnitude, self.size, self.ex, self.mode = magnitude, size, ex, mode
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor): 
        if self.magnitude == 0: return o
        size = ifnone(self.size, int(round((1 + self.magnitude) * o.shape[-1])))
        output = F.interpolate(o, size=size, mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)
        return output

# %% ../../nbs/017_data.transforms.ipynb 76
class TSRandomSize(RandTransform):
    "Randomly resizes the sequence length of a time series"
    order = 90
    def __init__(self, magnitude=0.1, ex=None, mode='linear', **kwargs):
        "mode:  'nearest' | 'linear' | 'area'"
        self.magnitude, self.ex, self.mode = magnitude, ex, mode
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0: return o
        size_perc = 1 + random_half_normal() * self.magnitude * (-1 if random.random() > .5 else 1)
        return F.interpolate(o, size=int(size_perc * o.shape[-1]), mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)

# %% ../../nbs/017_data.transforms.ipynb 78
class TSRandomLowRes(RandTransform):
    "Randomly resizes the sequence length of a time series to a lower resolution"
    order = 90
    def __init__(self, magnitude=.5, ex=None, mode='linear', **kwargs): 
        "mode:  'nearest' | 'linear' | 'area'"
        self.magnitude, self.ex, self.mode = magnitude, ex, mode
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor): 
        if not self.magnitude or self.magnitude <= 0: return o
        size_perc = 1 - (np.random.rand() * (1 - self.magnitude))
        return F.interpolate(o, size=int(size_perc * o.shape[-1]), mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)

# %% ../../nbs/017_data.transforms.ipynb 79
class TSDownUpScale(RandTransform):
    "Downscales a time series and upscales it again to previous sequence length"
    order = 90
    def __init__(self, magnitude=0.5, ex=None, mode='linear', **kwargs): 
        "mode:  'nearest' | 'linear' | 'area'"
        self.magnitude, self.ex, self.mode = magnitude, ex, mode
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor): 
        if not self.magnitude or self.magnitude <= 0 or self.magnitude >= 1: return o
        output = F.interpolate(o, size=int((1 - self.magnitude) * o.shape[-1]), mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)
        output = F.interpolate(output, size=o.shape[-1], mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 81
class TSRandomDownUpScale(RandTransform):
    "Randomly downscales a time series and upscales it again to previous sequence length"
    order = 90
    def __init__(self, magnitude=.5, ex=None, mode='linear', **kwargs): 
        "mode:  'nearest' | 'linear' | 'area'"
        self.magnitude, self.ex, self.mode = magnitude, ex, mode
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor): 
        if not self.magnitude or self.magnitude <= 0 or self.magnitude >= 1: return o
        scale_factor = 0.5 + 0.5 * np.random.rand() 
        output = F.interpolate(o, size=int(scale_factor * o.shape[-1]), mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)
        output = F.interpolate(output, size=o.shape[-1], mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 83
class TSRandomConv(RandTransform):
    """Applies a convolution with a random kernel and random weights with required_grad=False"""
    order = 90
    def __init__(self, magnitude=0.05, ex=None, ks=[1, 3, 5, 7], **kwargs):
        self.magnitude, self.ex, self.ks = magnitude, ex, ks
        self.conv = nn.Conv1d(1, 1, 1, bias=False)
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor):
        if not self.magnitude or self.magnitude <= 0 or self.ks is None: return o
        ks = np.random.choice(self.ks, 1)[0] if is_listy(self.ks) else self.ks
        c_in = o.shape[1]
        weight = nn.Parameter(torch.zeros(c_in, c_in, ks, device=o.device, requires_grad=False))
        nn.init.kaiming_normal_(weight)
        self.conv.weight = weight
        self.conv.padding = ks // 2
        output = (1 - self.magnitude) * o + self.magnitude * self.conv(o)
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# %% ../../nbs/017_data.transforms.ipynb 85
class TSRandom2Value(RandTransform):
    "Randomly sets selected variables of type `TSTensor` to predefined value (default: np.nan)"
    order = 90
    def __init__(self, magnitude=0.1, sel_vars=None, sel_steps=None, static=False, value=np.nan, **kwargs):
        assert not (sel_steps is not None and static), "you must choose either static or sel_steps"
        if is_listy(sel_vars) and is_listy(sel_steps):
            sel_vars = np.asarray(sel_vars)[:, None]
        self.sel_vars, self.sel_steps = sel_vars, sel_steps
        if sel_vars is None:
            self._sel_vars = slice(None)
        else:
            self._sel_vars = sel_vars
        if sel_steps is None or static:
            self._sel_steps = slice(None)
        else:
            self._sel_steps = sel_steps
        self.magnitude, self.static, self.value = magnitude , static, value
        super().__init__(**kwargs)

    def encodes(self, o:TSTensor):
        if not self.magnitude or self.magnitude <= 0 or self.magnitude > 1: return o
        if self.static:
            if self.sel_vars is not None:
                if self.magnitude == 1:
                    o[:, self._sel_vars] = o[:, self._sel_vars].fill_(self.value)
                    return o
                else:
                    vals = torch.zeros_like(o)
                    vals[:, self._sel_vars] = torch.rand(*vals[:, self._sel_vars, 0].shape, device=o.device).unsqueeze(-1)
            else:
                if self.magnitude == 1:
                    return o.fill_(self.value) 
                else:
                    vals = torch.rand(*o.shape[:-1], device=o.device).unsqueeze(-1)
        elif self.sel_vars is not None or self.sel_steps is not None:
            if self.magnitude == 1:
                o[:, self._sel_vars, self._sel_steps] = o[:, self._sel_vars, self._sel_steps].fill_(self.value)
                return o
            else:
                vals = torch.zeros_like(o)
                vals[:, self._sel_vars, self._sel_steps] = torch.rand(*vals[:, self._sel_vars, self._sel_steps].shape, device=o.device)
        else:
            if self.magnitude == 1:
                return o.fill_(self.value) 
            else:
                vals = torch.rand_like(o)
        mask = vals > (1 - self.magnitude)
        return o.masked_fill(mask, self.value)

# %% ../../nbs/017_data.transforms.ipynb 96
class TSMask2Value(RandTransform):
    "Randomly sets selected variables of type `TSTensor` to predefined value (default: np.nan)"
    order = 90
    def __init__(self, mask_fn, value=np.nan, sel_vars=None, **kwargs):
        self.sel_vars = sel_vars
        self.mask_fn = mask_fn
        self.value = value
        super().__init__(**kwargs)

    def encodes(self, o:TSTensor):
        mask = self.mask_fn(o)
        if self.sel_vars is not None:
            mask[:, self.sel_vars] = False
        return o.masked_fill(mask, self.value)

# %% ../../nbs/017_data.transforms.ipynb 98
all_TS_randaugs = [
    
    TSIdentity, 
    
    # Noise
    (TSMagAddNoise, 0.1, 1.),
    (TSGaussianNoise, .01, 1.),
    (partial(TSMagMulNoise, ex=0), 0.1, 1),
    (partial(TSTimeNoise, ex=0), 0.1, 1.),
    (partial(TSRandomFreqNoise, ex=0), 0.1, 1.),
    partial(TSShuffleSteps, ex=0),
    (TSRandomTimeScale, 0.05, 0.5), 
    (TSRandomTimeStep, 0.05, 0.5), 
    (partial(TSFreqDenoise, ex=0), 0.1, 1.),
    (TSRandomLowRes, 0.05, 0.5),
    (TSInputDropout, 0.05, .5),
    
    # Magnitude
    (partial(TSMagWarp, ex=0), 0.02, 0.2),
    (TSMagScale, 0.2, 1.),
    (partial(TSMagScalePerVar, ex=0), 0.2, 1.),
    (partial(TSRandomConv, ex=0), .05, .2),
    partial(TSBlur, ex=0),
    partial(TSSmooth, ex=0),
    partial(TSDownUpScale, ex=0),
    partial(TSRandomDownUpScale, ex=0), 
    (TSRandomTrend, 0.1, 0.5), 
    TSVerticalFlip, 
    (TSVarOut, 0.05, 0.5), 
    (TSCutOut, 0.05, 0.5), 
    
    # Time
    (partial(TSTimeWarp, ex=0), 0.02, 0.2),
    (TSWindowWarp, 0.05, 0.5),
    (TSRandomSize, 0.05, 1.),
    TSHorizontalFlip, 
    (TSTranslateX, 0.1, 0.5),
    (TSRandomShift, 0.02, 0.2), 
    (TSRandomZoomIn, 0.05, 0.5), 
    (TSWindowSlicing, 0.05, 0.2),
    (TSRandomZoomOut, 0.05, 0.5),
    (TSRandomLookBackOut, 0.1, 1.),
    (TSRandomResizedLookBack, 0.1, 1.),
    (TSTimeStepOut, 0.01, 0.2),
    (TSRandomCropPad, 0.05, 0.5), 
    (TSRandomResizedCrop, 0.05, 0.5),
    (TSMaskOut, 0.01, 0.2),
]

# %% ../../nbs/017_data.transforms.ipynb 99
class RandAugment(RandTransform):
    order = 90
    def __init__(self, tfms:list, N:int=1, M:int=3, **kwargs):
        '''
        tfms   : list of tfm functions (not called)
        N      : number of tfms applied to each batch (usual values 1-3)
        M      : tfm magnitude multiplier (1-10, usually 3-5). Only works if tfms are tuples (tfm, min, max)
        kwargs : RandTransform kwargs
        '''
        super().__init__(**kwargs)
        if not isinstance(tfms, list): tfms = [tfms]
        self.tfms, self.N, self.magnitude = tfms, min(len(tfms), N), M / 10
        self.n_tfms, self.tfms_idxs = len(tfms), np.arange(len(tfms))

    def encodes(self, o:(NumpyTensor, TSTensor)):
        if not self.N or not self.magnitude: return o
        tfms = self.tfms if self.n_tfms==1 else L(self.tfms)[np.random.choice(np.arange(self.n_tfms), self.N, replace=False)]
        tfms_ = []
        for tfm in tfms:
            if isinstance(tfm, tuple):
                t, min_val, max_val = tfm
                tfms_ += [t(magnitude=self.magnitude * float(max_val - min_val) + min_val)]
            else:  tfms_ += [tfm()]
        output = compose_tfms(o, tfms_, split_idx=self.split_idx)
        return output

# %% ../../nbs/017_data.transforms.ipynb 101
class TestTfm(RandTransform):
    "Utility class to test the output of selected tfms during training"
    def __init__(self, tfm, magnitude=1., ex=None, **kwargs): 
        self.tfm, self.magnitude, self.ex = tfm, magnitude, ex
        self.tfmd, self.shape = [], []
        super().__init__(**kwargs)
    def encodes(self, o: TSTensor): 
        if not self.magnitude or self.magnitude <= 0: return o
        output = self.tfm(o, split_idx=self.split_idx)
        self.tfmd.append(torch.equal(o, output))
        self.shape.append(o.shape)
        return output

# %% ../../nbs/017_data.transforms.ipynb 102
def get_tfm_name(tfm):
    if isinstance(tfm, tuple): tfm = tfm[0]
    if hasattr(tfm, "func"): tfm = tfm.func
    if hasattr(tfm, "__name__"): return tfm.__name__
    elif hasattr(tfm, "__class__") and hasattr(tfm.__class__, "__name__"): return tfm.__class__.__name__
    else: return tfm
