pysmo.functions

def clone_to_mini[TMini: _AnyMini](
    mini_cls: type[TMini], source: "_AnyProto", update: dict | None = None
) -> TMini:
    """Create a new instance of a Mini class from a matching other one.

    This function creates a clone of an existing class by
    [copying][copy.copy] the attributes defined in `mini_cls` from the source
    to the target. Attributes only present in the source object are ignored,
    potentially resulting in a smaller and more performant object.

    If the source instance is missing an attribute for which a default is
    defined in the target class, then that default value for that attribute is
    used.

    Args:
        mini_cls: The type of Mini class to create.
        source: The instance to clone (must contain all attributes present
            in `mini_cls`).
        update: Update or add attributes in the returned `mini_cls` instance.

    Returns:
        A new Mini instance type mini_cls.

    Raises:
        AttributeError: If the `source` instance does not contain all
            attributes in `mini_cls` (unless they are provided with the
            `update` keyword argument).

    Examples:
        Create a [`MiniSeismogram`][pysmo.MiniSeismogram] from a
        [`SacSeismogram`][pysmo.classes.SacSeismogram] instance with
        a new `begin_time`.

        ```python
        >>> from pysmo.functions import clone_to_mini
        >>> from pysmo import MiniSeismogram
        >>> from pysmo.classes import SAC
        >>> from pandas import Timestamp
        >>> from datetime import timezone
        >>> now = Timestamp.now(timezone.utc)
        >>> sac_seismogram = SAC.from_file("example.sac").seismogram
        >>> mini_seismogram = clone_to_mini(MiniSeismogram, sac_seismogram, update={"begin_time": now})
        >>> all(sac_seismogram.data == mini_seismogram.data)
        True
        >>> mini_seismogram.begin_time == now
        True
        >>>
        ```

    Tip: See Also
        [`copy_from_mini`][pysmo.functions.copy_from_mini]: Copy attributes
        from a Mini instance to matching other one.
    """

    update = update or dict()

    if all(
        map(
            lambda x: (
                hasattr(source, x.name) or x.name in update or x.default is not NOTHING
            ),
            fields(mini_cls),
        )
    ):
        clone_dict = {
            attr.name: (
                update[attr.name]
                if attr.name in update
                else copy(getattr(source, attr.name, attr.default))
            )
            for attr in fields(mini_cls)
        }
        # TODO: why do we need cast here for mypy?
        return cast(TMini, mini_cls(**clone_dict))

    raise AttributeError(
        f"Unable to create clone: {source} not compatible with {mini_cls}."
    )

def copy_from_mini(
    source: "_AnyMini", target: "_AnyProto", update: dict | None = None
) -> None:
    """Copy attributes from a Mini instance to matching other one.

    This function [copies][copy.copy] all attributes in the `source` Mini class
    instance to a compatible `target` instance.

    Args:
        source: The Mini instance to copy attributes from.
        target: Compatible target instance.
        update: Update or add attributes in the target instance.

    Raises:
        AttributeError: If the `target` instance does not contain all
            attributes in the `source` instance (unless they are
            provided with the `update` keyword argument).

    Tip: See Also
        [`clone_to_mini`][pysmo.functions.clone_to_mini]: Create a new
        instance of a Mini class from a matching other one.
    """

    update = update or dict()

    attributes = unstructure(source).keys() | set()
    attributes.update(update.keys())

    if all(map(lambda x: hasattr(target, x), attributes)):
        for attribute in attributes:
            if attribute in update:
                setattr(target, attribute, update[attribute])
            else:
                setattr(target, attribute, copy(getattr(source, attribute)))
    else:
        raise AttributeError(
            f"Unable to copy to target: {type(target)} not compatible with {type(source)}."
        )

