"""Helpers to utilize existing stft / istft tests for testing `ShortTimeFFT`.

This module provides the functions stft_compare() and istft_compare(), which,
compares the output between the existing (i)stft() and the shortTimeFFT based
_(i)stft_wrapper() implementations in this module.

For testing add the following imports to the file ``tests/test_spectral.py``::

    from ._scipy_spectral_test_shim import stft_compare as stft
    from ._scipy_spectral_test_shim import istft_compare as istft

and remove the existing imports of stft and istft.

The idea of these wrappers is not to provide a backward-compatible interface
but to demonstrate that the ShortTimeFFT implementation is at least as capable
as the existing one and delivers comparable results. Furthermore, the
wrappers highlight the different philosophies of the implementations,
especially in the border handling.
"""
import platform
from typing import cast, Literal

import numpy as np
from numpy.testing import assert_allclose

from scipy.signal import ShortTimeFFT
from scipy.signal import get_window, stft, istft
from scipy.signal._arraytools import const_ext, even_ext, odd_ext, zero_ext
from scipy.signal._short_time_fft import FFT_MODE_TYPE
from scipy.signal._spectral_py import _triage_segments


def _stft_wrapper(x, fs=1.0, window='hann', nperseg=256, noverlap=None,
                  nfft=None, detrend=False, return_onesided=True,
                  boundary='zeros', padded=True, axis=-1, scaling='spectrum'):
    """Wrapper for the SciPy `stft()` function based on `ShortTimeFFT` for
    unit testing.

    Handling the boundary and padding is where `ShortTimeFFT` and `stft()`
    differ in behavior. Parts of `_spectral_helper()` were copied to mimic
    the` stft()` behavior.

    This function is meant to be solely used by `stft_compare()`.
    """
    if scaling not in ('psd', 'spectrum'):  # same errors as in original stft:
        raise ValueError(f"Parameter {scaling=} not in ['spectrum', 'psd']!")

    # The following lines are taken from the original _spectral_helper():
    boundary_funcs = {'even': even_ext,
                      'odd': odd_ext,
                      'constant': const_ext,
                      'zeros': zero_ext,
                      None: None}

    if boundary not in boundary_funcs:
        raise ValueError(f"Unknown boundary option '{boundary}', must be one" +
                         f" of: {list(boundary_funcs.keys())}")
    if x.size == 0:
        return np.empty(x.shape), np.empty(x.shape), np.empty(x.shape)

    if nperseg is not None:  # if specified by user
        nperseg = int(nperseg)
        if nperseg < 1:
            raise ValueError('nperseg must be a positive integer')

    # parse window; if array like, then set nperseg = win.shape
    win, nperseg = _triage_segments(window, nperseg,
                                    input_length=x.shape[axis])

    if nfft is None:
        nfft = nperseg
    elif nfft < nperseg:
        raise ValueError('nfft must be greater than or equal to nperseg.')
    else:
        nfft = int(nfft)

    if noverlap is None:
        noverlap = nperseg//2
    else:
        noverlap = int(noverlap)
    if noverlap >= nperseg:
        raise ValueError('noverlap must be less than nperseg.')
    nstep = nperseg - noverlap
    n = x.shape[axis]

    # Padding occurs after boundary extension, so that the extended signal ends
    # in zeros, instead of introducing an impulse at the end.
    # I.e. if x = [..., 3, 2]
    # extend then pad -> [..., 3, 2, 2, 3, 0, 0, 0]
    # pad then extend -> [..., 3, 2, 0, 0, 0, 2, 3]

    if boundary is not None:
        ext_func = boundary_funcs[boundary]
        # Extend by nperseg//2 in front and back:
        x = ext_func(x, nperseg//2, axis=axis)

    if padded:
        # Pad to integer number of windowed segments
        # I.e make x.shape[-1] = nperseg + (nseg-1)*nstep, with integer nseg
        x = np.moveaxis(x, axis, -1)

        # This is an edge case where shortTimeFFT returns one more time slice
        # than the Scipy stft() shorten to remove last time slice:
        if n % 2 == 1 and nperseg % 2 == 1 and noverlap % 2 == 1:
            x = x[..., : -1]

        nadd = (-(x.shape[-1]-nperseg) % nstep) % nperseg
        zeros_shape = list(x.shape[:-1]) + [nadd]
        x = np.concatenate((x, np.zeros(zeros_shape)), axis=-1)
        x = np.moveaxis(x, -1, axis)

    #  ... end original _spectral_helper() code.
    scale_to = {'spectrum': 'magnitude', 'psd': 'psd'}[scaling]

    if np.iscomplexobj(x) and return_onesided:
        return_onesided = False
    # using cast() to make mypy happy:
    fft_mode = cast(FFT_MODE_TYPE, 'onesided' if return_onesided else 'twosided')

    ST = ShortTimeFFT(win, nstep, fs, fft_mode=fft_mode, mfft=nfft,
                      scale_to=scale_to, phase_shift=None)

    k_off = nperseg // 2
    p0 = 0  # ST.lower_border_end[1] + 1
    nn = x.shape[axis] if padded else n+k_off+1
    # number of frames akin to legacy stft computation
    p1 = (x.shape[axis] - nperseg) // nstep + 1 

    detr = None if detrend is False else detrend
    Sxx = ST.stft_detrend(x, detr, p0, p1, k_offset=k_off, axis=axis)
    t = ST.t(nn, 0, p1 - p0, k_offset=0 if boundary is not None else k_off)
    if x.dtype in (np.float32, np.complex64):
        Sxx = Sxx.astype(np.complex64)

    return ST.f, t, Sxx


def _istft_wrapper(Zxx, fs=1.0, window='hann', nperseg=None, noverlap=None,
                   nfft=None, input_onesided=True, boundary=True, time_axis=-1,
                   freq_axis=-2, scaling='spectrum') -> \
        tuple[np.ndarray, np.ndarray, tuple[int, int]]:
    """Wrapper for the SciPy `istft()` function based on `ShortTimeFFT` for
        unit testing.

    Note that only option handling is implemented as far as to handle the unit
    tests. E.g., the case ``nperseg=None`` is not handled.

    This function is meant to be solely used by `istft_compare()`.
    """
    # *** Lines are taken from _spectral_py.istft() ***:
    if Zxx.ndim < 2:
        raise ValueError('Input stft must be at least 2d!')

    if freq_axis == time_axis:
        raise ValueError('Must specify differing time and frequency axes!')

    nseg = Zxx.shape[time_axis]

    if input_onesided:
        # Assume even segment length
        n_default = 2*(Zxx.shape[freq_axis] - 1)
    else:
        n_default = Zxx.shape[freq_axis]

    # Check windowing parameters
    if nperseg is None:
        nperseg = n_default
    else:
        nperseg = int(nperseg)
        if nperseg < 1:
            raise ValueError('nperseg must be a positive integer')

    if nfft is None:
        if input_onesided and (nperseg == n_default + 1):
            # Odd nperseg, no FFT padding
            nfft = nperseg
        else:
            nfft = n_default
    elif nfft < nperseg:
        raise ValueError('nfft must be greater than or equal to nperseg.')
    else:
        nfft = int(nfft)

    if noverlap is None:
        noverlap = nperseg//2
    else:
        noverlap = int(noverlap)
    if noverlap >= nperseg:
        raise ValueError('noverlap must be less than nperseg.')
    nstep = nperseg - noverlap

    # Get window as array
    if isinstance(window, str) or type(window) is tuple:
        win = get_window(window, nperseg)
    else:
        win = np.asarray(window)
        if len(win.shape) != 1:
            raise ValueError('window must be 1-D')
        if win.shape[0] != nperseg:
            raise ValueError(f'window must have length of {nperseg}')

    outputlength = nperseg + (nseg-1)*nstep
    # *** End block of: Taken from _spectral_py.istft() ***

    # Using cast() to make mypy happy:
    fft_mode = cast(FFT_MODE_TYPE, 'onesided' if input_onesided else 'twosided')
    scale_to = cast(Literal['magnitude', 'psd'],
                    {'spectrum': 'magnitude', 'psd': 'psd'}[scaling])

    ST = ShortTimeFFT(win, nstep, fs, fft_mode=fft_mode, mfft=nfft,
                      scale_to=scale_to, phase_shift=None)

    if boundary:
        j = nperseg if nperseg % 2 == 0 else nperseg - 1
        k0 = ST.k_min + nperseg // 2
        k1 = outputlength - j + k0
    else:
        raise NotImplementedError("boundary=False does not make sense with" +
                                  "ShortTimeFFT.istft()!")

    x = ST.istft(Zxx, k0=k0, k1=k1, f_axis=freq_axis, t_axis=time_axis)
    t = np.arange(k1 - k0) * ST.T
    k_hi = ST.upper_border_begin(k1 - k0)[0]
    # using cast() to make mypy happy:
    return t, x, (ST.lower_border_end[0], k_hi)


def stft_compare(x, fs=1.0, window='hann', nperseg=256, noverlap=None,
                 nfft=None, detrend=False, return_onesided=True,
                 boundary='zeros', padded=True, axis=-1, scaling='spectrum'):
    """Assert that the results from the existing `stft()` and `_stft_wrapper()`
    are close to each other.

    For comparing the STFT values an absolute tolerance of the floating point
    resolution was added to circumvent problems with the following tests:
    * For float32 the tolerances are much higher in
      TestSTFT.test_roundtrip_float32()).
    * The TestSTFT.test_roundtrip_scaling() has a high relative deviation.
      Interestingly this did not appear in Scipy 1.9.1 but only in the current
      development version.
    """
    kw = dict(x=x, fs=fs, window=window, nperseg=nperseg, noverlap=noverlap,
              nfft=nfft, detrend=detrend, return_onesided=return_onesided,
              boundary=boundary, padded=padded, axis=axis, scaling=scaling)
    f, t, Zxx = stft(**kw)
    f_wrapper, t_wrapper, Zxx_wrapper = _stft_wrapper(**kw)

    e_msg_part = " of `stft_wrapper()` differ from `stft()`."
    assert_allclose(f_wrapper, f, err_msg=f"Frequencies {e_msg_part}")
    assert_allclose(t_wrapper, t, err_msg=f"Time slices {e_msg_part}")

    # Adapted tolerances to account for:
    atol = np.finfo(Zxx.dtype).resolution * 2
    assert_allclose(Zxx_wrapper, Zxx, atol=atol,
                    err_msg=f"STFT values {e_msg_part}")
    return f, t, Zxx


def istft_compare(Zxx, fs=1.0, window='hann', nperseg=None, noverlap=None,
                  nfft=None, input_onesided=True, boundary=True, time_axis=-1,
                  freq_axis=-2, scaling='spectrum'):
    """Assert that the results from the existing `istft()` and
    `_istft_wrapper()` are close to each other.

    Quirks:
    * If ``boundary=False`` the comparison is skipped, since it does not
      make sense with ShortTimeFFT.istft(). Only used in test
      TestSTFT.test_roundtrip_boundary_extension().
    * If ShortTimeFFT.istft() decides the STFT is not invertible, the
      comparison is skipped, since istft() only emits a warning and does not
      return a correct result. Only used in
      ShortTimeFFT.test_roundtrip_not_nola().
    * For comparing the signals an absolute tolerance of the floating point
      resolution was added to account for the low accuracy of float32 (Occurs
      only in TestSTFT.test_roundtrip_float32()).
    """
    kw = dict(Zxx=Zxx, fs=fs, window=window, nperseg=nperseg,
              noverlap=noverlap, nfft=nfft, input_onesided=input_onesided,
              boundary=boundary, time_axis=time_axis, freq_axis=freq_axis,
              scaling=scaling)

    t, x = istft(**kw)
    if not boundary:  # skip test_roundtrip_boundary_extension():
        return t, x  # _istft_wrapper does() not implement this case
    try:  # if inversion fails, istft() only emits a warning:
        t_wrapper, x_wrapper, (k_lo, k_hi) = _istft_wrapper(**kw)
    except ValueError as v:  # Do nothing if inversion fails:
        if v.args[0] == "Short-time Fourier Transform not invertible!":
            return t, x
        raise v

    e_msg_part = " of `istft_wrapper()` differ from `istft()`"
    assert_allclose(t, t_wrapper, err_msg=f"Sample times {e_msg_part}")

    # Adapted tolerances to account for resolution loss:
    atol = np.finfo(x.dtype).resolution*2  # instead of default atol = 0
    rtol = 1e-7  # default for np.allclose()

    # Relax atol on 32-Bit platforms a bit to pass CI tests.
    #  - Not clear why there are discrepancies (in the FFT maybe?)
    #  - Not sure what changed on 'i686' since earlier on those test passed
    if x.dtype == np.float32 and platform.machine() == 'i686':
        # float32 gets only used by TestSTFT.test_roundtrip_float32() so
        # we are using the tolerances from there to circumvent CI problems
        atol, rtol = 1e-4, 1e-5
    elif platform.machine() in ('aarch64', 'i386', 'i686'):
        atol = max(atol, 1e-12)  # 2e-15 seems too tight for 32-Bit platforms

    assert_allclose(x_wrapper[k_lo:k_hi], x[k_lo:k_hi], atol=atol, rtol=rtol,
                    err_msg=f"Signal values {e_msg_part}")
    return t, x