# SPDX-FileCopyrightText: Copyright © 2025 Damian Anders <damian.anders@tu-dresden.de>
#
# SPDX-License-Identifier: AGPL-3.0-or-later

"""Functions for calculating the (inverse) Fourier transform with correct normalisation for real signals."""

# abgeleitet aus commit c1b69177062af7b1b97ebb4e8fafb01651e3abad

import numpy as np
import numpy.typing as npt
from scipy import signal as ssig

__all__ = ["rfft_tukey_corrected", "irfft_corrected"]

def rfft_tukey_corrected(x: npt.NDArray[np.float64],
						 fs: float,
						 axis: int = 0,
						 tukey: float = 0.05,
) -> tuple[np.ndarray, np.ndarray]:
	"""
	Applies a Tukey window to the input signal along the specified axis, performs a normalized
	real-valued fast Fourier transform (RFFT), and corrects the spectrum for a real-valued input
	signal. The function also computes the corresponding frequency values.

	Parameters
	----------
	x : ndarray
		Input signal data array.
	fs : float
		Sampling frequency of the input signal.
	axis : int, default=0
		The axis over which the computation is performed. Default is 0.
	tukey : float, default=0.05
		The alpha parameter for the Tukey window. Determines the fraction of the window inside
		the cosine taper. Default is 0.05.

	Returns
	-------
	tuple of (np.ndarray, np.ndarray)
		- The first element is the spectrum of the input signal after applying the Tukey window and
		  performing the RFFT.
		- The second element is the array of frequency values corresponding to the computed spectrum.
	"""
	shape_to = np.ones(x.ndim, dtype=np.int_)	# benötigte shape des Tukey-Fenster-Vektors
	shape_to[axis] = -1

	if tukey > 0:
		fenster = ssig.windows.tukey(x.shape[axis], alpha=tukey, sym=False	# Tukey-Fenster-Vektor erstellen
									 ).reshape(tuple(shape_to))
	else:
		fenster = 1.0

	spektrum =  np.fft.rfft(x * fenster, axis=axis, norm='forward')	# rFFT berechnen

	spektrum = _helperNormalization(spektrum, axis, inverse=False)	# für reelles Eingangssignal normieren

	freq = np.fft.rfftfreq(x.shape[axis], d=1/fs)	# Frequenzlinienvektor erstellen

	return spektrum, freq


def irfft_corrected(spektrum: npt.NDArray[np.complex128],
					axis: int = 0,
) -> np.ndarray:
	"""
	Computes the inverse real-to-complex FFT on the input spectrum after reversing the normalization.

	This function first handles the reversal of normalization on the input spectrum and then applies
	the inverse FFT (`irfft`) operation along the specified axis. The normalization reversal ensures
	that the input spectrum is correctly prepared before applying the FFT transformation.

	Parameters
	----------
	spektrum : npt.NDArray[np.complex128]
	    The input frequency-domain spectrum, represented as a NumPy array of complex numbers.

	axis : int, optional
	    The axis along which the reverse FFT is performed. The default is 0.

	Returns
	-------
	np.ndarray
	    The real-valued time-domain signal array resulting from the inverse FFT transformation.
	"""
	spektrum = _helperNormalization(spektrum, axis, inverse=True)	# Normierung rückgängig machen

	return np.fft.irfft(spektrum, norm='forward', axis=axis)


def _helperNormalization(spektrum: npt.NDArray[np.complex128],
						 axis: int,
						 inverse: bool = False,
) -> np.ndarray:
	"""
	Handles the normalization of a spectrum along a specified axis. The function modifies
	the spectrum differently depending on whether it is being prepared for an inverse Fast
	Fourier Transform (iFFT) or a forward FFT. The axis of normalization is temporarily
	moved to axis 0 during the operation and restored afterward.

	Parameters
	----------
	spektrum : numpy.ndarray
	    The spectrum to be normalized, typically in the form of a multi-dimensional
	    array. Its contents and shape will determine the behavior of the normalization.
	axis : int
	    The axis along which the normalization is applied. This axis will be temporarily
	    moved to axis 0 during the normalization process.
	inverse : bool, optional
	    A flag indicating whether the normalization is being undone for iFFT
	    (True) or applied for a forward FFT (False). By default, this is False.

	Returns
	-------
	numpy.ndarray
	    The normalized spectrum with its original axis order restored.
	"""
	spektrum = np.moveaxis(spektrum, axis, 0)	# zu normierende Achse zu Achse 0 machen
	if inverse:
		normFaktor = 0.5	# Normierung für iFFT rückgängig machen
	else:
		normFaktor = 2	# Spektrum normieren

	spektrum[1:-1] = normFaktor * spektrum[1:-1]
	spektrum = np.moveaxis(spektrum, 0, axis)  # und Achsen zurücktauschen

	return spektrum