def crop[T: Seismogram](
    seismogram: T, begin_time: Timestamp, end_time: Timestamp, clone: bool = False
) -> None | T:
    """Shorten a seismogram by providing new begin and end times.

    This function calculates the indices corresponding to the provided new
    begin and end times using [`time2index`][pysmo.functions.time2index], then
    slices the seismogram `data` array accordingly and updates the
    `begin_time`.

    Args:
        seismogram: [`Seismogram`][pysmo.Seismogram] object.
        begin_time: New begin time.
        end_time: New end time.
        clone: Operate on a clone of the input seismogram.

    Returns:
        Cropped [`Seismogram`][pysmo.Seismogram] object if called with `clone=True`.

    Raises:
        ValueError: If new begin time is after new end time.

    Examples:
        ```python
        >>> from pysmo.functions import crop
        >>> from pysmo.classes import SAC
        >>> from pandas import Timedelta
        >>> sac_seis = SAC.from_file("example.sac").seismogram
        >>> new_begin_time = sac_seis.begin_time + Timedelta(seconds=10)
        >>> new_end_time = sac_seis.end_time - Timedelta(seconds=10)
        >>> crop(sac_seis, new_begin_time, new_end_time)
        >>>
        ```
    """

    if begin_time > end_time:
        raise ValueError("New begin_time cannot be after new end_time")

    start_index = time2index(seismogram, begin_time)
    end_index = time2index(seismogram, end_time)

    if clone is True:
        seismogram = deepcopy(seismogram)

    seismogram.data = seismogram.data[start_index : end_index + 1]
    seismogram.begin_time += seismogram.delta * start_index

    return seismogram if clone is True else None

def detrend[T: Seismogram](seismogram: T, clone: bool = False) -> None | T:
    """Remove linear and/or constant trends from a seismogram.

    Args:
        seismogram: Seismogram object.
        clone: Operate on a clone of the input seismogram.

    Returns:
        Detrended [`Seismogram`][pysmo.Seismogram] object if called with `clone=True`.

    Examples:
        ```python
        >>> import numpy as np
        >>> import pytest
        >>> from pysmo.functions import detrend
        >>> from pysmo.classes import SAC
        >>> sac_seis = SAC.from_file("example.sac").seismogram
        >>> 0 == pytest.approx(np.mean(sac_seis.data), abs=1e-11)
        np.False_
        >>> detrend(sac_seis)
        >>> 0 == pytest.approx(np.mean(sac_seis.data), abs=1e-11)
        np.True_
        >>>
        ```
    """
    if clone is True:
        seismogram = deepcopy(seismogram)

    seismogram.data = scipy.signal.detrend(seismogram.data)

    if clone is True:
        return seismogram
    return None

def normalize[T: Seismogram](
    seismogram: T,
    t1: Timestamp | None = None,
    t2: Timestamp | None = None,
    clone: bool = False,
) -> None | T:
    """Normalize a seismogram with its absolute max value.

    Args:
        seismogram: Seismogram object.
        t1: Optionally restrict searching of maximum to time after this time.
        t2: Optionally restrict searching of maximum to time before this time.
        clone: Operate on a clone of the input seismogram.

    Returns:
        Normalized [`Seismogram`][pysmo.Seismogram] object if `clone=True`.

    Examples:
        ```python
        >>> import numpy as np
        >>> from pysmo.functions import normalize
        >>> from pysmo.classes import SAC
        >>> sac_seis = SAC.from_file("example.sac").seismogram
        >>> normalize(sac_seis)
        >>> -1 <= np.max(sac_seis.data) <= 1
        np.True_
        >>>
        ```
    """

    if clone is True:
        seismogram = deepcopy(seismogram)

    start_index, end_index = None, None

    if t1 is not None:
        start_index = time2index(seismogram, t1)

    if t2 is not None:
        end_index = time2index(seismogram, t2)

    seismogram.data /= np.max(np.abs(seismogram.data[start_index:end_index]))

    if clone is True:
        return seismogram

    return None

