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

__all__ = ['TSStandardize', 'TSNormalize', 'TSIdentity', 'TSShuffle_HLs', 'TSMagNoise', 'TSMagMulNoise',
           'random_curve_generator', 'random_cum_curve_generator', 'random_cum_noise_generator', 'TSTimeNoise',
           'TSMagWarp', 'TSTimeWarp', 'TSMagScale', 'TSMagScaleVar', 'TSZoomIn', 'TSZoomOut', 'TSScale',
           'TSRandomTimeStep', 'TSBlur', 'TSSmooth', 'maddest', 'TSDenoise', 'TSRandomNoise', 'TSLookBack', 'TSVarOut',
           'TSCutOut', 'TSTimeStepOut', 'TSCrop', 'TSRandomCrop', 'TSRandomResizedCrop', 'TSCenterCrop', 'TSMaskOut',
           'TSTranslateX', 'TSFlip', 'TSRandomFlip', 'TSShift', 'TSRandomRotate', 'TSNeg', 'TSRandomNeg', 'TSFreqNoise',
           'TSFreqWarp', 'TSFreqScale']

# Cell
from ..imports import *
from ..utils import *
from .external import *
from .core import *

# Cell
from scipy.interpolate import CubicSpline
from scipy.ndimage import convolve1d
import pywt

# Cell
class TSStandardize(Transform):
    "Standardize/destd batch of `NumpyTensor` or `TSTensor`"
    parameters, order = L('mean', 'std'), 99
    def __init__(self, mean=None, std=None, by_sample=False, by_var=False, verbose=False):
        self.mean = tensor(mean) if mean is not None else None
        self.std = tensor(std) if std is not None else None
        self.by_sample, self.by_var = by_sample, by_var
        if by_sample and by_var: self.axes = (2)
        elif by_sample: self.axes = (1, 2)
        elif by_var: self.axes = (0, 2)
        else: self.axes = ()
        self.verbose = verbose

    @classmethod
    def from_stats(cls, mean, std): return cls(mean, std)

    def setups(self, dl: DataLoader):
        if self.mean is None or self.std is None:
            pv(f'{self.__class__.__name__} setup mean={self.mean}, std={self.std}, by_sample={self.by_sample}, by_var={self.by_var}', self.verbose)
            x, *_ = dl.one_batch()
            self.mean, self.std = x.mean(self.axes, keepdim=self.axes!=()), x.std(self.axes, keepdim=self.axes!=()) + 1e-7
            pv(f'mean: {self.mean}  std: {self.std}\n', self.verbose)

    def encodes(self, x:(NumpyTensor, TSTensor)):
        if self.by_sample: self.mean, self.std = x.mean(self.axes, keepdim=self.axes!=()), x.std(self.axes, keepdim=self.axes!=()) + 1e-7
        return (x - self.mean) / self.std

# Cell
@patch
def mul_min(x:(torch.Tensor, TSTensor, NumpyTensor), axes=(), keepdim=False):
    if axes == (): return retain_type(x.min(), x)
    axes = reversed(sorted(axes if is_listy(axes) else [axes]))
    min_x = x
    for ax in axes: min_x, _ = min_x.min(ax, keepdim)
    return retain_type(min_x, x)

@patch
def mul_max(x:(torch.Tensor, TSTensor, NumpyTensor), axes=(), keepdim=False):
    if axes == (): return retain_type(x.max(), x)
    axes = reversed(sorted(axes if is_listy(axes) else [axes]))
    max_x = x
    for ax in axes: max_x, _ = max_x.max(ax, keepdim)
    return retain_type(max_x, x)

