# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/audio.03_augment.ipynb.

# %% ../../nbs/audio.03_augment.ipynb 1
# Contains code from:
# fastai - Apache License 2.0 - Copyright (c) 2023 fast.ai

# %% ../../nbs/audio.03_augment.ipynb 3
from __future__ import annotations

import random
import math
from collections import Counter
from fractions import Fraction
from functools import reduce
from itertools import chain, count, islice, repeat

import colorednoise
from primePy import primes

from torch import _VF
from torch.distributions import Bernoulli
import torchaudio.transforms as tatfms
import torchaudio.functional as TAF
from torchaudio.functional.functional import _get_sinc_resample_kernel, _apply_sinc_resample_kernel

from fastcore.transform import DisplayedTransform, retain_type

from fastai.data.load import DataLoader
from fastai.vision.augment import RandTransform, broadcast_vec

from ..transform import BatchRandTransform
from .core import TensorAudio, TensorSpec, TensorMelSpec
from ..imports import *

# %% auto 0
__all__ = ['Flip', 'Roll', 'AudioPadMode', 'RandomCropPad', 'VolumeMode', 'Volume', 'PeakNorm', 'VolumeOrPeakNorm', 'Noise',
           'VolumeBatch', 'PitchShift', 'PitchShiftTA', 'TimeStretch', 'PitchShiftOrTimeStretch', 'TimeMasking',
           'FrequencyMasking', 'AmplitudeToDBMode', 'AmplitudeToDB', 'AudioNormalize']

# %% ../../nbs/audio.03_augment.ipynb 6
def _get_audio_attr(x:TensorAudio|TensorSpec|TensorMelSpec|tuple[TensorAudio|TensorSpec|TensorMelSpec,...], attr:str) -> Any:
    if isinstance(x, tuple):
        return getattr(x[0], attr)
    else:
        return getattr(x, attr)

# %% ../../nbs/audio.03_augment.ipynb 8
class Flip(RandTransform):
    "Randomly flip `TensorAudio` with probability `p`"
    order, split_idx = 5, 0
    def __init__(self, p:float=0.5):
        super().__init__(p=p)

    def encodes(self, x:TensorAudio) -> Tensor:
        return x.fliplr()

# %% ../../nbs/audio.03_augment.ipynb 11
class Roll(RandTransform):
    "Randomly shift `TensorAudio` with rollover"
    order, split_idx = 6, 0
    def __init__(self,
        p:float=0.5, # Per-item probability
        max_roll:float=0.5, # Maximum shift
    ):
        store_attr('max_roll')
        super().__init__(p=p)

    def before_call(self,
        b:TensorAudio|tuple[TensorAudio,...],
        split_idx:int # Index of the train/valid dataset
    ):
        super().before_call(b, split_idx)
        if split_idx==0: self.shift = int(random.uniform(-1,1) * self.max_roll * _get_audio_attr(b, 'samples'))

    def encodes(self, x:TensorAudio) -> Tensor:
        return x.roll(shifts=self.shift, dims=-1)

# %% ../../nbs/audio.03_augment.ipynb 14
class AudioPadMode(Enum):
    "All methods of padding audio as attributes to get tab-completion and typo-proofing",
    Constant = 1
    ConstantPre = 2
    ConstantPost = 3
    Repeat = 4

