"""Applies data augmentation techniques to BatDetect2 training examples.

This module provides various functions and configurable components for applying
data augmentation to the training examples (`xr.Dataset` containing audio,
spectrogram, and target heatmaps) generated by the
`batdetect2.train.preprocess` module.

Data augmentation artificially increases the diversity of the training data by
applying random transformations, which generally helps improve the robustness
and generalization performance of trained models.

Augmentations included:
- Time-based: `select_subclip`, `warp_spectrogram`, `mask_time`.
- Amplitude/Noise-based: `mix_examples`, `add_echo`, `scale_volume`,
  `mask_frequency`.

Some augmentations modify the audio waveform and require recomputing the
spectrogram using the `PreprocessorProtocol`, while others operate directly
on the spectrogram or target heatmaps. The entire augmentation pipeline can be
configured using the `AugmentationsConfig` class, specifying a sequence of
augmentation steps, each with its own parameters and application probability.
The `build_augmentations` function constructs the final augmentation callable
from this configuration.
"""

from functools import partial
from typing import Annotated, Callable, List, Literal, Optional, Union

import numpy as np
import xarray as xr
from pydantic import Field
from soundevent import arrays, data

from batdetect2.configs import BaseConfig, load_config
from batdetect2.preprocess import PreprocessorProtocol
from batdetect2.train.types import Augmentation
from batdetect2.utils.arrays import adjust_width

__all__ = [
    "AugmentationConfig",
    "AugmentationsConfig",
    "DEFAULT_AUGMENTATION_CONFIG",
    "EchoAugmentationConfig",
    "ExampleSource",
    "FrequencyMaskAugmentationConfig",
    "MixAugmentationConfig",
    "TimeMaskAugmentationConfig",
    "VolumeAugmentationConfig",
    "WarpAugmentationConfig",
    "add_echo",
    "build_augmentations",
    "load_augmentation_config",
    "mask_frequency",
    "mask_time",
    "mix_examples",
    "scale_volume",
    "warp_spectrogram",
]

ExampleSource = Callable[[], xr.Dataset]
"""Type alias for a function that returns a training example (`xr.Dataset`).

Used by the `mix_examples` augmentation to fetch another example to mix with.
"""


class MixAugmentationConfig(BaseConfig):
    """Configuration for MixUp augmentation (mixing two examples)."""

    augmentation_type: Literal["mix_audio"] = "mix_audio"

    probability: float = 0.2
    """Probability of applying this augmentation to an example."""

    min_weight: float = 0.3
    """Minimum mixing weight (lambda) applied to the primary example."""

    max_weight: float = 0.7
    """Maximum mixing weight (lambda) applied to the primary example."""