def pad[T: Seismogram](
    seismogram: T,
    begin_time: Timestamp,
    end_time: Timestamp,
    mode: "_ModeKind | _ModeFunc" = "constant",
    clone: bool = False,
    **kwargs: Any,
) -> None | T:
    """Pad seismogram data.

    This function calculates the indices corresponding to the provided new
    begin and end times using [`time2index`][pysmo.functions.time2index], then
    pads the [`data`][pysmo.Seismogram.data] array using
    [`numpy.pad`][numpy.pad] and updates the
    [`begin_time`][pysmo.Seismogram.begin_time]. Note that the actual begin and
    end times are set by indexing, so they may be slightly different than the
    provided input begin and end times.

    Args:
        seismogram: [`Seismogram`][pysmo.Seismogram] object.
        begin_time: New begin time.
        end_time: New end time.
        mode: Pad mode to use (see [`numpy.pad`][numpy.pad] for all modes).
        clone: Operate on a clone of the input seismogram.
        kwargs: Keyword arguments to pass to [`numpy.pad`][numpy.pad].

    Returns:
        Padded [`Seismogram`][pysmo.Seismogram] object if called with `clone=True`.

    Raises:
        ValueError: If new begin time is after new end time.

    Examples:
        ```python
        >>> from pysmo.functions import pad
        >>> from pysmo.classes import SAC
        >>> from pandas import Timedelta
        >>> import numpy as np
        >>>
        >>> sac_seis = SAC.from_file("example.sac").seismogram
        >>> original_length = len(sac_seis.data)
        >>> sac_seis.data
        array([2302., 2313., 2345., ..., 2836., 2772., 2723.], shape=(180000,))
        >>> new_begin_time = sac_seis.begin_time - Timedelta(seconds=10)
        >>> new_end_time = sac_seis.end_time + Timedelta(seconds=10)
        >>> pad(sac_seis, new_begin_time, new_end_time)
        >>> np.isclose(len(sac_seis.data), original_length + 20 / sac_seis.delta.total_seconds())
        np.True_
        >>> sac_seis.data
        array([0., 0., 0., ..., 0., 0., 0.], shape=(181000,))
        >>>
        ```
    """

    if begin_time > end_time:
        raise ValueError("New begin_time cannot be after new end_time")

    start_index = time2index(seismogram, begin_time, allow_out_of_bounds=True)
    end_index = time2index(seismogram, end_time, allow_out_of_bounds=True)

    if clone is True:
        seismogram = deepcopy(seismogram)

    pad_before = max(0, -start_index)
    pad_after = max(0, end_index - (len(seismogram.data) - 1))

    if pad_before > 0 or pad_after > 0:
        seismogram.data = np.pad(
            seismogram.data,
            pad_width=(pad_before, pad_after),
            mode=mode,
            **kwargs,
        )
        seismogram.begin_time += seismogram.delta * min(0, start_index)

    return seismogram if clone else None

@beartype
def resample[T: Seismogram](
    seismogram: T, delta: PositiveTimedelta, clone: bool = False
) -> None | T:
    """Resample Seismogram data using the Fourier method.

    This function uses [`scipy.resample`][scipy.signal.resample] to resample
    the data to a new sampling interval. If the new sampling interval is
    identical to the current one, no action is taken.

    Args:
        seismogram: Seismogram object.
        delta: New sampling interval.
        clone: Operate on a clone of the input seismogram.

    Returns:
        Resampled [`Seismogram`][pysmo.Seismogram] object if called with `clone=True`.

    Examples:
        ```python
        >>> from pysmo.functions import resample
        >>> from pysmo.classes import SAC
        >>> sac_seis = SAC.from_file("example.sac").seismogram
        >>> len(sac_seis.data)
        180000
        >>> original_delta = sac_seis.delta
        >>> new_delta = original_delta * 2
        >>> resample(sac_seis, new_delta)
        >>> len(sac_seis.data)
        90000
        >>>
        ```
    """
    if clone is True:
        seismogram = deepcopy(seismogram)

    if delta != seismogram.delta:
        npts = int(len(seismogram.data) * seismogram.delta / delta)
        seismogram.data = scipy.signal.resample(seismogram.data, npts)
        seismogram.delta = delta

    if clone is True:
        return seismogram
    return None

