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

"""Signal resampling module for frequency harmonic alignment."""

import numpy as np
from scipy import interpolate as sip, signal as ssig
from numpy import typing as npt

__all__ = ["resample_to_harmonics"]

def resample_to_harmonics(x_vector: npt.NDArray[np.float64],
						  y_array: npt.NDArray[np.float64],
						  f_sampling: float,
						  f_base: float,
						  tukey_alpha: float,
) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], np.float64, np.int_]:
	"""
	Resample and pad a signal so that harmonics of a base frequency align with frequency bins.

	This function computes a new sampling rate and determines the necessary zero-padding
	length so that the harmonics of a specified base frequency (`f_base`) fall exactly on
	the frequency lines of a subsequent FFT. The function utilizes cubic spline interpolation
	for resampling and applies a Tukey window before padding to ensure a smooth transition.
	The primary focus of this technique is on padding the signal and changing the original
	sampling rate as little as possible. The sampling rate is only ever reduced, never
	increased.

	Parameters
	----------
	x_vector : npt.NDArray[np.float64]
		The input time vector array of the signal.
	y_array : npt.NDArray[np.float64]
		The signal data array corresponding to the time vector. Can be a 1D or 2D array.
	f_sampling : float
		The original sampling frequency of the input data.
	f_base : float
		The base frequency (e.g., rotational frequency) whose harmonics are to be aligned.
	tukey_alpha : float
		The shape parameter of the Tukey window applied to the data to ensure a smooth
		transition into the zero-padding.

	Returns
	-------
	y_array_resampled : npt.NDArray[np.float64]
		The resampled, windowed, and zero-padded signal data.
	x_array_resampled : npt.NDArray[np.float64]
		The new time vector array corresponding to the resampled and padded data.
	f_s_resampled : np.float64
		The new sampling frequency of the resampled data.
	steps : np.int_
		The number of frequency steps (bins) between two harmonics of the base frequency.

	Raises
	------
	ValueError
		If `x_vector` contains data in more than one axis.
		If `y_array` is not a 1D or 2D array.
		If `x_vector` and `y_array` do not match in length.
		If `tukey_alpha` is not within the interval [0, 1].
		If `f_base` or `f_sampling` are not strictly positive.
		If `f_base` is greater than or equal to 0.5 * `f_sampling`.
	"""
	if np.sum(np.asarray(x_vector.shape) > 1) > 1:
		raise ValueError("Only one axis of \"x_vector\" may contain data.")
	if y_array.ndim > 2:
		raise ValueError("\"y_array\" must be a 1D or 2D array.")
	if not np.any(np.asarray(y_array.shape) == x_vector.size):
		raise ValueError("\"x_vector\" and \"y_array\" do not match in length.")
	if (tukey_alpha < 0) or (tukey_alpha > 1):
		raise ValueError("\"tukey_alpha\" must be in the interval [0,1].")
	if (f_base <= 0) or (f_sampling <= 0):
		raise ValueError("\"f_base\" and \"f_sampling\" must be positive.")
	if f_base >= (f_sampling / 2):
		raise ValueError("\"f_base\" must be less than 0.5*\"f_sampling\".")

	x_vector = np.ravel(x_vector)
	y_array = np.atleast_2d(y_array)
	if y_array.shape[1] > y_array.shape[0]:
		y_array = y_array.T

	anz_daten = y_array.shape[0]
	d_f = f_sampling / anz_daten	# Frequenzauflösung

	steps = np.int_(np.floor(f_base / d_f))	# Mindestanzahl der Frequenzlinien zwischen zwei Harmonischen der Drehfrequenz
	steps = steps + 1	# ein weiterer voller Frequenzschritt soll zwischen zwei Harmonischen der Drehfrequenz liegen

	padding = np.int_(np.ceil(steps * f_sampling / f_base - anz_daten))	# Länge des Zeitsignal-Paddings
	if np.mod(anz_daten + padding, 2):
		padding = padding + 1

	# Abtastrate so berechnen, dass Harmonische der Drehfrequenz immer auf eine Frequenzlinie fallen
	f_s_resampled = f_base * (anz_daten + padding) / steps

	# Fenster erstellen, damit kein Sprung zu Padding
	if 0 < tukey_alpha <= 1:
		fenster = np.reshape(ssig.windows.tukey(anz_daten, tukey_alpha, sym=False), (-1, 1))
	else:	# kann durch Fehlerbehandlung nur noch 0 sein
		fenster = 1

	# Resampling ausführen:
	x_array_resampled = np.arange(anz_daten + padding) / f_s_resampled + x_vector[0]	# neuer Zeitlinienvektor
	interp_signal = sip.CubicSpline(x_vector, y_array, axis=0, bc_type='natural')	# Interpolationsfunktion
	y_array_resampled = interp_signal(x_array_resampled[:anz_daten]) * fenster	# interpolierte Daten

	# Padding hinzufügen
	y_array_resampled = np.pad(y_array_resampled, ((0, padding), (0, 0)))

	return y_array_resampled, x_array_resampled, f_s_resampled, steps