def mix_examples(
    example: xr.Dataset,
    other: xr.Dataset,
    preprocessor: PreprocessorProtocol,
    weight: Optional[float] = None,
    min_weight: float = 0.3,
    max_weight: float = 0.7,
) -> xr.Dataset:
    """Combine two training examples using MixUp augmentation.

    Performs a weighted linear combination of the audio waveforms from two
    examples (`example` and `other`). The spectrogram is then *recomputed*
    from the mixed audio using the provided `preprocessor`. Target heatmaps
    (detection, class) are combined by taking the element-wise maximum. Target
    size heatmaps are combined by element-wise addition.

    Parameters
    ----------
    example : xr.Dataset
        The primary training example.
    other : xr.Dataset
        The second training example to mix with `example`.
    preprocessor : PreprocessorProtocol
        The preprocessor used to recompute the spectrogram from mixed audio.
        Ensures consistency with original preprocessing.
    weight : float, optional
        The mixing weight (lambda) applied to the primary `example`. The weight
        for `other` will be `(1 - weight)`. If None, a random weight is chosen
        uniformly between `min_weight` and `max_weight`.
    min_weight : float, default=0.3
        Minimum value for the random weight lambda.
    max_weight : float, default=0.7
        Maximum value for the random weight lambda. Must be >= `min_weight`.

    Returns
    -------
    xr.Dataset
        A new dataset representing the mixed example, with combined audio,
        recomputed spectrogram, and combined target heatmaps. Attributes from
        the primary `example` are preserved.
    """
    if weight is None:
        weight = np.random.uniform(min_weight, max_weight)

    audio1 = example["audio"]
    audio2 = adjust_width(other["audio"].values, len(audio1))

    with xr.set_options(keep_attrs=True):
        combined = weight * audio1 + (1 - weight) * audio2

    spectrogram = preprocessor.compute_spectrogram(
        combined.rename({"audio_time": "time"})
    ).data

    # NOTE: The subclip's spectrogram might be slightly longer than the
    # spectrogram computed from the subclip's audio. This is due to a
    # simplification in the subclip process: It doesn't account for the
    # spectrogram parameters to precisely determine the corresponding audio
    # samples. To work around this, we pad the computed spectrogram with zeros
    # as needed.
    previous_width = len(example["time"])
    spectrogram = adjust_width(spectrogram, previous_width)

    detection_heatmap = xr.apply_ufunc(
        np.maximum,
        example["detection"],
        adjust_width(other["detection"].values, previous_width),
    )

    class_heatmap = xr.apply_ufunc(
        np.maximum,
        example["class"],
        adjust_width(other["class"].values, previous_width),
    )

    size_heatmap = example["size"] + adjust_width(
        other["size"].values, previous_width
    )

    return xr.Dataset(
        {
            "audio": combined,
            "spectrogram": xr.DataArray(
                data=spectrogram,
                dims=example["spectrogram"].dims,
                coords=example["spectrogram"].coords,
            ),
            "detection": detection_heatmap,
            "class": class_heatmap,
            "size": size_heatmap,
        },
        attrs=example.attrs,
    )


class EchoAugmentationConfig(BaseConfig):
    """Configuration for adding synthetic echo/reverb."""

    augmentation_type: Literal["add_echo"] = "add_echo"

    probability: float = 0.2
    """Probability of applying this augmentation."""

    max_delay: float = 0.005
    min_weight: float = 0.0
    max_weight: float = 1.0


def add_echo(
    example: xr.Dataset,
    preprocessor: PreprocessorProtocol,
    delay: Optional[float] = None,
    weight: Optional[float] = None,
    min_weight: float = 0.1,
    max_weight: float = 1.0,
    max_delay: float = 0.005,
) -> xr.Dataset:
    """Add a synthetic echo to the audio waveform.

    Creates an echo by adding a delayed and attenuated version of the original
    audio waveform back onto itself. The spectrogram is then recomputed from
    the modified audio. Target heatmaps remain unchanged.

    Parameters
    ----------
    example : xr.Dataset
        The input training example containing 'audio', 'spectrogram', etc.
    preprocessor : PreprocessorProtocol
        Preprocessor used to recompute the spectrogram from the modified audio.
    delay : float, optional
        The delay time in seconds for the echo. If None, chosen randomly
        between 0 and `max_delay`.
    weight : float, optional
        The relative amplitude (weight) of the echo compared to the original.
        If None, chosen randomly between `min_weight` and `max_weight`.
    min_weight : float, default=0.1
        Minimum value for the random echo weight.
    max_weight : float, default=1.0
        Maximum value for the random echo weight. Must be >= `min_weight`.
    max_delay : float, default=0.005
        Maximum value for the random echo delay in seconds. Must be >= 0.

    Returns
    -------
    xr.Dataset
        A new dataset with the echo added to the 'audio' variable and the
        'spectrogram' variable recomputed. Other variables (targets, attrs)
        are copied from the original example.
    """
    if delay is None:
        delay = np.random.uniform(0, max_delay)

    if weight is None:
        weight = np.random.uniform(min_weight, max_weight)

    audio = example["audio"]
    step = arrays.get_dim_step(audio, "audio_time")
    audio_delay = audio.shift(audio_time=int(delay / step), fill_value=0)

    with xr.set_options(keep_attrs=True):
        audio = audio + weight * audio_delay

    spectrogram = preprocessor.compute_spectrogram(
        audio.rename({"audio_time": "time"}),
    ).data

    # NOTE: The subclip's spectrogram might be slightly longer than the
    # spectrogram computed from the subclip's audio. This is due to a
    # simplification in the subclip process: It doesn't account for the
    # spectrogram parameters to precisely determine the corresponding audio
    # samples. To work around this, we pad the computed spectrogram with zeros
    # as needed.
    spectrogram = adjust_width(
        spectrogram,
        example["spectrogram"].sizes["time"],
    )

    return example.assign(
        audio=audio,
        spectrogram=xr.DataArray(
            data=spectrogram,
            dims=example["spectrogram"].dims,
            coords=example["spectrogram"].coords,
            attrs=example["spectrogram"].attrs,
        ),
    )


