pysmo.tools.signal

@register_filter
def bandpass[T: Seismogram](
    seismogram: T,
    freqmin: float = 0.1,
    freqmax: float = 0.5,
    corners: int = 2,
    zerophase: bool = False,
    clone: bool = False,
) -> T | None:
    """Apply a bandpass filter to the input seismogram.

    Args:
        seismogram: The input seismogram to be filtered.
        freqmin: The minimum frequency of the bandpass filter (in Hz).
        freqmax: The maximum frequency of the bandpass filter (in Hz).
        corners: The number of corners (poles) for the Butterworth filter.
        zerophase: If `True`, apply the filter in both forward and reverse
            directions to achieve zero phase distortion.
        clone: If `True`, return a new Seismogram object with the filtered
            data. If `False`, modify the input seismogram in place.

    Returns:
        A new Seismogram object containing the filtered data when called with `clone=True`.
    """
    fe = 0.5 / seismogram.delta.total_seconds()
    low = freqmin / fe
    high = freqmax / fe

    if not (0 < low < 1):
        raise ValueError(
            f"freqmin ({freqmin}) is invalid for sampling rate {1 / seismogram.delta.total_seconds()} Hz."
        )
    if not (0 < high < 1):
        raise ValueError(
            f"freqmax ({freqmax}) is invalid for sampling rate {1 / seismogram.delta.total_seconds()} Hz."
        )
    if freqmin >= freqmax:
        raise ValueError("freqmin must be less than freqmax.")

    sos = iirfilter(corners, [low, high], btype="band", ftype="butter", output="sos")

    if clone:
        seismogram = deepcopy(seismogram)

    if zerophase:
        seismogram.data = sosfiltfilt(sos, seismogram.data)
    else:
        seismogram.data = sosfilt(sos, seismogram.data)

    return seismogram if clone else None

@register_filter
def bandstop[T: Seismogram](
    seismogram: T,
    freqmin: float = 0.1,
    freqmax: float = 0.5,
    corners: int = 2,
    zerophase: bool = False,
    clone: bool = False,
) -> T | None:
    """Apply a bandstop filter to the input seismogram.

    Args:
        seismogram: The input seismogram to be filtered.
        freqmin: The minimum frequency of the bandstop filter (in Hz).
        freqmax: The maximum frequency of the bandstop filter (in Hz).
        corners: The number of corners (poles) for the Butterworth filter.
        zerophase: If `True`, apply the filter in both forward and reverse
            directions to achieve zero phase distortion.
        clone: If `True`, return a new Seismogram object with the filtered
            data. If `False`, modify the input seismogram in place.

    Returns:
        A new Seismogram object containing the filtered data when called with `clone=True`.
    """
    fe = 0.5 / seismogram.delta.total_seconds()
    low = freqmin / fe
    high = freqmax / fe

    if not (0 < low < 1):
        raise ValueError(
            f"freqmin ({freqmin}) is invalid for sampling rate {1 / seismogram.delta.total_seconds()} Hz."
        )
    if not (0 < high < 1):
        raise ValueError(
            f"freqmax ({freqmax}) is invalid for sampling rate {1 / seismogram.delta.total_seconds()} Hz."
        )
    if freqmin >= freqmax:
        raise ValueError("freqmin must be less than freqmax.")

    sos = iirfilter(
        corners, [low, high], btype="bandstop", ftype="butter", output="sos"
    )

    if clone:
        seismogram = deepcopy(seismogram)

    if zerophase:
        seismogram.data = sosfiltfilt(sos, seismogram.data)
    else:
        seismogram.data = sosfilt(sos, seismogram.data)

    return seismogram if clone else None