@beartype
def taper[T: Seismogram](
    seismogram: T,
    taper_width: NonNegativeTimedelta | UnitFloat,
    window_type: _WindowType = "hann",
    clone: bool = False,
) -> None | T:
    """Apply a symmetric taper to the ends of a Seismogram.

    The [`taper()`][pysmo.functions.taper] function applies a symmetric taper to
    the data of a [`Seismogram`][pysmo.Seismogram] object. The taper width is
    understood as the portion of the seismogram affected by the taper window
    function. It can be provided as an absolute duration (non-negative
    [`Timedelta`][pandas.Timedelta]), or as a fraction of seismogram length
    ([`float`][float] between `0` and `1`). Internally, absolute durations are
    converted to fractions by dividing by the total seismogram duration, and
    absolute durations should therefore not exceed the total seismogram
    duration.

    The shape of the windowing function is calculated by calling the scipy
    [`get_window()`][scipy.signal.windows.get_window] function using the number
    of samples corresponding to the fraction specified above, then it is split
    in half and applied to the beginning and end of the seismogram data. Thus
    `taper_width=0` corresponds to a rectangular window (i.e. no tapering), and
    `taper_width=1` to a symmetric taper applied to the entire length of the
    seismogram. A value of e.g. `0.5` applies the "ramp up" portion of the
    window to the first quarter of the seismogram, while the "ramp down" portion
    of the window is applied to the last quarter.

    Warning:
        The scipy [`get_window()`][scipy.signal.windows.get_window] function
        is a helper function that calculates a large variety of window shapes,
        which do not all make sense in this application (e.g. boxcar or tukey).
        Users are encouraged to read the documentation of the actual window
        functions available via
        [`get_window()`][scipy.signal.windows.get_window] to see if they can be
        split in the middle and used as "ramp up" and "ramp down" functions.

    Args:
        seismogram: Seismogram object.
        taper_width: Width of the taper to use.
        window_type: Function to calculate taper shape (see
            [`get_window`][scipy.signal.windows.get_window] for valid inputs).
        clone: Operate on a clone of the input seismogram.

    Returns:
        Tapered [`Seismogram`][pysmo.Seismogram] object if called with `clone=True`.

    Examples:
        ```python
        >>> from pysmo.functions import taper, detrend
        >>> from pysmo.classes import SAC
        >>> sac_seis = SAC.from_file("example.sac").seismogram
        >>> detrend(sac_seis)
        >>> sac_seis.data
        array([ 95.59652208, 106.59521819, 138.59391429, ..., 394.90004126,
               330.89873737, 281.89743348], shape=(180000,))
        >>> taper(sac_seis, 0.2)
        >>> sac_seis.data
        array([0.00000000e+00, 8.11814104e-07, 4.22204657e-06, ...,
               1.20300114e-05, 2.52007798e-06, 0.00000000e+00], shape=(180000,))
        >>>
        ```
    """

    nsamples: int
    if isinstance(taper_width, Timedelta):
        nsamples = taper_width // seismogram.delta
    else:
        nsamples = floor(len(seismogram.data) * taper_width)

    if nsamples > len(seismogram.data):
        raise ValueError(
            "'taper_width' is too large. Total taper width may exceed the duration of the seismogram."
        )

    if clone is True:
        seismogram = deepcopy(seismogram)

    # Need at least 2 samples to apply a taper
    if nsamples >= 2:
        window = scipy.signal.windows.get_window(window_type, nsamples, fftbins=False)
        ramp_samples = nsamples // 2
        seismogram.data[:ramp_samples] *= window[:ramp_samples]
        seismogram.data[-ramp_samples:] *= window[-ramp_samples:]

    if clone is True:
        return seismogram
    return None

@beartype
def time2index(
    seismogram: Seismogram,
    time: Timestamp,
    allow_out_of_bounds: bool = False,
) -> int:
    """
    Converts a specific timestamp to the corresponding data array index.

    Seismic data is sampled at discrete intervals. When a requested time does
    not align perfectly with a sample, this function selects the nearest
    index using the following rules:

    1. If the time is within 0.1% of a sample interval of an integer, it
       "snaps" to that integer to account for floating-point jitter.
    2. Use standard rounding (0.5 rounds up to the next index) otherwise.

    Args:
        seismogram: Seismogram object.
        time: The absolute time to convert.
        allow_out_of_bounds: If True, returns the calculated index even if it
            falls outside the seismogram's data range [0, len-1].

    Returns:
        The index of the sample closest to the provided time.

    Raises:
        ValueError: If the calculated index is outside the data array and
            `allow_out_of_bounds` is False.
    """
    # Calculate the fractional index position
    index_float = (time - seismogram.begin_time) / seismogram.delta

    # Snap to nearest integer if within a tiny tolerance (1e-3 samples)
    # This prevents 2.999999999 from being floored to 2 instead of 3.
    if np.isclose(index_float, np.round(index_float), atol=1e-3):
        index = int(np.round(index_float))
    # Standard rounding within the trace (0.5 rounds up)
    else:
        index = int(np.floor(index_float + 0.5))

    # Validation
    if 0 <= index < len(seismogram.data) or allow_out_of_bounds:
        return index

    raise ValueError(
        f"Calculated index {index} is out of bounds for seismogram with "
        f"{len(seismogram.data)} samples. (Target time: {time})"
    )