class VolumeAugmentationConfig(BaseConfig):
    """Configuration for random volume scaling of the spectrogram."""

    augmentation_type: Literal["scale_volume"] = "scale_volume"
    probability: float = 0.2
    min_scaling: float = 0.0
    max_scaling: float = 2.0


def scale_volume(
    example: xr.Dataset,
    factor: Optional[float] = None,
    max_scaling: float = 2,
    min_scaling: float = 0,
) -> xr.Dataset:
    """Scale the amplitude of the spectrogram by a random factor.

    Multiplies the entire spectrogram DataArray by a scaling factor chosen
    uniformly between `min_scaling` and `max_scaling`. This simulates changes
    in recording volume or distance to the sound source directly in the
    spectrogram domain. Audio and target heatmaps are unchanged.

    Parameters
    ----------
    example : xr.Dataset
        The input training example containing 'spectrogram'.
    factor : float, optional
        The scaling factor to apply. If None, chosen randomly between
        `min_scaling` and `max_scaling`.
    min_scaling : float, default=0.0
        Minimum value for the random scaling factor. Must be non-negative.
    max_scaling : float, default=2.0
        Maximum value for the random scaling factor. Must be >= `min_scaling`.

    Returns
    -------
    xr.Dataset
        A new dataset with the 'spectrogram' variable scaled.

    Raises
    ------
    ValueError
        If `min_scaling` > `max_scaling` or if `min_scaling` is negative.
    """
    if factor is None:
        factor = np.random.uniform(min_scaling, max_scaling)

    with xr.set_options(keep_attrs=True):
        return example.assign(spectrogram=example["spectrogram"] * factor)


class WarpAugmentationConfig(BaseConfig):
    augmentation_type: Literal["warp"] = "warp"
    probability: float = 0.2
    delta: float = 0.04