def delay(
    seismogram1: Seismogram,
    seismogram2: Seismogram,
    total_delay: bool = False,
    max_shift: pd.Timedelta | None = None,
    abs_max: bool = False,
) -> tuple[pd.Timedelta, float]:
    """Cross correlates two seismograms to determine signal delay.

    This function is a wrapper around the [`correlate`][scipy.signal.correlate]
    function. The default behaviour is to call the correlate function with
    `#!py mode="full"` using the full length data of the input seismograms.
    This is the most robust option, but also the slowest.

    If an approximate delay is known (e.g. because a particular phase is being
    targeted using a computed arrival time), the search space can be limited
    by setting the `max_shift` parameter to a value. The length of the
    seismogram data used for the cross-correlation is then set such that the
    calculated delay lies within +/- `max_shift`.

    Note:
        `max_shift` intentionally does *not* take the begin times of the
        seismograms into consideration. Thus, calling `#!py delay()` with
        `#!py total_delay=True` may return a delay that is larger than
        `max_shift`.

    Implications of setting the `max_shift` parameter are as follows:

    - This mode requires the seismograms to be of equal length.
    - If the true delay (i.e. the amount of time the seismograms _should_ be
      shifted by) lies within the `max_shift` range, and also produces the
      highest correlation, the delay time returned is identical for both
      methods.
    - If the true delay lies outside the `max_shift` range and produces the
      highest correlation, the delay time returned will be incorrect when
      `max_shift` is set.
    - In the event that the true delay lies within the `max_shift` range but
      the maximum signal correlation occurs outside, it will be correctly
      retrieved when the `max_shift` parameter is set, while not setting it
      yields an incorrect result.


    Args:
        seismogram1: First seismogram to use for cross correlation.
        seismogram2: Second seismogram to use for cross correlation.
        total_delay: Include the difference in `begin_time` in the delay.
        max_shift: Maximum (absolute) length of the delay.
        abs_max: Return the delay corresponding to absolute maximum.

    Returns:
        delay: Time delay of the second seismogram with respect to the first.
        cc: Normalised cross-correlation value of the overlapping
            seismograms *after* shifting (uses
            [`scipy.stats.mstats.pearsonr`][] for the calculation). This value
            ranges from -1 to 1, with 1 indicating a perfect correlation, 0
            indicating no correlation, and -1 indicating a perfect
            anti-correlation.

    Examples:
        To illustrate the `delay()` function: this reads a seismogram from a
        SAC file, then generates a second seismogram from it with a shift in
        the data and in the begin time. That makes the true delay known, so
        it can be compared against the computed delay:

        ```python
        >>> from pysmo import MiniSeismogram
        >>> from pysmo.classes import SAC
        >>> from pysmo.functions import detrend, clone_to_mini
        >>> from pysmo.tools.signal import delay
        >>> from datetime import timedelta
        >>> import numpy as np
        >>>
        >>> # Create a Seismogram from a SAC file and detrend it:
        >>> seis1 = SAC.from_file("example.sac").seismogram
        >>> detrend(seis1)
        >>>
        >>> # Create a second seismogram from the first with
        >>> # a different begin_time and a shift in the data:
        >>> seis2 = clone_to_mini(MiniSeismogram, seis1)
        >>> nroll = 1234
        >>> seis2.data = np.roll(seis2.data, nroll)
        >>> begin_time_delay = timedelta(seconds=100)
        >>> seis2.begin_time += begin_time_delay
        >>>
        >>> # The signal delay is the number of samples shifted * delta:
        >>> (signal_delay := nroll * seis1.delta).total_seconds()
        61.7
        >>>
        >>> # Call the delay function with the two seismograms and verify
        >>> # that the caclulated_delay is equal to the known signal delay:
        >>> calculated_delay, _ = delay(seis1, seis2)
        >>> calculated_delay == signal_delay
        True
        >>>
        ```

        Since the true delay is known, this can mimic a scenario where an
        approximate delay is known before the cross-correlation, which can be
        used to limit the search space and speed up the calculation. Here,
        `max_shift` is set to the known signal delay plus 1 second:

        ```python
        >>> max_shift = signal_delay + timedelta(seconds=1)
        >>> calculated_delay, _ = delay(seis1, seis2, max_shift=max_shift)
        >>> # As before, the calculated delay should be equal to the signal delay:
        >>> calculated_delay == signal_delay
        True
        >>>
        ```

        Setting `total_delay=True` also takes into account the difference
        in `begin_time` between the two seismograms:

        ```python
        >>> calculated_delay, _ = delay(seis1, seis2, total_delay=True, max_shift=signal_delay+timedelta(seconds=1))
        >>> # With `total_delay=True`, the calculated delay should be equal to
        >>> # the signal delay plus the begin time difference:
        >>> calculated_delay == signal_delay + (seis2.begin_time - seis1.begin_time)
        True
        >>>
        ```

        To demonstrate the `abs_max` parameter, the second seismogram's data
        is sign-flipped:

        ```python
        >>> seis2.data *= -1
        >>> calculated_delay, cc = delay(seis1, seis2)
        >>> # Without `abs_max=True`, the calculated delay is no longer equal
        >>> # to the true signal delay (as expected):
        >>> calculated_delay == signal_delay
        False
        >>> # The normalised cross-correlation value is also not very high
        >>> cc
        np.float64(0.5094230)
        >>>
        >>> calculated_delay, cc = delay(seis1, seis2, abs_max=True)
        >>> # with `abs_max=True`, the signal delay is again retrieved:
        >>> calculated_delay == signal_delay
        True
        >>> # And, as the signals are completely opposite, the normalised
        >>> # cross-correlation value is -1:
        >>> np.testing.assert_approx_equal(cc, -1)
        >>>
        ```
    """

    _check_same_delta(seismogram1, seismogram2)

    if max_shift is not None and len(seismogram1.data) != len(seismogram2.data):
        raise ValueError(
            "Input seismograms must be of equal length when using `max_shift`."
        )

    data1, data2 = seismogram1.data, seismogram2.data
    delta = seismogram1.delta

    if max_shift is not None:
        max_lag_in_samples = math.ceil(max_shift / delta)
        data1 = np.pad(data1, max_lag_in_samples)
        corr = _correlate(data1, data2, mode="valid")
    else:
        corr = _correlate(data1, data2, mode="full")

    corr_index = np.argmax(corr)

    if abs_max and np.max(corr) < -1 * np.min(corr):
        corr_index = np.argmin(corr)

    if max_shift is not None:
        shift = -int(corr_index - max_lag_in_samples)
    else:
        shift = int(len(data2) - 1 - corr_index)

    delay = shift * delta

    # find overlapping parts of seismograms after alignment
    if shift < 0:
        data1 = data1[-shift:]
    else:
        data2 = data2[shift:]
    if len(data1) > len(data2):
        data1 = data1[: len(data2)]
    else:
        data2 = data2[: len(data1)]

    cc, _ = _pearsonr(data1, data2)

    if total_delay:
        delay += seismogram2.begin_time - seismogram1.begin_time

    return delay, cc