class TSNormalize(Transform):
    "Normalize/denorm batch of `NumpyTensor` or `TSTensor`"
    parameters, order = L('min', 'max'), 99

    def __init__(self, min=None, max=None, range_min=-1, range_max=1, by_sample=True, by_var=False, verbose=False):
        self.min = tensor(min) if min is not None else None
        self.max = tensor(max) if max is not None else None
        self.range_min, self.range_max = range_min, range_max
        self.by_sample, self.by_var = by_sample, by_var
        if by_sample and by_var: self.axes = (2)
        elif by_sample: self.axes = (1, 2)
        elif by_var: self.axes = (0, 2)
        else: self.axes = ()
        self.verbose = verbose

    @classmethod
    def from_stats(cls, min, max, range_min=0, range_max=1): return cls(min, max, self.range_min, self.range_max)

    def setups(self, dl: DataLoader):
        if self.min is None or self.max is None:
            pv(f'{self.__class__.__name__} setup min={self.min}, max={self.max}, range_min={self.range_min}, range_max={self.range_max}, by_sample={self.by_sample}, by_var={self.by_var}',  self.verbose)
            x, *_ = dl.one_batch()
            self.min, self.max = x.mul_min(self.axes, keepdim=self.axes!=()), x.mul_max(self.axes, keepdim=self.axes!=())
            pv(f'min: {self.min}  max: {self.max}\n', self.verbose)

    def encodes(self, x:(NumpyTensor, TSTensor)):
        if self.by_sample: self.min, self.max = x.mul_min(self.axes, keepdim=self.axes!=()), x.mul_max(self.axes, keepdim=self.axes!=())
        return ((x - self.min) / (self.max - self.min)) * (self.range_max - self.range_min) + self.range_min

# Cell
class TSIdentity(Transform):
    "Applies the identity tfm to a `TSTensor` batch"
    order = 90
    def __init__(self, magnitude=0., **kwargs): self.magnitude = magnitude
    def encodes(self, o: TSTensor): return o