def warp_spectrogram(
    example: xr.Dataset,
    factor: Optional[float] = None,
    delta: float = 0.04,
) -> xr.Dataset:
    """Apply time warping by resampling the time axis.

    Stretches or compresses the time axis of the spectrogram and all target
    heatmaps by a random `factor` (chosen between `1-delta` and `1+delta`).
    Uses linear interpolation for the spectrogram and nearest neighbor for the
    discrete-like target heatmaps. Updates the 'time' coordinate accordingly.
    The audio waveform is not modified.

    Parameters
    ----------
    example : xr.Dataset
        Input training example with 'spectrogram', 'detection', 'class',
        'size', and 'time' coordinates.
    factor : float, optional
        The warping factor. If None, chosen randomly between `1-delta` and
        `1+delta`. Values > 1 stretch time, < 1 compress time.
    delta : float, default=0.04
        Controls the range `[1-delta, 1+delta]` for the random warp factor.
        Must be >= 0 and < 1.

    Returns
    -------
    xr.Dataset
        A new dataset with time-warped spectrogram and target heatmaps, and
        updated 'time' coordinates.
    """
    if factor is None:
        factor = np.random.uniform(1 - delta, 1 + delta)

    start_time, end_time = arrays.get_dim_range(
        example,  # type: ignore
        "time",
    )
    duration = end_time - start_time

    new_time = np.linspace(
        start_time,
        start_time + duration * factor,
        example.time.size,
    )

    with xr.set_options(keep_attrs=True):
        spectrogram = (
            example["spectrogram"]
            .interp(
                coords={"time": new_time},
                method="linear",
                kwargs=dict(
                    fill_value=0,
                ),
            )
            .clip(min=0)
        )

        detection = example["detection"].interp(
            time=new_time,
            method="nearest",
            kwargs=dict(
                fill_value=0,
            ),
        )

        classification = example["class"].interp(
            time=new_time,
            method="nearest",
            kwargs=dict(
                fill_value=0,
            ),
        )

        size = example["size"].interp(
            time=new_time,
            method="nearest",
            kwargs=dict(
                fill_value=0,
            ),
        )

        return example.assign(
            {
                "time": new_time,
                "spectrogram": spectrogram,
                "detection": detection,
                "class": classification,
                "size": size,
            }
        )


def mask_axis(
    array: xr.DataArray,
    dim: str,
    start: float,
    end: float,
    mask_value: Union[float, Callable[[xr.DataArray], float]] = np.mean,
) -> xr.DataArray:
    """Mask values along a specified dimension.

    Sets values in the DataArray to `mask_value` where the coordinate along
    `dim` falls within the range (`start`, `end`). Values outside this range
    are kept. Used as a helper for time/frequency masking.

    Parameters
    ----------
    array : xr.DataArray
        The input DataArray (e.g., spectrogram).
    dim : str
        The name of the dimension along which to mask
        (e.g., "time", "frequency").
    start : float
        The starting coordinate value for the mask range.
    end : float
        The ending coordinate value for the mask range (exclusive, typically).
        Values >= start and < end will be masked.
    mask_value : float or Callable[[xr.DataArray], float], default=np.mean
        The value to use for masking. Can be a fixed float (e.g., 0.0) or a
        callable (like `np.mean`, `np.median`) that computes the value from
        the input `array`.

    Returns
    -------
    xr.DataArray
        The DataArray with the specified range along `dim` masked.

    Raises
    ------
    ValueError
        If `dim` is not found in the array's dimensions or coordinates.
    """
    if dim not in array.dims:
        raise ValueError(f"Axis {dim} not found in array")

    coord = array.coords[dim]
    condition = (coord < start) | (coord > end)

    if callable(mask_value):
        mask_value = mask_value(array)

    with xr.set_options(keep_attrs=True):
        return array.where(condition, other=mask_value)


class TimeMaskAugmentationConfig(BaseConfig):
    augmentation_type: Literal["mask_time"] = "mask_time"
    probability: float = 0.2
    max_perc: float = 0.05
    max_masks: int = 3


def mask_time(
    example: xr.Dataset,
    max_perc: float = 0.05,
    max_mask: int = 3,
) -> xr.Dataset:
    """Apply random time masking (SpecAugment) to the spectrogram.

    Randomly selects a number of time intervals (up to `max_masks`) and masks
    (sets to the mean value) the spectrogram within those intervals. The width
    of each mask is chosen randomly up to `max_perc` of the total duration.
    Only the 'spectrogram' variable is modified.

    Parameters
    ----------
    example : xr.Dataset
        Input training example containing 'spectrogram' and 'time' coordinate.
    max_perc : float, default=0.05
        Maximum width of a single mask as a fraction of total duration.
    max_masks : int, default=3
        Maximum number of time masks to apply.

    Returns
    -------
    xr.Dataset
        Dataset with time masking applied to the 'spectrogram'.
    """
    num_masks = np.random.randint(1, max_mask + 1)
    start_time, end_time = arrays.get_dim_range(
        example,  # type: ignore
        "time",
    )

    spectrogram = example["spectrogram"]
    for _ in range(num_masks):
        mask_size = np.random.uniform(0, max_perc) * (end_time - start_time)
        start = np.random.uniform(start_time, end_time - mask_size)
        end = start + mask_size
        spectrogram = mask_axis(spectrogram, "time", start, end)

    with xr.set_options(keep_attrs=True):
        return example.assign(spectrogram=spectrogram)