def differentiate[T: Seismogram](seismogram: T, clone: bool = False) -> T | None:
    r"""Differentiate a seismogram in the frequency domain.

    Multiplies the FFT of `seismogram.data` by $i\omega$ at each
    [`rfftfreq`][numpy.fft.rfftfreq] bin, then inverse transforms back to the
    time domain. The DC bin is correctly zeroed ($i\omega = 0$), since a
    constant offset differentiates to zero.

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

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

    Raises:
        ValueError: If `seismogram.data` is empty.

    Examples:
        A synthetic sine wave with a known closed-form derivative
        ($\omega \cos(\omega t)$) is used to verify the result:

        ```python
        >>> import numpy as np
        >>> import pandas as pd
        >>> from pysmo import MiniSeismogram
        >>> from pysmo.tools.signal import differentiate
        >>> dt = 0.1
        >>> npts = 250  # an exact number of cycles of the 1 Hz signal below
        >>> t = np.arange(npts) * dt
        >>> omega = 2 * np.pi * 1.0
        >>> seismogram = MiniSeismogram(
        ...     begin_time=pd.Timestamp("2010-02-27T06:30:00Z"),
        ...     delta=pd.Timedelta(seconds=dt),
        ...     data=np.sin(omega * t),
        ... )
        >>> velocity = differentiate(seismogram, clone=True)
        >>> expected = omega * np.cos(omega * t)
        >>> np.allclose(velocity.data, expected, atol=1e-6)
        True
        >>>
        ```
    """
    if len(seismogram.data) == 0:
        raise ValueError("Cannot differentiate an empty seismogram.")

    if clone:
        seismogram = deepcopy(seismogram)

    dt = seismogram.delta.total_seconds()
    npts = len(seismogram.data)
    freqs = np.fft.rfftfreq(npts, d=dt)
    omega = 2 * np.pi * freqs

    spectrum = np.fft.rfft(seismogram.data)
    spectrum *= 1j * omega
    seismogram.data = np.fft.irfft(spectrum, n=npts)

    return seismogram if clone else None

