"""This notebooks presents classical bioacoustic audio augmenting techniques"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/01_classical_audio_augmentations.ipynb.

# %% auto 0
__all__ = ['pitch_shift', 'time_shift']

# %% ../../../nbs/01_classical_audio_augmentations.ipynb 2
import librosa
import numpy as np
from typing import Literal, Any, assert_never
import python_stretch

# %% ../../../nbs/01_classical_audio_augmentations.ipynb 6
def pitch_shift(
    wav_data: np.ndarray,
    sample_rate: int,
    n_steps: float,
    method: Literal['librosa', 'python_stretch'] = 'python_stretch'
) -> np.ndarray:
    """
    Applies pitch shift to the wav data.

    Args:
        wav_data: Audio data from a wav file.
        sample_rate: Sample rate of the audio.
        n_steps: Number of semitones to shift (positive = higher, negative = lower).
        method: librosa or pitch strech

    Returns:
        np.ndarray: Pitch-shifted wav data.
    """
    if method == 'librosa':
        return librosa.effects.pitch_shift(wav_data, sr=sample_rate, n_steps=n_steps)
    elif method == 'python_stretch':
        ndim = wav_data.ndim
        data = wav_data.copy()
        if ndim == 1:
            data = data[np.newaxis, :]

        stretch = python_stretch.Signalsmith.Stretch()
        stretch.preset(data.shape[0], sample_rate)
        stretch.setTransposeSemitones(n_steps)
        samples = stretch.process(data)

        return samples[0] if samples.ndim > ndim else samples

    assert_never()

# %% ../../../nbs/01_classical_audio_augmentations.ipynb 17
def time_shift(
    wav_data: np.ndarray,
    sample_rate: int,
    rate: float,
    method: Literal['librosa', 'python_stretch'] = 'python_stretch'
) -> np.ndarray:
    """
    Applies pitch shift to the wav data.

    Args:
        wav_data: Audio data from a wav file.
        sample_rate: Sample rate of the audio.
        rate: rate to be shifted.
        method: librosa or pitch strech

    Returns:
        np.ndarray: Pitch-shifted wav data.
    """
    if method == 'librosa':
        return librosa.effects.time_stretch(wav_data, rate=rate)
    elif method == 'python_stretch':
        ndim = wav_data.ndim
        data = wav_data.copy()
        if ndim == 1:
            data = data[np.newaxis, :]

        stretch = python_stretch.Signalsmith.Stretch()
        stretch.preset(data.shape[0], sample_rate)
        stretch.setTimeFactor(rate)
        samples = stretch.process(data)

        return samples[0] if samples.ndim > ndim else samples

    assert_never()