class FrequencyMaskAugmentationConfig(BaseConfig):
    augmentation_type: Literal["mask_freq"] = "mask_freq"
    probability: float = 0.2
    max_perc: float = 0.10
    max_masks: int = 3


def mask_frequency(
    example: xr.Dataset,
    max_perc: float = 0.10,
    max_masks: int = 3,
) -> xr.Dataset:
    """Apply random frequency masking (SpecAugment) to the spectrogram.

    Randomly selects a number of frequency intervals (up to `max_masks`) and
    masks (sets to the mean value) the spectrogram within those intervals. The
    height of each mask is chosen randomly up to `max_perc` of the total
    frequency range. Only the 'spectrogram' variable is modified.

    Parameters
    ----------
    example : xr.Dataset
        Input training example containing 'spectrogram' and 'frequency'
        coordinate.
    max_perc : float, default=0.10
        Maximum height of a single mask as a fraction of total frequency range.
    max_masks : int, default=3
        Maximum number of frequency masks to apply.

    Returns
    -------
    xr.Dataset
        Dataset with frequency masking applied to the 'spectrogram'.

    Raises
    ------
    ValueError
        If frequency coordinate info is missing or invalid.
    """
    num_masks = np.random.randint(1, max_masks + 1)
    min_freq, max_freq = arrays.get_dim_range(
        example,  # type: ignore
        "frequency",
    )

    spectrogram = example["spectrogram"]
    for _ in range(num_masks):
        mask_size = np.random.uniform(0, max_perc) * (max_freq - min_freq)
        start = np.random.uniform(min_freq, max_freq - mask_size)
        end = start + mask_size
        spectrogram = mask_axis(spectrogram, "frequency", start, end)

    with xr.set_options(keep_attrs=True):
        return example.assign(spectrogram=spectrogram)


AugmentationConfig = Annotated[
    Union[
        MixAugmentationConfig,
        EchoAugmentationConfig,
        VolumeAugmentationConfig,
        WarpAugmentationConfig,
        FrequencyMaskAugmentationConfig,
        TimeMaskAugmentationConfig,
    ],
    Field(discriminator="augmentation_type"),
]
"""Type alias for the discriminated union of individual augmentation config."""


class AugmentationsConfig(BaseConfig):
    """Configuration for a sequence of data augmentations.

    Attributes
    ----------
    steps : List[AugmentationConfig]
        An ordered list of configuration objects, each defining a single
        augmentation step (e.g., MixAugmentationConfig,
        TimeMaskAugmentationConfig). Each step's configuration must include an
        `augmentation_type` field and a `probability` field, along with
        type-specific parameters. The augmentations will be applied
        (probabilistically) in the sequence defined by this list.
    """

    steps: List[AugmentationConfig] = Field(default_factory=list)


class MaybeApply:
    """Applies an augmentation function probabilistically."""

    def __init__(
        self,
        augmentation: Augmentation,
        probability: float = 0.2,
    ):
        """Initialize the wrapper.

        Parameters
        ----------
        augmentation : Augmentation (Callable[[xr.Dataset], xr.Dataset])
            The augmentation function to potentially apply.
        probability : float, default=0.5
            The probability (0.0 to 1.0) of applying the augmentation.
        """
        self.augmentation = augmentation
        self.probability = probability

    def __call__(self, example: xr.Dataset) -> xr.Dataset:
        """Apply the wrapped augmentation with configured probability.

        Parameters
        ----------
        example : xr.Dataset
            The input training example.

        Returns
        -------
        xr.Dataset
            The potentially augmented training example.
        """
        if np.random.random() > self.probability:
            return example

        return self.augmentation(example)


