"""This notebook presents augmentation techniques that combine multiple more than one sample: Mixup, CutMix, SpecMix and more"""

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

# %% auto 0
__all__ = ['mixup', 'cutmix']

# %% ../../nbs/03_combination_augmentations.ipynb 2
import copy
import warnings
from typing import TypeAlias, Sequence

import numpy as np
from numpy.typing import NDArray

from bioacoustics_augmentation.utils.audio import Spectrogram, AudioMetadata
from bioacoustics_augmentation.utils.image import BoundingBox

# %% ../../nbs/03_combination_augmentations.ipynb 5
FloatArray: TypeAlias = NDArray[np.floating]
IntArray: TypeAlias = NDArray[np.integer]

# %% ../../nbs/03_combination_augmentations.ipynb 6
def mixup(
    primary_spec: FloatArray,
    secondary_spec: FloatArray,
    primary_label: FloatArray,
    secondary_label: FloatArray,
    alpha: float = 1.0,
    *,
    random_generator: np.random.Generator | None = None,
) -> tuple[FloatArray, FloatArray, float]:
    if random_generator is None:
        random_generator = np.random.default_rng()

    if not (alpha and alpha > 0):
        lambda_ = 1.0
        warnings.warn("Mixup is being used with invalid alpha. No Mixup will be applied.")
    else:
        lambda_ = random_generator.beta(alpha, alpha)

    mixed_data = lambda_ * primary_spec + (1.0 - lambda_) * secondary_spec
    mixed_label = lambda_ * primary_label + (1.0 - lambda_) * secondary_label

    return mixed_data, mixed_label, lambda_

# %% ../../nbs/03_combination_augmentations.ipynb 8
def _rand_bbox(
    size: tuple[int, ...],
    lambda_: float,
    random_generator: np.random.Generator
) -> BoundingBox:
    """
    Generate random bounding box coordinates for CutMix
    """
    width = size[-1]
    height = size[-2] if len(size) > 1 else 1

    cut_ratio = np.sqrt(1.0 - lambda_)
    cut_width = int(width * cut_ratio)
    cut_height = int(height * cut_ratio)

    # Uniformly sample the center of the patch
    center_x = random_generator.integers(0, width)
    center_y = random_generator.integers(0, height)

    x1 = np.clip(center_x - cut_width // 2, 0, width)
    y1 = np.clip(center_y - cut_height // 2, 0, height)
    x2 = np.clip(center_x + cut_width // 2, 0, width)
    y2 = np.clip(center_y + cut_height // 2, 0, height)

    return BoundingBox(x1=x1, y1=y1, x2=x2, y2=y2)

# %% ../../nbs/03_combination_augmentations.ipynb 9
def _paste_patch(
    source_sample: FloatArray,
    destination_sample: FloatArray,
    bbox: BoundingBox
) -> None:
    """
    Pastes a patch from the source sample to the destination sample.
    """
    if destination_sample.ndim == 1:
        # Handles 1D data like sequences (Batch, Features)
        destination_sample[bbox.x1:bbox.x2] = source_sample[bbox.x1:bbox.x2]
    else:
        # Handles 2D+ data like images (Batch, Channels, Height, Width)
        slicing = (..., slice(bbox.y1, bbox.y2), slice(bbox.x1, bbox.x2))
        destination_sample[slicing] = source_sample[slicing]

# %% ../../nbs/03_combination_augmentations.ipynb 10
def cutmix(
    primary_spec: FloatArray,
    secondary_spec: FloatArray,
    primary_label: FloatArray,
    secondary_label: FloatArray,
    alpha: float = 1.0,
    *,
    random_generator: np.random.Generator | None = None,
) -> tuple[FloatArray, FloatArray, float]:
    if random_generator is None:
        random_generator = np.random.default_rng()

    if not (alpha and alpha > 0):
        warnings.warn("CutMix is used with an invalid alpha. No augmentation will be applied.")
        return primary_spec, primary_label, 1.0

    lambda_ = random_generator.beta(alpha, alpha)
    mixed_data = primary_spec.copy()

    bbox = _rand_bbox(primary_spec.shape, lambda_, random_generator)

    if primary_spec.ndim == 1:
        total_area = primary_spec.shape[-1]
    else:
        total_area = primary_spec.shape[-1] * primary_spec.shape[-2]

    if total_area > 0:
        actual_lambda = 1.0 - (bbox.area() / total_area)

        _paste_patch(
            destination_sample=mixed_data,
            source_sample=secondary_spec,
            bbox=bbox
        )

        mixed_label = actual_lambda * primary_label + (1.0 - actual_lambda) * secondary_label
    else:
        mixed_label = primary_label.copy()

    return mixed_data, mixed_label, lambda_