# Cell
class TSShuffle_HLs(Transform):
    "Randomly shuffles His/Lows of an OHLC `TSTensor` batch"
    order = 90
    def __init__(self, magnitude=.1, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        timesteps = o.shape[-1] // 4
        pos_rand_list = np.random.choice(np.arange(timesteps),size=random.randint(0, 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

# Cell
class TSMagNoise(Transform):
    "Applies additive noise on the y-axis for each step of a `TSTensor` batch"
    order = 90
    def __init__(self, magnitude=.02, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        if self.magnitude <= 0: return o
        seq_len = o.shape[-1]
        noise = torch.normal(0, self.magnitude, (1, seq_len), dtype=o.dtype, device=o.device)
        output = o + noise
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

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

# Cell
def random_curve_generator(o, magnitude=.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=.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=.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)

# Cell
class TSTimeNoise(Transform):
    "Applies noise to each step in the x-axis of a `TSTensor` batch based on smooth random curve"
    order = 90
    def __init__(self, magnitude=.1, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        if 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

# Cell
class TSMagWarp(Transform):
    "Applies warping to the y-axis of a `TSTensor` batch based on a smooth random curve"
    order = 90
    def __init__(self, magnitude=.02, ord=4, ex=None, **kwargs): self.magnitude, self.ord, self.ex = magnitude, ord, ex
    def encodes(self, o: TSTensor):
        if 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

# Cell
class TSTimeWarp(Transform):
    "Applies time warping to the x-axis of a `TSTensor` batch based on a smooth random curve"
    order = 90
    def __init__(self, magnitude=.02, ord=4, ex=None, **kwargs): self.magnitude, self.ord, self.ex = magnitude, ord, ex
    def encodes(self, o: TSTensor):
        if self.magnitude <= 0: return o
        seq_len = o.shape[-1]
        f = CubicSpline(np.arange(seq_len), 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

# Cell
class TSMagScale(Transform):
    "Applies scaling to each step in the y-axis of a `TSTensor` batch based on a smooth random curve"
    order = 90
    def __init__(self, magnitude=.02, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        if self.magnitude <= 0: return o
        scale = 1 + 2 * (torch.rand(1, device=o.device) - .5) * self.magnitude
        output = o * scale
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# Cell
class TSMagScaleVar(Transform):
    "Applies scaling to each variable and step in the y-axis of a `TSTensor` batch based on smooth random curves"
    order = 90
    def __init__(self, magnitude=.02, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        if self.magnitude <= 0: return o
        scale = 1 + 2 * (torch.rand((o.shape[-2], 1), device=o.device) - .5) * self.magnitude
        output = o * scale
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# Cell
class TSZoomIn(Transform):
    "Amplifies a sequence focusing on a random section of the steps"
    order = 90
    def __init__(self, magnitude=.02, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        if 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(seq_len * lambd)
        start = 0 if win_len == seq_len else np.random.randint(0, seq_len - win_len)
        f = CubicSpline(np.arange(win_len), o[..., start : start + win_len].cpu(), axis=-1)
        output = o.new(f(np.linspace(0, win_len - 1, num=seq_len)))
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# Cell
class TSZoomOut(Transform):
    "Compresses a sequence on the x-axis"
    order = 90
    def __init__(self, magnitude=.02, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        if 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(seq_len * lambd)
        if win_len == seq_len: start = 0
        else: start = np.random.randint(0, seq_len - win_len)
        f = CubicSpline(np.arange(o.shape[-1]), o.cpu(), axis=-1)
        output = torch.zeros_like(o, dtype=o.dtype, device=o.device)
        output[..., start:start + win_len] = o.new(f(np.linspace(0, seq_len - 1, num=win_len)))
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# Cell
class TSScale(Transform):
    "Randomly amplifies/ compresses a sequence on the x-axis"
    order = 90
    def __init__(self, magnitude=.02, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        if self.magnitude <= 0: return o
        if np.random.rand() <= .5: return TSZoomIn(magnitude=self.magnitude, ex=self.ex)(o)
        else: return TSZoomOut(magnitude=self.magnitude, ex=self.ex)(o)

# Cell
class TSRandomTimeStep(Transform):
    "Compresses a sequence on the x-axis by randomly selecting sequence steps"
    order = 90
    def __init__(self, magnitude=.02, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        if self.magnitude <= 0: return o
        seq_len = o.shape[-1]
        new_seq_len = int(seq_len * max(.5, (1 - np.random.rand() * self.magnitude)))
        timesteps = np.sort(np.random.choice(np.arange(seq_len),new_seq_len, replace=False))
        f = CubicSpline(np.arange(len(timesteps)), o[..., timesteps].cpu(), axis=-1)
        output = o.new(f(np.linspace(0, new_seq_len - 1, num=seq_len)))
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# Cell
class TSBlur(Transform):
    "Blurs a sequence applying a filter of type [1, 0..., 1]"
    order = 90
    def __init__(self, magnitude=.05, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        if self.magnitude <= 0: return o
        if self.magnitude == 3:  filterargs = np.array([1, 0, 1])
        else:
            magnitude = tuple((3, 3 + int(self.magnitude * 4)))
            n_zeros = int(np.random.choice(np.arange(magnitude[0], magnitude[1] + 1, 2))) - 2
            filterargs = np.array([1] + [0] * n_zeros + [1])
        w = filterargs * np.random.rand(len(filterargs))
        w = w / w.sum()
        output = o.new(convolve1d(o.cpu(), w, mode='nearest'))
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# Cell
class TSSmooth(Transform):
    "Smoothens a sequence applying a filter of type [1, 5..., 1]"
    order = 90
    def __init__(self, magnitude=.05, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        if self.magnitude <= 0: return o
        if self.magnitude == 3:  filterargs = np.array([1, 5, 1])
        else:
            magnitude = tuple((3, 3 + int(self.magnitude * 4)))
            n_ones = int(np.random.choice(np.arange(magnitude[0], magnitude[1] + 1, 2))) // 2
            filterargs = np.array([1] * n_ones + [5] + [1] * n_ones)
        w = filterargs * np.random.rand(len(filterargs))
        w = w / w.sum()
        output = o.new(convolve1d(o.cpu(), w, mode='nearest'))
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# Cell
def maddest(d, axis=None): #Mean Absolute Deviation
    return np.mean(np.absolute(d - np.mean(d, axis)), axis)

class TSDenoise(Transform):
    "Denoises a sequence applying a wavelet decomposition method"
    order = 90
    def __init__(self, magnitude=.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
    def encodes(self, o: TSTensor):
        if 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)
        if self.thr is None:
            # 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))
            coeff[1:] = (pywt.threshold(c, value=uthr, mode=self.thr_mode) for c in coeff[1:])
        elif self.thr == 'random': coeff[1:] = (pywt.threshold(c, value=np.random.rand(), mode=self.thr_mode) for c in coeff[1:])
        else: coeff[1:] = (pywt.threshold(c, value=self.thr, 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

# Cell
class TSRandomNoise(Transform):
    "Applys random noise using a wavelet decomposition method"
    order = 90
    def __init__(self, magnitude=.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
    def encodes(self, o: TSTensor):
        if 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() - .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

# Cell
class TSLookBack(Transform):
    "Selects a random number of sequence steps starting from the end"
    order = 90
    def __init__(self, magnitude=.1, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        if self.magnitude <= 0: return o
        seq_len = o.shape[-1]
        lambd = np.random.beta(self.magnitude, self.magnitude)
        lambd = min(lambd, 1 - lambd)
        lookback_per = int(lambd * seq_len)
        output = o.clone()
        output[..., :lookback_per] = 0
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# Cell
class TSVarOut(Transform):
    "Set the value of a random number of variables to zero"
    order = 90
    def __init__(self, magnitude=.1, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        if 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(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

# Cell
class TSCutOut(Transform):
    "Sets a random section of the sequence to zero"
    order = 90
    def __init__(self, magnitude=.1, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        if 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(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

# Cell
class TSTimeStepOut(Transform):
    "Sets random sequence steps to zero"
    order = 90
    def __init__(self, magnitude=.1, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        if 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(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

# Cell
class TSCrop(Transform):
    "Crops a section of the sequence of a predefined length"
    order = 90
    def __init__(self, magnitude=.50, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        if self.magnitude <= 0: return o
        magnitude = min(.5, self.magnitude)
        seq_len = o.shape[-1]
        win_len = int(seq_len * (1 - magnitude))
        start = np.random.randint(0, seq_len - win_len)
        end = start + win_len
        output = torch.zeros_like(o, dtype=o.dtype, device=o.device)
        output[..., start - end :] = o[..., start : end]
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# Cell
class TSRandomCrop(Transform):
    "Crops a section of the sequence of a random length"
    order = 90
    def __init__(self, magnitude=.05, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        if 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(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

# Cell
class TSRandomResizedCrop(Transform):
    "Crops a section of the sequence of a random length"
    order = 90
    def __init__(self, magnitude=.01, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        if 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(seq_len * lambd)
        if win_len == seq_len: return o
        start = np.random.randint(0, seq_len - win_len)
        f = CubicSpline(np.arange(win_len), o[..., start : start + win_len].cpu(), axis=-1)
        return o.new(f(np.linspace(0, win_len, num=seq_len)))

# Cell
class TSCenterCrop(Transform):
    "Crops a section of the sequence of a random length from the center"
    order = 90
    def __init__(self, magnitude=.5, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        if 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(seq_len * lambd)
        start = seq_len // 2 - win_len // 2
        end = start + win_len
        start = max(0, start)
        end = min(end, seq_len)
        output = torch.zeros_like(o, dtype=o.dtype, device=o.device)
        output[..., start : end] = o[..., start : end]
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# Cell
class TSMaskOut(Transform):
    "Set a random number of steps to zero"
    order = 90
    def __init__(self, magnitude=.05, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        if self.magnitude <= 0: return o
        seq_len = o.shape[-1]
        mask = torch.rand_like(o) <= self.magnitude
        output = o.clone()
        output[mask] = 0
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# Cell
class TSTranslateX(Transform):
    "Moves a selected sequence window a random number of steps"
    order = 90
    def __init__(self, magnitude=.05, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        if 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(seq_len * lambd * self.magnitude)
        if shift == 0: return o
        if np.random.rand() < .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

# Cell
class TSFlip(Transform):
    "Flips the sequence along the x-axis"
    order = 90
    def __init__(self, magnitude=None, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        output = torch.flip(o, [-1])
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# Cell
class TSRandomFlip(Transform):
    "Flips the sequence along the x-axis"
    order = 90
    def __init__(self, magnitude=None, ex=None, p=0.5, **kwargs):
        self.magnitude, self.ex, self.p = magnitude, ex, p
    def encodes(self, o:TSTensor):
        if random.random() < self.p: return o
        output = torch.flip(o, [-1])
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# Cell
class TSShift(Transform):
    "Shifts and splits a sequence"
    order = 90
    def __init__(self, magnitude=None, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        pos = np.random.randint(0, o.shape[-1])
        output = torch.cat((o[..., pos:], o[..., :pos]), dim=-1)
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# Cell
class TSRandomRotate(Transform):
    "Randomly rotates the sequence along the z-axis"
    order = 90
    def __init__(self, magnitude=.1, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor):
        if self.magnitude <= 0: return o
        flat_x = o.view(o.shape[0], -1)
        ran = flat_x.max(dim=-1, keepdim=True).values - flat_x.min(dim=-1, keepdim=True).values
        trend = torch.linspace(0, 1, o.shape[-1], device=o.device) * ran
        t = (1 + self.magnitude * 2 * (np.random.rand() - .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

# Cell
class TSNeg(Transform):
    "Applies a negative value to the time sequence"
    order = 90
    def __init__(self, magnitude=None, ex=None, **kwargs): self.magnitude, self.ex = magnitude, ex
    def encodes(self, o: TSTensor): return - o

# Cell
class TSRandomNeg(Transform):
    "Randomly applies a negative value to the time sequence"
    order = 90
    def __init__(self, magnitude=None, ex=None, p=.5, **kwargs): self.magnitude, self.ex, self.p = magnitude, ex, p
    def encodes(self, o: TSTensor):
        if self.p < random.random(): return o
        return - o

# Cell
class TSFreqNoise(Transform):
    "Applies noise based on a wavelet decomposition"
    order = 90
    def __init__(self, magnitude=.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
    def encodes(self, o: TSTensor):
        if self.magnitude <= 0: return o
        seq_len = o.shape[-1]
        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 + 2 * (np.random.rand() - .5) * self.magnitude for c in coeff[1:]]
        output = o.new(pywt.waverec(coeff, self.wavelet, mode=self.mode)[..., :seq_len])
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# Cell
class TSFreqWarp(Transform):
    "Applies warp based on a wavelet decomposition"
    order = 90
    def __init__(self, magnitude=.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
    def encodes(self, o: TSTensor):
        if self.magnitude <= 0: return o
        seq_len = o.shape[-1]
        self.level = 1 if self.level is None else self.level
        new_x = random_cum_noise_generator(o[:o.shape[-1] // 2], magnitude=self.magnitude)
        coeff = pywt.wavedec(o.cpu(), self.wavelet, mode=self.mode, level=self.level)
        coeff[1:] = [CubicSpline(np.arange(c.shape[-1]), c, axis=-1)(new_x[:c.shape[-1]]) for c in coeff[1:]]
        output = o.new(pywt.waverec(coeff, self.wavelet, mode=self.mode)[..., :seq_len])
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output

# Cell
class TSFreqScale(Transform):
    "Modifies the scale based on a wavelet decomposition"
    order = 90
    def __init__(self, magnitude=.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
    def encodes(self, o: TSTensor):
        if self.magnitude <= 0: return o
        seq_len = o.shape[-1]
        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() - .5) * self.magnitude) for c in coeff[1:]]
        output = o.new(pywt.waverec(coeff, self.wavelet, mode=self.mode)[..., :seq_len])
        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]
        return output