class AudioMixer:
    """Callable class for MixUp augmentation, handling example fetching.

    Wraps the `mix_examples` logic, providing the necessary `example_source`
    to fetch a second example dynamically when called.

    Parameters
    ----------
    min_weight : float
        Minimum mixing weight (lambda) for the primary example.
    max_weight : float
        Maximum mixing weight (lambda) for the primary example.
    example_source : ExampleSource (Callable[[], xr.Dataset])
        A function that, when called, returns another random training example
        dataset (`xr.Dataset`) to be mixed with the input example.
    preprocessor : PreprocessorProtocol
        The preprocessor needed to recompute the spectrogram after mixing
        audio.
    """

    def __init__(
        self,
        min_weight: float,
        max_weight: float,
        example_source: ExampleSource,
        preprocessor: PreprocessorProtocol,
    ):
        """Initialize the AudioMixer."""
        self.min_weight = min_weight
        self.example_source = example_source
        self.max_weight = max_weight
        self.preprocessor = preprocessor

    def __call__(self, example: xr.Dataset) -> xr.Dataset:
        """Fetch another example and perform mixup.

        Parameters
        ----------
        example : xr.Dataset
            The primary input training example.

        Returns
        -------
        xr.Dataset
            The mixed training example. Returns the original example if
            fetching the second example fails.
        """
        other = self.example_source()
        return mix_examples(
            example,
            other,
            self.preprocessor,
            min_weight=self.min_weight,
            max_weight=self.max_weight,
        )


def build_augmentation_from_config(
    config: AugmentationConfig,
    preprocessor: PreprocessorProtocol,
    example_source: Optional[ExampleSource] = None,
) -> Augmentation:
    """Factory function to build a single augmentation from its config.

    Takes a configuration object for one augmentation step (which includes the
    `augmentation_type` discriminator) and returns the corresponding functional
    augmentation callable (e.g., a `partial` function or a callable class like
    `AudioMixer`).

    Parameters
    ----------
    config : AugmentationConfig
        Configuration object for a single augmentation (e.g., instance of
        `MixAugmentationConfig`, `EchoAugmentationConfig`, etc.).
    preprocessor : PreprocessorProtocol
        The preprocessor object, required by augmentations that modify audio
        and need to recompute the spectrogram (e.g., mixup, echo).
    example_source : ExampleSource, optional
        A callable that provides other training examples. Required only if
        the configuration includes `MixAugmentationConfig` (`augmentation_type`
        is "mix_audio").

    Returns
    -------
    Augmentation (Callable[[xr.Dataset], xr.Dataset])
        A callable function that takes a training example (`xr.Dataset`) and
        returns a potentially augmented version.

    Raises
    ------
    ValueError
        If `config.augmentation_type` is "mix_audio" but `example_source` is
        None.
    NotImplementedError
        If `config.augmentation_type` does not match any known augmentation
        type.
    """
    if config.augmentation_type == "mix_audio":
        if example_source is None:
            raise ValueError(
                "Mix audio augmentation ('mix_audio') requires an "
                "'example_source' callable to be provided."
            )

        return AudioMixer(
            example_source=example_source,
            preprocessor=preprocessor,
            min_weight=config.min_weight,
            max_weight=config.max_weight,
        )

    if config.augmentation_type == "add_echo":
        return partial(
            add_echo,
            preprocessor=preprocessor,
            max_delay=config.max_delay,
            min_weight=config.min_weight,
            max_weight=config.max_weight,
        )

    if config.augmentation_type == "scale_volume":
        return partial(
            scale_volume,
            max_scaling=config.max_scaling,
            min_scaling=config.min_scaling,
        )

    if config.augmentation_type == "warp":
        return partial(
            warp_spectrogram,
            delta=config.delta,
        )

    if config.augmentation_type == "mask_time":
        return partial(
            mask_time,
            max_perc=config.max_perc,
            max_mask=config.max_masks,
        )

    if config.augmentation_type == "mask_freq":
        return partial(
            mask_frequency,
            max_perc=config.max_perc,
            max_masks=config.max_masks,
        )

    raise NotImplementedError(
        "Invalid or unimplemented augmentation type: "
        f"{config.augmentation_type}"
    )