@register_filter
def envelope[T: Seismogram](
    seismogram: T, fc: float, alpha: float, clone: bool = False
) -> T | None:
    """Calculates the envelope of a gaussian filtered seismogram.

    Args:
        seismogram: Seismogram object.
        fc: Centre frequency of the Gaussian filter (in Hz).
        alpha: Dimensionless shape parameter controlling the filter width.
            Larger values produce a narrower (more selective) filter.
        clone: Operate on a clone of the input seismogram.

    Returns:
        Seismogram containing the envelope.

    Examples:
        ```python
        >>> from pysmo.classes import SAC
        >>> from pysmo.tools.signal import envelope
        >>> seis = SAC.from_file("example.sac").seismogram
        >>> fc = 0.02 # Centre Gaussian filter at 0.02 Hz (50s period)
        >>> alpha = 50 # Set alpha (which determines filterwidth) to 50
        >>> envelope_seis = envelope(seis, fc, alpha, clone=True)
        >>>
        ```
    """
    if clone:
        seismogram = deepcopy(seismogram)
    seismogram.data = _gauss(seismogram, fc, alpha)[0]
    return seismogram if clone else None

def filter[T: Seismogram](
    seismogram: T,
    filter_name: FilterName,
    clone: bool = False,
    **filter_options: bool | int | float,
) -> T | None:
    """Apply a specified filter to the input seismogram.

    This function is a convenience wrapper that calls other filters in this module.

    Args:
        seismogram: The input seismogram to be filtered.
        filter_name: The type of filter to apply.
        clone: If `True`, return a new Seismogram object with the filtered
            data. If `False`, modify the input seismogram in place.
        **filter_options: Filter parameters passed to the specified filter
            function.

    Returns:
        A new Seismogram object containing the filtered data when called with `clone=True`.

    Raises:
        ValueError: If `filter_name` is not a registered filter.

    Examples:
        ```python
        >>> from pysmo.classes import SAC
        >>> from pysmo.tools.signal import filter
        >>> seis = SAC.from_file("example.sac").seismogram
        >>>
        >>> # create a new filtered seismogram with a lowpass filter
        >>> filtered_seis = filter(seis, "lowpass", freqmax=0.5, clone=True)
        >>>
        >>> # or update in place with a bandpass filter
        >>> filter(seis, "bandpass", freqmin=0.1, freqmax=0.5)
        >>>
        ```
    """

    try:
        filter_func = _FILTER_REGISTRY[filter_name]
    except KeyError:
        # This fallback handles cases where FilterName is updated but
        # the function isn't decorated yet.
        valid_filters = ", ".join(_FILTER_REGISTRY.keys())
        raise ValueError(
            f"Filter '{filter_name}' is not registered. Available: {valid_filters}"
        )

    if clone:
        return filter_func(seismogram, clone=True, **filter_options)
    filter_func(seismogram, clone=False, **filter_options)
    return None

@register_filter
def gauss[T: Seismogram](
    seismogram: T, fc: float, alpha: float, clone: bool = False
) -> T | None:
    """Returns a gaussian filtered seismogram.

    Args:
        seismogram: Seismogram object.
        fc: Centre frequency of the Gaussian filter (in Hz).
        alpha: Dimensionless shape parameter controlling the filter width.
            Larger values produce a narrower (more selective) filter.
        clone: Operate on a clone of the input seismogram.

    Returns:
        Gaussian filtered seismogram.

    Examples:
        ```python
        >>> from pysmo.classes import SAC
        >>> from pysmo.tools.signal import gauss
        >>> seis = SAC.from_file("example.sac").seismogram
        >>> fc = 0.02 # Centre Gaussian filter at 0.02 Hz (50s period)
        >>> alpha = 50 # Set alpha (which determines filterwidth) to 50
        >>> gauss_seis = gauss(seis, fc, alpha, clone=True)
        >>>
        ```
    """
    if clone:
        seismogram = deepcopy(seismogram)
    seismogram.data = _gauss(seismogram, fc, alpha)[1]
    return seismogram if clone else None

