"""This notebooks presents classical bioacoustic augmenting techniques, i.e, methods that manipulates the sound or spectogram."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/02_classical_spectrogram_augmentions.ipynb.

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

# %% ../../../nbs/02_classical_spectrogram_augmentions.ipynb 2
import scipy
import librosa
import numpy as np
from typing import Literal, Any, assert_never
from tps import ThinPlateSpline

# %% ../../../nbs/02_classical_spectrogram_augmentions.ipynb 5
def time_shift(spectrogram: Spectrogram, shift: int):
    shifted_spectrogram = spectrogram.copy()

    freq_bins, time_frames = shifted_spectrogram.data.shape
    spectrogram1 = shifted_spectrogram.data[:, :time_frames - shift]
    spectrogram2 = shifted_spectrogram.data[:, time_frames - shift:]

    shifted_spectrogram.data = np.concatenate((spectrogram2, spectrogram1), axis=1)

    return shifted_spectrogram

# %% ../../../nbs/02_classical_spectrogram_augmentions.ipynb 6
def pitch_shift(spectrogram: Spectrogram, shift: int):
    shifted_spectrogram = spectrogram.copy()

    freq_bins, time_frames = shifted_spectrogram.data.shape
    spectrogram1 = shifted_spectrogram.data[:freq_bins - shift, :]
    spectrogram2 = shifted_spectrogram.data[freq_bins - shift:, :]
    shifted_spectrogram.data = np.concatenate((spectrogram2, spectrogram1), axis=0)

    return shifted_spectrogram

# %% ../../../nbs/02_classical_spectrogram_augmentions.ipynb 17
def random_shifts(
    spectrogram: Spectrogram,
    time_shift_range: tuple[int, int] | None = None,
    pitch_shift_range: tuple[int, int] | None = None,
) -> Spectrogram:
    """
    Randomly applies shift (both in time and pitch) to the original spectrogram.

    The amount of the time and pitch shifts is randomly picked from the
    ranges time_shift_range and pitch_shift_range.

    Parameters
    ----------
    spectrograms: Spectrogram
        Input spectrograms organized as a list or single Spectrogram object.
    time_shift_range: tuple[int, int] | None
        Min and max values for time shift represented as [min, max].
    pitch_shift_range: tuple[int, int] | None
        Min and max values for pitch shift represented as [min, max].

    Returns
    -------
    new_spectrograms : Spectrogram
        Spectrogram set of shifted spectrograms.
    """
    if time_shift_range is None:
        time_shift_range = [0, 0]

    if pitch_shift_range is None:
        pitch_shift_range = [0, 0]

    min_pitch_shift, max_pitch_shift = pitch_shift_range
    rand_pitch_shift = min_pitch_shift + np.random.rand() * (max_pitch_shift - min_pitch_shift)
    rand_pitch_shift = int(rand_pitch_shift)

    min_time_shift, max_time_shift = time_shift_range
    rand_time_shift = np.random.randint(min_time_shift, max_time_shift + 1)

    shifted_spectrogram = time_shift(spectrogram, rand_time_shift)
    shifted_spectrogram = pitch_shift(shifted_spectrogram, rand_pitch_shift)

    return shifted_spectrogram

# %% ../../../nbs/02_classical_spectrogram_augmentions.ipynb 29
def _sparse_image_warp(
    spectrogram: Spectrogram,
    source_control_point_locations: np.ndarray,
    dest_control_point_locations: np.ndarray,
    regularization_weight: float = 0.0,
) -> Spectrogram:
    """
    Sparse image warping using correspondences between sparse control points.

    It applies a non-linear warp to the image, where the warp is specified by
    the source and destination locations of a number of control points.

    Args:
        spectrogram: The input Spectrogram object.
        source_control_point_locations: An array of shape `[num_points, 2]` with source (y, x) coordinates.
        dest_control_point_locations: An array of shape `[num_points, 2]` with destination (y, x) coordinates.
        regularization_weight: Weight on smoothness regularizer in interpolation

    Returns:
        A tuple containing:
        - warped_spectrogram: The warped Spectrogram object.
        - flow_field: The dense flow field as a NumPy array of shape `[height, width, 2]`.
    """
    image_data = spectrogram.data
    height, width = image_data.shape[:2]

    # Handle 2D spectrogram data by adding a temporary channel dimension for warping logic
    if image_data.ndim == 2:
        image = image_data[..., np.newaxis]
    else:
        image = image_data

    tps = ThinPlateSpline(alpha=regularization_weight)
    tps.fit(dest_control_point_locations, source_control_point_locations)

    output_indices = np.indices((height, width), dtype=np.float32).transpose(1, 2, 0)
    source_indices = tps.transform(output_indices.reshape(-1, 2)).reshape(height, width, 2)

    map_coords = source_indices.transpose(2, 0, 1)

    warped_channels = [
        scipy.ndimage.map_coordinates(image[..., c], map_coords, order=1, mode='reflect')
        for c in range(image.shape[2])
    ]
    warped_image = np.stack(warped_channels, axis=-1)

    if image_data.ndim == 2:
        warped_image = np.squeeze(warped_image, axis=-1)

    warped_spectrogram = spectrogram.copy()
    warped_spectrogram.data = warped_image
    warped_spectrogram.stft = None  # Invalidate STFT as the phase is no longer correct

    return warped_spectrogram

# %% ../../../nbs/02_classical_spectrogram_augmentions.ipynb 30
def _apply_time_warp(
    spectrogram: Spectrogram,
    warp_param: int,
    random_generator: np.random.Generator
) -> Spectrogram:
    """
    Applies time warping by creating control points as described in the SpecAugment paper.
    """
    height, width = spectrogram.data.shape

    # Source point is a random point on the center horizontal line
    source_y = height // 2
    source_x = random_generator.integers(warp_param, width - warp_param)

    # Destination point is the source point shifted horizontally by a random amount
    dest_y = source_y
    dest_x = source_x + random_generator.integers(-warp_param, warp_param)

    # The paper mentions fixing 6 boundary points:
    # 4 corners and the middles of the vertical edges.
    # These points do not move.
    source_points = np.array([
        [0, 0], [height - 1, 0],                    # Left corners
        [0, width - 1], [height - 1, width - 1],    # Right corners
        [height // 2, 0], [height // 2, width - 1], # Vertical middles
        [source_y, source_x]                        # The point to be warped
    ])

    dest_points = np.array([
        [0, 0], [height - 1, 0],
        [0, width - 1], [height - 1, width - 1],
        [height // 2, 0], [height // 2, width - 1],
        [dest_y, dest_x]                            # The new location of the warped point
    ])

    return _sparse_image_warp(spectrogram, source_points, dest_points)

# %% ../../../nbs/02_classical_spectrogram_augmentions.ipynb 34
def _frequency_masking(
    spectrogram: Spectrogram,
    num_masks: int,
    mask_param: int,  # Frequency mask parameter (`F` as mentioned in the paper)
    mask_value: float,
    random_generator: np.random.Generator,
) -> Spectrogram:
    augmented_spec = spectrogram.copy()
    freq_bins, time_frames = augmented_spec.data.shape

    for _ in range(num_masks):
        # Choose a random mask width `f` from a uniform distribution [0, F)
        f = random_generator.integers(0, mask_param + 1)

        # Choose a random starting frequency bin `f0`
        if freq_bins - f > 0:
            f0 = random_generator.integers(0, freq_bins - f)

            # Apply the mask
            augmented_spec.data[f0:f0 + f, :] = mask_value

    return augmented_spec

# %% ../../../nbs/02_classical_spectrogram_augmentions.ipynb 38
def _time_masking(
    spectrogram: Spectrogram,
    num_masks: int,
    mask_param: int,  # Time mask parameter (`T` as mentioned in the paper)
    mask_value: float,
    random_generator: np.random.Generator,
) -> Spectrogram:
    augmented_spec = spectrogram.copy()
    _, time_frames = augmented_spec.data.shape

    for _ in range(num_masks):
        # Choose a random mask width `t` from a uniform distribution [0, T)
        t = random_generator.integers(0, mask_param + 1)

        # Choose a random starting time frame `t0`
        if time_frames - t > 0:
            t0 = random_generator.integers(0, time_frames - t)

            # Apply the mask
            augmented_spec.data[:, t0:t0 + t] = mask_value

    return augmented_spec

# %% ../../../nbs/02_classical_spectrogram_augmentions.ipynb 42
def apply_spec_augment(
    spectrogram: Spectrogram,
    time_warp_param: int = 80,
    freq_mask_param: int = 27,
    time_mask_param: int = 100,
    num_freq_masks: int = 1,
    num_time_masks: int = 1,
    mask_value: float = 0.0,
    random_generator: np.random.Generator | None = None,
) -> Spectrogram:
    """
    Applies SpecAugment transformations to a spectrogram using helper functions.

    Args:
        spectrogram: The input Spectrogram object.
        time_warp_param: The maximum time warp distance (W). Set to 0 to disable.
        freq_mask_param: The maximum width (F) of the frequency mask.
        time_mask_param: The maximum width (T) of the time mask.
        num_freq_masks: The number of frequency masks to apply.
        num_time_masks: The number of time masks to apply.
        mask_value: The value to fill the masked regions with. The paper uses the mean,
            which is 0 for normalized spectrograms.
        random_generator: A NumPy random generator for reproducibility.

    Returns:
        A new Spectrogram object with augmentations applied.
    """
    if random_generator is None:
        random_generator = np.random.default_rng()

    augmented_spec = spectrogram.copy()

    if time_warp_param > 0:
        augmented_spec = _apply_time_warp(
            spectrogram=augmented_spec,
            warp_param=time_warp_param,
            random_generator=random_generator,
        )

    augmented_spec = _frequency_masking(
        spectrogram=augmented_spec,
        num_masks=num_freq_masks,
        mask_param=freq_mask_param,
        mask_value=mask_value,
        random_generator=random_generator,
    )

    augmented_spec = _time_masking(
        spectrogram=augmented_spec,
        num_masks=num_time_masks,
        mask_param=time_mask_param,
        mask_value=mask_value,
        random_generator=random_generator,
    )

    return augmented_spec