DEFAULT_AUGMENTATION_CONFIG: AugmentationsConfig = AugmentationsConfig(
    steps=[
        MixAugmentationConfig(),
        EchoAugmentationConfig(),
        VolumeAugmentationConfig(),
        WarpAugmentationConfig(),
        TimeMaskAugmentationConfig(),
        FrequencyMaskAugmentationConfig(),
    ]
)


def build_augmentations(
    preprocessor: PreprocessorProtocol,
    config: Optional[AugmentationsConfig] = None,
    example_source: Optional[ExampleSource] = None,
) -> Augmentation:
    """Build a composite augmentation pipeline function from configuration.

    Takes the overall `AugmentationsConfig` (containing a list of individual
    augmentation steps), builds the callable function for each step using
    `build_augmentation_from_config`, wraps each function with `MaybeApply`
    to handle its application probability, and returns a single
    `Augmentation` function that applies the entire sequence.

    Parameters
    ----------
    config : AugmentationsConfig
        The configuration object detailing the sequence of augmentation steps.
    preprocessor : PreprocessorProtocol
        The preprocessor object, needed for audio-modifying augmentations.
    example_source : ExampleSource, optional
        A callable providing other examples, required if 'mix_audio' is used.

    Returns
    -------
    Augmentation (Callable[[xr.Dataset], xr.Dataset])
        A single callable function that takes a training example (`xr.Dataset`)
        and applies the configured sequence of augmentations probabilistically,
        returning the augmented example. Returns the original example if
        `config.steps` is empty.

    Raises
    ------
    ValueError
        If 'mix_audio' is configured but `example_source` is not provided.
    NotImplementedError
        If an unknown `augmentation_type` is encountered in `config.steps`.
    """
    config = config or DEFAULT_AUGMENTATION_CONFIG

    augmentations = []

    for step_config in config.steps:
        augmentation = build_augmentation_from_config(
            step_config,
            preprocessor=preprocessor,
            example_source=example_source,
        )
        augmentations.append(
            MaybeApply(
                augmentation=augmentation,
                probability=step_config.probability,
            )
        )

    return partial(_apply_augmentations, augmentations=augmentations)


def load_augmentation_config(
    path: data.PathLike, field: Optional[str] = None
) -> AugmentationsConfig:
    """Load the augmentations configuration from a file.

    Reads a configuration file (YAML) and validates it against the
    `AugmentationsConfig` schema, potentially extracting data from a nested
    field.

    Parameters
    ----------
    path : PathLike
        Path to the configuration file.
    field : str, optional
        Dot-separated path to a nested section within the file containing the
        augmentations configuration (e.g., "training.augmentations"). If None,
        the entire file content is used.

    Returns
    -------
    AugmentationsConfig
        The loaded and validated augmentations configuration object.

    Raises
    ------
    FileNotFoundError
        If the config file path does not exist.
    yaml.YAMLError
        If the file content is not valid YAML.
    pydantic.ValidationError
        If the loaded config data does not conform to `AugmentationsConfig`.
    KeyError, TypeError
        If `field` specifies an invalid path.
    """
    return load_config(path, schema=AugmentationsConfig, field=field)


def _apply_augmentations(
    example: xr.Dataset,
    augmentations: List[Augmentation],
):
    """Apply a sequence of augmentation functions to an example."""
    for augmentation in augmentations:
        example = augmentation(example)
    return example