@register_filter
def highpass[T: Seismogram](
    seismogram: T,
    freqmin: float = 0.1,
    corners: int = 2,
    zerophase: bool = False,
    clone: bool = False,
) -> T | None:
    """Apply a highpass filter to the input seismogram.

    Args:
        seismogram: The input seismogram to be filtered.
        freqmin: The minimum frequency of the highpass filter (in Hz).
        corners: The number of corners (poles) for the Butterworth filter.
        zerophase: If `True`, apply the filter in both forward and reverse
            directions to achieve zero phase distortion.
        clone: If `True`, return a new Seismogram object with the filtered
            data. If `False`, modify the input seismogram in place.

    Returns:
        A new Seismogram object containing the filtered data when called with `clone=True`.
    """
    fe = 0.5 / seismogram.delta.total_seconds()
    low = freqmin / fe

    if not (0 < low < 1):
        raise ValueError(
            f"freqmin ({freqmin}) is invalid for sampling rate {1 / seismogram.delta.total_seconds()} Hz."
        )

    sos = iirfilter(corners, low, btype="high", ftype="butter", output="sos")

    if clone:
        seismogram = deepcopy(seismogram)

    if zerophase:
        seismogram.data = sosfiltfilt(sos, seismogram.data)
    else:
        seismogram.data = sosfilt(sos, seismogram.data)

    return seismogram if clone else None

def integrate[T: Seismogram](seismogram: T, clone: bool = False) -> T | None:
    r"""Integrate a seismogram in the frequency domain.

    Divides the FFT of `seismogram.data` by $i\omega$ at each
    [`rfftfreq`][numpy.fft.rfftfreq] bin ($\omega > 0$), then inverse
    transforms back to the time domain. The DC bin is set to `0.0` rather
    than divided by. Working in the frequency domain avoids the unbounded
    low-frequency drift a cumulative time-domain integrator introduces.

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

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

    Raises:
        ValueError: If `seismogram.data` is empty.

    Examples:
        A synthetic cosine wave with a known closed-form integral
        ($\sin(\omega t)$) is used to verify the result:

        ```python
        >>> import numpy as np
        >>> import pandas as pd
        >>> from pysmo import MiniSeismogram
        >>> from pysmo.tools.signal import integrate
        >>> dt = 0.1
        >>> npts = 250  # an exact number of cycles of the 1 Hz signal below
        >>> t = np.arange(npts) * dt
        >>> omega = 2 * np.pi * 1.0
        >>> seismogram = MiniSeismogram(
        ...     begin_time=pd.Timestamp("2010-02-27T06:30:00Z"),
        ...     delta=pd.Timedelta(seconds=dt),
        ...     data=omega * np.cos(omega * t),
        ... )
        >>> displacement = integrate(seismogram, clone=True)
        >>> expected = np.sin(omega * t)
        >>> np.allclose(displacement.data, expected, atol=1e-6)
        True
        >>>
        ```
    """
    if len(seismogram.data) == 0:
        raise ValueError("Cannot integrate an empty seismogram.")

    if clone:
        seismogram = deepcopy(seismogram)

    dt = seismogram.delta.total_seconds()
    npts = len(seismogram.data)
    freqs = np.fft.rfftfreq(npts, d=dt)
    omega = 2 * np.pi * freqs

    spectrum = np.fft.rfft(seismogram.data)
    integrated_spectrum = np.zeros_like(spectrum)
    integrated_spectrum[1:] = spectrum[1:] / (1j * omega[1:])
    seismogram.data = np.fft.irfft(integrated_spectrum, n=npts)

    return seismogram if clone else None

@register_filter
def lowpass[T: Seismogram](
    seismogram: T,
    freqmax: float = 0.5,
    corners: int = 2,
    zerophase: bool = False,
    clone: bool = False,
) -> T | None:
    """Apply a lowpass filter to the input seismogram.

    Args:
        seismogram: The input seismogram to be filtered.
        freqmax: The maximum frequency of the lowpass filter (in Hz).
        corners: The number of corners (poles) for the Butterworth filter.
        zerophase: If `True`, apply the filter in both forward and reverse
            directions to achieve zero phase distortion.
        clone: If `True`, return a new Seismogram object with the filtered
            data. If `False`, modify the input seismogram in place.

    Returns:
        A new Seismogram object containing the filtered data when called with `clone=True`.
    """
    fe = 0.5 / seismogram.delta.total_seconds()
    high = freqmax / fe

    if not (0 < high < 1):
        raise ValueError(
            f"freqmax ({freqmax}) is invalid for sampling rate {1 / seismogram.delta.total_seconds()} Hz."
        )

    sos = iirfilter(corners, high, btype="low", ftype="butter", output="sos")

    if clone:
        seismogram = deepcopy(seismogram)

    if zerophase:
        seismogram.data = sosfiltfilt(sos, seismogram.data)
    else:
        seismogram.data = sosfilt(sos, seismogram.data)

    return seismogram if clone else None