# %% ../../nbs/audio.03_augment.ipynb 15
@patch
def crop_pad(x:TensorAudio,
    end_len:int,
    crop_start:int|None=None,
    pad_len:int=0,
    padmode:AudioPadMode=AudioPadMode.Repeat,
    constant:float=0
):
    if crop_start is not None:
        x = x[...,crop_start:crop_start+end_len]
    elif pad_len:
        if padmode==AudioPadMode.ConstantPre:
            x = _VF.constant_pad_nd(x, (pad_len, 0), constant)
        elif padmode==AudioPadMode.ConstantPost:
            x = _VF.constant_pad_nd(x, (0, pad_len), constant)
        elif padmode==AudioPadMode.Repeat:
            if pad_len >= x.samples:
                x = x.repeat(1, 2+max(pad_len//x.samples, 2))
            else:
                x = torch.cat([x, x[...,0:pad_len]], dim=-1)
        elif padmode==AudioPadMode.Constant:
            if pad_len % 2 == 0: pad = pad_len//2
            else:                pad = pad_len//2 + 1
            x = _VF.constant_pad_nd(x, (pad, pad), constant)

    if x.samples > end_len:
        x = x[...,0:end_len]
    return x

# %% ../../nbs/audio.03_augment.ipynb 16
class RandomCropPad(RandTransform):
    "Randomly resize `TensorAudio` to specified length. Center crop during valid."
    split_idx, order = None, 25
    def __init__(self,
        duration:Numeric|None=None, # Crop length in seconds
        samples:int|None=None, # Crop length in samples
        padmode:AudioPadMode=AudioPadMode.Repeat, # How to pad if necessary
        constant:Numeric=0, # Value for `AudioPadMode.Constant`
    ):
        if duration is None and samples is None:
            raise ValueError('One of `duration` or `samples` must be set')
        if constant >= 1 or constant <= -1:
            raise ValueError(f'{constant=}` must be between [-1,1]')
        if padmode not in [vol for vol in AudioPadMode]:
            raise ValueError(f"{padmode=} is not valid")
        store_attr()
        super().__init__()

    def before_call(self,
        b:TensorAudio|tuple[TensorAudio,...],
        split_idx:int # Index of the train/valid dataset
    ):
        self.samples = self.samples if self.samples is not None else int(self.duration*_get_audio_attr(b, 'sr'))
        self.orig_samples = _get_audio_attr(b, 'samples')
        if split_idx:
            self.crop_start = (self.orig_samples-self.samples)//2 if self.samples < self.orig_samples else None
        else:
            self.crop_start = torch.randint(0,self.orig_samples-self.samples, (1,)) if self.samples < self.orig_samples else None

        self.pad_len = (self.samples-self.orig_samples) if self.samples > self.orig_samples else 0

    def encodes(self, x:TensorAudio) -> Tensor:
        return x.crop_pad(self.samples, self.crop_start, self.pad_len, self.padmode, self.constant)

# %% ../../nbs/audio.03_augment.ipynb 24
class VolumeMode(Enum):
    "All Volume modes as attributes to get tab-completion and typo-proofing",
    DB = 1
    Power = 2
    Amplitude = 3

# %% ../../nbs/audio.03_augment.ipynb 25
class Volume(RandTransform):
    "Randomly change `TensorAudio`'s volume"
    order, split_idx = 30, 0
    def __init__(self,
        p:float=0.75, # Per-item probability
        gain:Numeric|None=None, # Gain is a positive amplitude ratio, a power (voltage squared), or in decibels. If none, randomly select from `gain_range`.
        gain_range:tuple[Numeric,Numeric]=(-18, 6), # Random range for gain
        volmode:VolumeMode=VolumeMode.DB # One of `VolumeMode.DB`, `VolumeMode.Amplitude`, or `VolumeMode.Power`
    ):
        super().__init__(p=p)
        store_attr(but='p')
        if volmode not in [vol for vol in VolumeMode]:
            raise ValueError(f"{volmode=} is not valid")
        if volmode != VolumeMode.DB:
            self.gain_range = (max(gain_range[0],0), gain_range[1])
        self._random_gain = gain is None

    def before_call(self,
        b:TensorAudio|tuple[TensorAudio,...],
        split_idx:int # Index of the train/valid dataset
    ):
        super().before_call(b, split_idx)
        if split_idx==0 and self._random_gain: self.gain = random.uniform(*self.gain_range)

    def encodes(self, x:TensorAudio) -> Tensor:
        if self.volmode == VolumeMode.DB:
            x = TAF.gain(x, self.gain)
        elif self.volmode == VolumeMode.Amplitude:
            x = x * self.gain
        elif self.volmode == VolumeMode.Power:
            x = TAF.gain(x, 10 * math.log10(self.gain))

        return torch.clamp(x, -1, 1)

# %% ../../nbs/audio.03_augment.ipynb 28
class PeakNorm(RandTransform):
    "Randomly apply constant gain so `TensorAudio`'s loudest level is between -1 and 1 with probability `p`"
    order, split_idx = 31, 0
    def __init__(self, p:float=0.1):
        super().__init__(p=p)
        store_attr(but='p')

    def encodes(self, x:TensorAudio) -> Tensor:
            return x / torch.max(torch.abs(x))

# %% ../../nbs/audio.03_augment.ipynb 31
class VolumeOrPeakNorm(RandTransform):
    "Randomly apply `Volume` or `Peak` to `TensorAudio` in one transform"
    order, split_idx = 30, 0
    def __init__(self,
        p:float=0.75, # Per-item probability
        peak_p:float=0.1, # Probability of applying `Peak`
        gain:Numeric|None=None, # Gain is a positive amplitude ratio, a power (voltage squared), or in decibels. If none, randomly select from `gain_range`.
        gain_range:tuple[Numeric,Numeric]=(-18, 6), # Random range for gain
        volmode:VolumeMode=VolumeMode.DB # One of `VolumeMode.DB`, `VolumeMode.Amplitude`, or `VolumeMode.Power`
    ):
        super().__init__(p=p)
        store_attr(but='p')
        if volmode not in [vol for vol in VolumeMode]:
            raise ValueError(f"{volmode=} is not valid")
        if volmode != VolumeMode.DB:
            self.gain_range = (max(gain_range[0],0), gain_range[1])
        self._random_gain = gain is None

    def before_call(self,
        b:TensorAudio|tuple[TensorAudio,...],
        split_idx:int # Index of the train/valid dataset
    ):
        super().before_call(b, split_idx)
        if split_idx==0:
            self.peak = torch.randn(1) < self.peak_p
            if not self.peak and self._random_gain:
                self.gain = random.uniform(*self.gain_range)

    def encodes(self, x:TensorAudio) -> Tensor:
        if self.peak:
            x = x / torch.max(torch.abs(x))
        else:
            if self.volmode == VolumeMode.DB:
                x = TAF.gain(x, self.gain)
            elif self.volmode == VolumeMode.Amplitude:
                x = x * self.gain
            elif self.volmode == VolumeMode.Power:
                x = TAF.gain(x, 10 * math.log10(self.gain))

        return torch.clamp(x, -1, 1)

# %% ../../nbs/audio.03_augment.ipynb 33
# Noise contains code from https://github.com/fastaudio/fastaudio:
# fastaudio - MIT License - Copyright (c) 2020 scart97

class NoiseColor(Enum):
    """Exponent beta values for named noise colors (see `colored_noise`)."""
    Violet = -2
    Blue = -1
    White = 0
    Pink = 1
    Brown = 2

# %% ../../nbs/audio.03_augment.ipynb 34
class Noise(RandTransform):
    "Adds noise of specified color and level to `TensorAudio` relative to mean audio level"
    order, split_idx = 60, 0
    def __init__(self,
        p=0.25, # Per-item probability
        noise_level:float|None=None, # Loudness of noise, if None randomly selects from `noise_range`
        noise_range:tuple[float,float]=(0.01,0.1), # Range of noise loudness values
        color:NoiseColor|None=None, # Color of noise to add, if None randomly selects from `NoiseColor`
    ):
        super().__init__(p=p)
        store_attr(but='p')
        self._color_choice = [nc for nc in NoiseColor]
        if color is not None and color not in self._color_choice:
            raise ValueError(f"{color=} is not valid")
        self._random_color = color is None
        self._random_noise = noise_level is None

    def before_call(self,
        b:TensorAudio|tuple[TensorAudio,...],
        split_idx:int # Index of the train/valid dataset
    ):
        super().before_call(b, split_idx)
        if split_idx==0:
            if self._random_color: self.color = random.choice(self._color_choice)
            if self._random_noise: self.noise_level = random.uniform(*self.noise_range)

    def encodes(self, x:TensorAudio) -> Tensor:
        # if it's white noise, use torch.randn_like for speed
        if self.color == NoiseColor.White:
            noise = torch.randn_like(x)
        else:
            noise = torch.from_numpy(colorednoise.powerlaw_psd_gaussian(exponent=self.color.value, size=x.samples)).float()
        x += noise * x.abs().mean() * self.noise_level
        return x

# %% ../../nbs/audio.03_augment.ipynb 38
class VolumeBatch(BatchRandTransform):
    "Randomly change `TensorAudio`'s volume on the GPU"
    order, split_idx = 10, 0
    def __init__(self,
        p:float=0.5, # Per-item probability
        gain:Numeric|None=None, # Gain is a positive amplitude ratio, a power (voltage squared), or in decibels. If none, randomly select from `gain_range`.
        gain_range:tuple[Numeric,Numeric]=(-18, 6), # Random range for gain
        volmode:VolumeMode=VolumeMode.DB # One of `VolumeMode.DB`, `VolumeMode.Amplitude`, or `VolumeMode.Power`
    ):
        super().__init__(p=p)
        store_attr(but='p')
        if volmode not in [vol for vol in VolumeMode]:
            raise ValueError(f"{volmode=} is not valid")
        if volmode != VolumeMode.DB:
            self.gain_range = (max(gain_range[0],0), gain_range[1])
        self._random_gain = gain is None

    def before_call(self,
        b:TensorAudio|tuple[TensorAudio,...],
        split_idx:int # Index of the train/valid dataset
    ):
        super().before_call(b, split_idx)
        if split_idx==0 and self._random_gain: self.gain = random.uniform(*self.gain_range)

    def encodes(self, x:TensorAudio) -> Tensor:
        if self.volmode == VolumeMode.DB:
            x = TAF.gain(x, self.gain)
        elif self.volmode == VolumeMode.Amplitude:
            x = x * self.gain
        elif self.volmode == VolumeMode.Power:
            x = TAF.gain(x, 10 * math.log10(self.gain))

        return torch.clamp(x, -1, 1)

# %% ../../nbs/audio.03_augment.ipynb 40
# PitchShift & TimeStretch contains code from https://github.com/KentoNishi/torch-pitch-shift & https://github.com/KentoNishi/torch-time-stretch
# torch-pitch-shift & torch-time-stretch - MIT License - Copyright (c) 2021 Kento Nishi

# %% ../../nbs/audio.03_augment.ipynb 41
# https://stackoverflow.com/a/46623112/9325832
def _combinations_without_repetition(r, iterable=None, values=None, counts=None):
    if iterable:
        values, counts = zip(*Counter(iterable).items())

    f = lambda i, c: chain.from_iterable(map(repeat, i, c))
    n = len(counts)
    indices = list(islice(f(count(), counts), r))
    if len(indices) < r:
        return
    while True:
        yield tuple(values[i] for i in indices)
        for i, j in zip(reversed(range(r)), f(reversed(range(n)), reversed(counts))):
            if indices[i] != j:
                break
        else:
            return
        j = indices[i] + 1
        for i, j in zip(range(i, r), f(count(j), counts[j:])):
            indices[i] = j

# %% ../../nbs/audio.03_augment.ipynb 42
def _get_fast_shifts(
    sample_rate:int, # Audio sample rate
    condition:Callable[...,bool]=lambda x: x >= 0.5 and x <= 2 and x != 1, # Function to validate fast shift ratios
) -> list[Fraction]: # Fast pitch-shift target ratios
    "Search for pitch-shift targets that can be computed quickly for a given sample rate"
    fast_shifts = set()
    factors = primes.factors(sample_rate)
    products = []
    for i in range(1, len(factors) + 1):
        products.extend(
            [
                reduce(lambda x, y: x * y, x)
                for x in _combinations_without_repetition(i, iterable=factors)
            ]
        )
    for i in products:
        for j in products:
            f = Fraction(i, j)
            if condition(f):
                fast_shifts.add(f)
    return list(fast_shifts)

# %% ../../nbs/audio.03_augment.ipynb 43
def _semitones_to_ratio(semitones:float) -> Fraction:
    "Convert semitonal shifts into a Fraction indicating a pitch shift ratio"
    return Fraction(2.0 ** (semitones / 12.0))

# %% ../../nbs/audio.03_augment.ipynb 44
def _get_fast_stretches(
    sample_rate:int, # Audio sample rate
    condition:Callable[...,bool]=lambda x: x >= 0.5 and x <= 2 and x != 1, # Function to validate fast stretch ratios.
) -> list[Fraction]: # Fast time-stretch target ratios
    "Search for time-stretch targets that can be computed quickly for a given sample rate"
    fast_shifts = set()
    factors = primes.factors(sample_rate)
    products = []
    for i in range(1, len(factors) + 1):
        products.extend(
            [
                reduce(lambda x, y: x * y, x)
                for x in _combinations_without_repetition(i, iterable=factors)
            ]
        )
    for i in products:
        for j in products:
            f = Fraction(i, j)
            if condition(f):
                fast_shifts.add(f)
    return list(fast_shifts)

# %% ../../nbs/audio.03_augment.ipynb 46
@patch
def pitch_shift(x:TensorAudio, n_fft, hop_length, shift, sr, new_sr, gcd, kernel, width, padmode, constant, device):
    shape = x.shape
    x = x.reshape(shape[0] * shape[1], shape[2])
    x = torch.stft(x, n_fft, hop_length, return_complex=True)
    phase_advance = torch.linspace(0, math.pi * hop_length, x.shape[1], device=device)[..., None]
    x = TAF.phase_vocoder(x, float(1 / shift), phase_advance)
    phase_advance = None
    x = torch.istft(x, n_fft, hop_length)
    x = retain_type(_apply_sinc_resample_kernel(x, sr, new_sr, gcd, kernel, width), typ=TensorAudio)
    crop_start = torch.randint(0, x.shape[-1]-shape[-1], (1,)) if shape[-1] < x.shape[-1]  else None
    pad_len = (shape[-1]-x.shape[-1]) if shape[-1] > x.shape[-1] else 0
    return x.crop_pad(shape[-1], crop_start, pad_len, padmode, constant).reshape(*shape)

# %% ../../nbs/audio.03_augment.ipynb 47
class PitchShift(BatchRandTransform):
    "Fast shift of `TensorAudio`'s pitch."
    order, split_idx = 20, 0
    def __init__(self,
        p:float=0.2, # Per-item probability
        semitones:tuple[float,float]=(-4.0,4.0), # Random pitch shift range in semitones to compute efficient shifts
        bins_per_octave:int=12, # Number of steps per octave
        padmode:AudioPadMode=AudioPadMode.Repeat, # How to pad if necessary
        constant:Numeric=0, # Value for `AudioPadMode.Constant`
        split:int|None=None # Apply transform to `split` items at a time. Use to prevent GPU OOM.
    ):
        super().__init__(p=p)
        store_attr(but='p')
        self.sr = 0
        if padmode not in [vol for vol in AudioPadMode]:
            raise ValueError(f"{padmode=} is not valid")

    def before_call(self,
        b:TensorAudio|tuple[TensorAudio,...],
        split_idx:int # Index of the train/valid dataset
    ):
        super().before_call(b, split_idx)
        if not self.sr:
            self.sr = _get_audio_attr(b, 'sr')
            self.n_fft = self.sr // 64
            self.hop_length = self.n_fft // 32
            self._fast_shifts = _get_fast_shifts(self.sr, self._validate_fast_shifts)
            self.type = _get_audio_attr(b, 'dtype')
            self.device = _get_audio_attr(b, 'device')
        if split_idx==0:
            self.shift = random.choice(self._fast_shifts)
            if not isinstance(self.shift, Fraction):
                self.shift = 2.0 ** (float(self.shift) / self.bins_per_octave)
            self.new_sr = int(self.sr/self.shift)
            self.gcd = math.gcd(self.sr, self.new_sr)
            self.kernel, self.width = _get_sinc_resample_kernel(self.sr, self.new_sr, self.gcd,
                                                                6, 0.99, 'sinc_interpolation',
                                                                None, self.device, self.type)

    def encodes(self, x:TensorAudio) -> Tensor:
        if self.split:
            xs = []
            for xi in retain_type(x.split(self.split), TensorAudio):
                xs.append(xi.pitch_shift(self.n_fft, self.hop_length, self.shift, self.sr, self.new_sr,
                                         self.gcd, self.kernel, self.width, self.padmode, self.constant, self.device))
            return torch.cat(xs, dim=0)
        else:
            return x.pitch_shift(self.n_fft, self.hop_length, self.shift, self.sr, self.new_sr,
                                 self.gcd, self.kernel, self.width, self.padmode, self.constant, self.device)

    def _validate_fast_shifts(self, x):
        return x >= _semitones_to_ratio(self.semitones[0]) and x <= _semitones_to_ratio(self.semitones[1]) and x != 1

# %% ../../nbs/audio.03_augment.ipynb 50
class PitchShiftTA(BatchRandTransform):
    "Shift the `TensorAudio`'s pitch using TorchAudio. Can be slower than `PitchShift`"
    order, split_idx = 20, 0
    def __init__(self,
        p:float=0.2, # Per-item probability
        n_steps:int|None=None, # The (fractional) steps to shift waveform
        n_step_range:tuple[int,int] = (2,6), # Random `n_steps` range if `n_steps` is None
        bins_per_octave:int=12, # Number of steps per octave
        n_fft:int=512, # Size of FFT, creates n_fft // 2 + 1 bins
        win_length:int|None=None, # Window size. Defaults to `n_fft`
        hop_length:int|None=None, # Length of hop between STFT windows. Defaults to `win_length` // 4
        window_fn:Callable[..., Tensor]=torch.hann_window, # Window tensor applied/multiplied to each frame/window
        wkwargs:dict|None=None, # Args for `window_fn`
    ):
        super().__init__(p=p)
        store_attr(but='p,window_fn,wkwargs')
        self.sr=0
        self.win_length = win_length if win_length is not None else n_fft
        self.hop_length = hop_length if hop_length is not None else self.win_length // 4
        self.window = window_fn(self.win_length) if wkwargs is None else window_fn(self.win_length, **wkwargs)

    def before_call(self,
        b:TensorAudio|tuple[TensorAudio,...],
        split_idx:int # Index of the train/valid dataset
    ):
        super().before_call(b, split_idx)
        if not self.sr:
            self.sr = _get_audio_attr(b, 'sr')
        if split_idx==0:
            self.steps = random.uniform(*self.n_step_range) if self.n_steps is None else self.n_steps

    def encodes(self, x:TensorAudio) -> Tensor:
        return TAF.pitch_shift(x, self.sr, self.steps, self.bins_per_octave, self.n_fft,
                               self.win_length, self.hop_length, self.window)

# %% ../../nbs/audio.03_augment.ipynb 53
@patch
def time_stretch(x:TensorAudio, n_fft, hop_length, stretch, padmode, constant, device):
    shape = x.shape
    x = x.reshape(shape[0] * shape[1], shape[2])
    x = torch.stft(x, n_fft, hop_length, return_complex=True)
    phase_advance = torch.linspace(0, math.pi * hop_length, x.shape[1], device=device)[..., None]
    x = TAF.phase_vocoder(x, float(1/stretch), phase_advance)
    phase_advance = None
    x = torch.istft(x, n_fft, hop_length)
    crop_start = torch.randint(0, x.shape[-1]-shape[-1], (1,)) if shape[-1] < x.shape[-1]  else None
    pad_len = (shape[-1]-x.shape[-1]) if shape[-1] > x.shape[-1] else 0
    return x.crop_pad(shape[-1], crop_start, pad_len, padmode, constant).reshape(*shape)

# %% ../../nbs/audio.03_augment.ipynb 54
class TimeStretch(BatchRandTransform):
    "Fast time stretch of `TensorAudio`"
    order, split_idx = 25, 0
    def __init__(self,
        p:float=0.2, # Per-item probability
        stretch_rates:tuple[float,float]=(0.5,2.0), # Random time stretch range to compute efficient stretches. Defaults to 50%-200% speed
        padmode:AudioPadMode=AudioPadMode.Repeat, # How to pad if necessary
        constant:Numeric=0, # Value for `AudioPadMode.Constant`
        split:int|None=None # Apply transform to `split` items at a time. Use to prevent GPU OOM.
    ):
        super().__init__(p=p)
        store_attr(but='p')
        self.sr = 0
        if padmode not in [vol for vol in AudioPadMode]:
            raise ValueError(f"{padmode=} is not valid")

    def before_call(self,
        b:TensorAudio|tuple[TensorAudio,...],
        split_idx:int # Index of the train/valid dataset
    ):
        super().before_call(b, split_idx)
        if not self.sr:
            self.sr = _get_audio_attr(b, 'sr')
            self.n_fft = self.sr // 64
            self.hop_length = self.n_fft // 32
            self._fast_stretchs = _get_fast_stretches(self.sr, self._validate_fast_stretches)
            self.device = _get_audio_attr(b, 'device')
        if split_idx==0:
            self.stretch = float(1/random.choice(self._fast_stretchs))

    def encodes(self, x:TensorAudio) -> Tensor:
        if self.split:
            xs = []
            for xi in retain_type(x.split(self.split), TensorAudio):
                xs.append(xi.time_stretch(self.n_fft, self.hop_length, self.stretch, self.padmode, self.constant, self.device))
            return torch.cat(xs, dim=0)
        else:
            return x.time_stretch(self.n_fft, self.hop_length, self.stretch, self.padmode, self.constant, self.device)

    def _validate_fast_stretches(self, x):
        return x >= self.stretch_rates[0] and x <= self.stretch_rates[1] and x != 1

# %% ../../nbs/audio.03_augment.ipynb 57
class PitchShiftOrTimeStretch(BatchRandTransform):
    "Randomly apply either `PitchShift` or `TimeStretch` to `TensorAudio` to minimize distortion"
    order, split_idx = 25, 0
    def __init__(self,
        p:float=0.4, # Per-item probability
        semitones:tuple[float,float]=(-4.0,4.0), # Random pitch shift range in semitones to compute efficient shifts
        bins_per_octave:int=12, # Number of steps per octave
        stretch_rates:tuple[float,float]=(0.5,2.0), # Random time stretch range to compute efficient stretches. Defaults to 50%-200% speed
        padmode:AudioPadMode=AudioPadMode.Repeat, # How to pad if necessary
        constant:Numeric=0, # Value for `AudioPadMode.Constant`
        split:int|None=None # Apply transform to `split` items at a time. Use to prevent OOM.
    ):
        super().__init__(p=p)
        if padmode not in [vol for vol in AudioPadMode]:
            raise ValueError(f"{padmode=} is not valid")
        self.bernoulli = Bernoulli(.5)
        self.shift   = PitchShift(1., semitones, bins_per_octave, padmode, constant, split)
        self.stretch = TimeStretch(1., stretch_rates, padmode, constant, split)

    def encodes(self, x:TensorAudio) -> Tensor:
        self.shift_smpl = self.bernoulli.sample((x.shape[0],))
        self.stretchs_smpl = (self.shift_smpl - 1) *-1
        self.shift_idxs, self.stretch_idxs = self.shift_smpl.bool(), self.stretchs_smpl.bool()
        if self.shift_smpl.sum() > 0:
            x[self.shift_idxs] = self.shift(x[self.shift_idxs], split_idx=0)
        if self.stretchs_smpl.sum() > 0:
            x[self.stretch_idxs] = self.stretch(x[self.stretch_idxs], split_idx=0)
        return x

# %% ../../nbs/audio.03_augment.ipynb 62
class TimeMasking(BatchRandTransform):
    "Randomly apply masking to a `TensorSpec` or `TensorMelSpec` in the time domain"
    split_idx, order = 0, 81
    def __init__(self,
        p:float=0.25, # Per-item probability
        max_mask:float=0.2, # Maximum possible length of the mask [0, max_mask)
        iid_masks:bool=True, # Apply different masks to each example/channel in the batch
        mask_value:int|None=0 # If None, random value between batch min and max
    ):
        store_attr(but='p')
        super().__init__(p=p)

    def encodes(self, x:TensorSpec|TensorMelSpec) -> Tensor:
        mask_param = int(self.max_mask * x.shape[-1])
        mv = random.randint(int(x.min), int(x.max)) if self.mask_value is None else self.mask_value
        if self.iid_masks and x.dim() == 4:
            return TAF.mask_along_axis_iid(x, mask_param, mv, 3)
        else:
            return TAF.mask_along_axis(x, mask_param, mv, 2)

# %% ../../nbs/audio.03_augment.ipynb 65
class FrequencyMasking(BatchRandTransform):
    "Randpmly apply masking to a `TensorSpec` or `TensorMelSpec` in the frequency domain"
    split_idx, order = 0, 82
    def __init__(self,
        p:float=0.25, # Per-item probability
        max_mask:float=0.2, # Maximum possible length of the mask [0, max_mask)
        iid_masks:bool=True, # Apply different masks to each example/channel in the batch
        mask_value:int|None=0 # If None, random value between batch min and max
    ):
        store_attr(but='p')
        super().__init__(p=p)

    def encodes(self, x:TensorSpec|TensorMelSpec) -> Tensor:
        mask_param = int(self.max_mask * x.shape[-2])
        mv = random.randint(int(x.min), int(x.max)) if self.mask_value is None else self.mask_value
        if self.iid_masks and x.dim() == 4:
            return TAF.mask_along_axis_iid(x, mask_param, mv, 2)
        else:
            return TAF.mask_along_axis(x, mask_param, mv, 1)

# %% ../../nbs/audio.03_augment.ipynb 69
class AmplitudeToDBMode(Enum):
    "All AmplitudeToDB modes as attributes to get tab-completion and typo-proofing",
    Power = 'power'
    Magnitude = 'magnitude'

# %% ../../nbs/audio.03_augment.ipynb 70
class AmplitudeToDB(DisplayedTransform):
    "Turn a `TensorSpec` or `TensorMelSpec` from the power/amplitude scale to the decibel scale"
    order = 98
    def __init__(self,
        top_db:float|None=None, # Minimum negative cut-off in decibels. A reasonable number is 80.
        mode:AmplitudeToDBMode=AmplitudeToDBMode.Power # Power is elementwise square of magnitude
    ):
        self.amdb = tatfms.AmplitudeToDB(mode.value, top_db)

    def encodes(self, x:TensorSpec|TensorMelSpec) -> Tensor:
        return self.amdb(x)

    def to(self, *args, **kwargs):
        device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)
        self.amdb.to(device)

# %% ../../nbs/audio.03_augment.ipynb 74
class AudioNormalize(DisplayedTransform):
    "Normalize/denorm batch of `TensorSpec` or `TensorMelSpec`."
    parameters,order = L('mean', 'std'),99
    def __init__(self, mean=None, std=None, n_spec=1, axes=(0,2,3)):
        store_attr()
        self.c = 0

    @classmethod
    def from_stats(cls, mean, std, n_spec=1, dim=1, ndim=4, cuda=True):
        return cls(*broadcast_vec(dim, ndim, mean, std, cuda=cuda), n_spec=n_spec)

    def setups(self, dl:DataLoader):
        if self.mean is None or self.std is None:
            x,*_ = dl.one_batch()
            self.mean,self.std = x.mean(self.axes, keepdim=True),x.std(self.axes, keepdim=True)+1e-7

    def encodes(self, x:TensorSpec|TensorMelSpec):
        if self.n_spec > 1:
            i = self.c*x.channels
            j = (self.c+1)*x.channels
            self.c = self.c + 1 if self.c < self.n_spec-1 else 0
            return (x-self.mean[:,i:j,...]) / self.std[:,i:j,...]
        else:
            return (x-self.mean) / self.std

    # ToDo: figure out how to handle decodes
    # def decodes(self, x:TensorSpec|TensorMelSpec):
    #     f = to_cpu if x.device.type=='cpu' else noop
    #     return (x*f(self.std) + f(self.mean))