@beartype
def window[T: Seismogram](
    seismogram: T,
    window_begin_time: Timestamp,
    window_end_time: Timestamp,
    ramp_width: NonNegativeTimedelta | NonNegativeNumber,
    window_type: _WindowType = "hann",
    same_shape: bool = False,
    clone: bool = False,
) -> None | T:
    """Returns an optionally padded and tapered window of a seismogram.

    This function combines the [`crop`][pysmo.functions.crop],
    [`detrend`][pysmo.functions.detrend], [`taper`][pysmo.functions.taper], and
    optionally [`pad`][pysmo.functions.pad] functions to return a 'windowed'
    seismogram. Its purpose is to focus on a specific time window of interest,
    while also (optionally) preserving the original seismogram length and
    tapering the signal before and after the window.

    Tip:
        Note that the window defined by `window_begin_time` and
        `window_end_time` *excludes* the tapered sections, so the total length
        of the window will be the provided window length plus the tapered
        sections of the signal. This behaviour is a bit different from
        [`taper()`][pysmo.functions.taper], where the taper is applied to the
        entire signal. In a sense the tapering here is applied to the 'outside'
        of the region of interest rather than the 'inside'.

    Args:
        seismogram: Seismogram object.
        window_begin_time: Begin time of the window.
        window_end_time: End time of the window.
        ramp_width: Duration of the taper on *each side*.

            - If `float`: calculated as a fraction of the window length.
            - If `Timedelta`: used as absolute duration.

            Note: Total duration = window length + (2 * `ramp_width`).
        window_type: Taper method to use (see [`taper`][pysmo.functions.taper]).
        same_shape: If True, pad the seismogram to its original length after
            windowing.
        clone: Operate on a clone of the input seismogram.

    Returns:
        Windowed [`Seismogram`][pysmo.Seismogram] object if called with `clone=True`.

    Examples:
        In this example we focus on a window starting 500 seconds after the
        `begin_time` of the seismogram and lasting for 1000 seconds. Setting the
        `ramp_width` to 250 seconds means that the actual window will start 250
        seconds earlier and end 250 seconds later than the specified window
        begin and end times.

        ```python
        >>> from pysmo.functions import window, detrend
        >>> from pysmo.classes import SAC
        >>> from pysmo.tools.plotutils import plotseis
        >>> from pandas import Timedelta
        >>>
        >>> sac_seis = SAC.from_file("example.sac").seismogram
        >>> ramp_width = Timedelta(seconds=250)
        >>> window_begin_time = sac_seis.begin_time + Timedelta(seconds=500)
        >>> window_end_time = window_begin_time + Timedelta(seconds=1000)
        >>> windowed_seis = window(sac_seis, window_begin_time, window_end_time, ramp_width, same_shape=True, clone=True)
        >>> detrend(sac_seis)
        >>> fig = plotseis(sac_seis, windowed_seis)
        >>>
        ```

        <!-- invisible-code-block: python
        ```
        >>> import matplotlib.pyplot as plt
        >>> plt.close("all")
        >>> if savedir:
        ...     plt.style.use("dark_background")
        ...     fig = plotseis(sac_seis, windowed_seis)
        ...     fig.savefig(savedir / "functions_window-dark.png", transparent=True)
        ...
        ...     plt.style.use("default")
        ...     fig = plotseis(sac_seis, windowed_seis)
        ...     fig.savefig(savedir / "functions_window.png", transparent=True)
        >>>
        ```
        -->

        <figure markdown="span">
        ![Functions window](../../images/functions/functions_window.png#only-light){ loading=lazy }
        ![Functions window](../../images/functions/functions_window-dark.png#only-dark){ loading=lazy }
        </figure>
    """

    begin_time, end_time = seismogram.begin_time, seismogram.end_time

    ramp_duration: Timedelta
    if isinstance(ramp_width, (float, int)):
        ramp_duration = ramp_width * (window_end_time - window_begin_time)
    else:
        ramp_duration = ramp_width

    window_begin_time -= ramp_duration
    window_end_time += ramp_duration

    if clone is True:
        seismogram = crop(seismogram, window_begin_time, window_end_time, clone=True)
    else:
        crop(seismogram, window_begin_time, window_end_time)
    detrend(seismogram)
    taper(seismogram, taper_width=ramp_duration * 2, window_type=window_type)
    if same_shape is True:
        pad(seismogram, begin_time, end_time)

    if clone is True:
        return seismogram
    return None