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

"""
Analyze eccentric motion data measured over time using numerical and signal processing techniques.

This module provides tools for fitting geometric models to shaft runout and eccentricity data.
It utilizes differential evolution optimization to estimate rotational frequency,
eccentricity magnitude, and sensor configurations from multi-channel displacement signals.

Key Features
------------
- Automatic boundary estimation for optimizer parameters using spectral analysis.
- Support for exact geometric models ('exc') and simplified cosine models ('cos').
- Integrated signal conditioning including Tukey windowing and lowpass filtering.
- Support for multi-sensor setups with relative angular position estimation.

Mathematical Models
-------------------
The module supports two primary models for sensor displacement $\\Delta s$:
1.  **Exact ('exc')**: Accounts for the geometric projection of eccentricity on a circular
	path, relevant when eccentricity is a significant fraction of the shaft radius.
2.  **Cosine ('cos')**: A first-order approximation where displacement is modeled as a
	pure sinusoid, suitable for small eccentricities.
"""

import logging
from typing import Literal

import numba as nb
import numpy as np
import numpy.typing as npt
from scipy import optimize as sop

from ._initial_guess import *
from ._protocols import FitFunction
from .rfft_corrected import irfft_corrected, rfft_tukey_corrected
from ._set_position_zero import _set_position_zero

logging.basicConfig(format='%(levelname)s:%(message)s', level=logging.WARNING)

__all__ = ["fit_eccentricity"]

def fit_eccentricity(time_values: npt.NDArray[np.float64],
					 signal: npt.NDArray[np.float64],
					 radius: float = 1.0,
					 fit_function: FitFunction | Literal['cos', 'exc'] = 'exc',
					 bounds: sop.Bounds | None = None,
					 angle_tol: float = 5,
					 position_zero: bool = True,
					 tukey_alpha: float = 0.2,
					 f_lowpass: float | None = None,
) -> sop.OptimizeResult:
	"""
	Fit a mathematical model to eccentric motion data using differential evolution.

	This function estimates rotational frequency, eccentricity magnitude, and sensor-specific
	parameters (angular positions and offsets) by minimizing the squared error between
	measured signals and a geometric model. It supports both an exact eccentric geometry
	model and a simplified cosine approximation.

	Parameters
	----------
	time_values : npt.NDArray[np.float64]
		One-dimensional array of time stamps for the measurements. Must be uniformly sampled.
	signal : npt.NDArray[np.float64]
		Measurement data. Can be a 1D array (single channel) or a 2D array where each
		column represents a different sensor channel.
	radius : float, default=1.0
		The nominal radius of the rotating shaft/object, used as a constant in the
		'exc' fitness function.
	fit_function : {'exc', 'cos'} or FitFunction, default='exc'
		The model to fit. 'exc' uses exact eccentric geometry; 'cos' uses a
		cosine approximation. Also accepts a custom callable following the
		`FitFunction` protocol.
	bounds : scipy.optimize.Bounds or None, default=None
		Search space boundaries for the optimizer. If None, boundaries are
		automatically estimated via `_bounds_runoutFit`.
	angle_tol : float, default=5
		Half-width of the angular search window around the estimated sensor
		positions, in degrees.
	position_zero : bool, default=True
		If True, shifts the optimized angular positions so the first sensor
		is at 0 radians, making others relative to the first.
	tukey_alpha : float, default=0.2
		The alpha parameter for the Tukey window applied during initial spectral
		analysis and f_lowpass filtering.
	f_lowpass : float or None, default=None
		Optional cutoff frequency in Hz. If provided, the signal is filtered in
		the frequency domain before fitting.

	Returns
	-------
	scipy.optimize.OptimizeResult
		The optimization result. The `x` attribute contains:
		`[fRot, eccentricity, angle_s1, offset_s1, ..., angle_sN, offset_sN]`.

	See Also
	--------
	_bounds_runoutFit : Internal helper for boundary estimation.
	_fitness_exzenter : The exact geometric fitness function.
	_fitness_cos : The simplified cosine fitness function.
	"""
	# res.x = [fRot, betragExzentrizitaet, winkelSensor1, offsetSensor1, ..., winkelSensorN, offsetSensorN]

	if np.squeeze(signal).ndim == 1:				# wenn nur ein Kanal,
		signal = np.reshape(signal, (-1, 1))	#	2D-Spaltenvektor erzwingen

	if isinstance(fit_function, str):
		if fit_function == 'cos':
			funDE = _fitness_cos
		elif fit_function == 'exc':
			funDE = _fitness_exzenter
		else:
			raise ValueError("Valid values for `fit_function` are “cos”, “exc” or a valid fitness function.")
	else:
		funDE = fit_function

	fs = 1/np.float64(np.mean(np.diff(time_values)))  # Samplingfrequenz
	if np.bool_(np.mod(signal.shape[0], 2)):	# ungerade Signallänge funktioniert nicht für FFT
		if f_lowpass is not None:
			logging.warning(f"Signal length is uneven. The last sample will be cut off.")	# Warnung nur ausgeben, wenn es das ausgegebene Signal beeinflusst
		anzSamples = signal.shape[0] - 1
	else:
		anzSamples = signal.shape[0]

	spektren, freq = rfft_tukey_corrected(signal[:anzSamples, :], fs, axis=0, tukey=tukey_alpha) # Spektren der Signale berechnen

	if f_lowpass is not None and (f_lowpass >= (fs / 2.0)):
		logging.warning(f"Invalid value {f_lowpass:.3g} for f_lowpass filter. Filtering disabled in the current function call.")
		f_lowpass = None

	if f_lowpass is None:
		xFit = time_values  # Werte unverändert übernehmen
		yFit = signal
	else:
		idx_lowpass = np.argmax(freq > f_lowpass)  # Index der ersten zu unterdrückenden Frequenz
		spektren[idx_lowpass:, :] = 0  # Amplituden oberhalb der Grenzfrequenz unterdrücken

		yArray_lowpass = irfft_corrected(spektren)  # inverse FFT der maskierten Kanal-Spektren
		tukeyMaske = np.zeros(anzSamples, dtype=np.bool_)  # Maske für Zeitlinien, die vom Tukey-Fenster beeinflusst werden, initialisieren
		idx_tukey = np.int_(np.ceil(anzSamples * tukey_alpha / 2))  # jeweilige Anzahl der beeinflussten Zeitlinien am Beginn & Ende des Signals
		tukeyMaske[:idx_tukey] = True  # beeinflusste Zeitlinien am Beginn & Ende maskieren
		tukeyMaske[-idx_tukey:] = True
		xFit = np.delete(time_values[:anzSamples], tukeyMaske)  # beeinflusste Zeitlinien aus Zeitlinienvektor entfernen
		yFit = np.delete(yArray_lowpass, tukeyMaske, axis=0)  # und aus Signal entfernen

	x0 = None
	if bounds is None:  # wenn keine Grenzen übergeben wurden:
		bounds, x0 = _bounds_runoutFit(signal, freq, spektren, angle_tol)  # plausible Grenzwerte automatisch ermitteln

	res = sop.differential_evolution(funDE,  					# Modellfunktion fitten:
									 bounds=bounds,				#	mit gewählten Grenzen
									 args=(xFit, yFit, radius),	#		und während des Fits konstant bleibenden Werten;
									 x0=x0,  					#	Startwerte, falls übergeben
									 init='sobol',				#	Verfahren zur Erstellung der Initial-Population; besonderes Verhalten für 'sobol' s. Doku;
									 polish=True)				#	bestes Individuum nach Abschluss des evolutionären Algorithmus
																#		als Ausgangspunkt für Gradientenverfahren verwenden

	slack = tuple(bounds.residual(res.x))
	if np.any(slack[0] == 0) or np.any(slack[1] == 0):
		logging.warning("One or more position parameters match the solver’s bound for that parameter.")

	if position_zero:
		res.x[2::2] = _set_position_zero(res.x[2::2])	# Position des ersten Sensors auf Null setzen, weitere Positionen relativ dazu

	return res


def _bounds_runoutFit(signal: npt.NDArray[np.float64],
					  freq: npt.NDArray[np.float64],
					  spektra: npt.NDArray[np.complex128],
					  angle_tol: float,
) -> tuple[sop.Bounds, npt.NDArray[np.float64]]:
	"""
	Estimates plausible boundary constraints and starting values for the runout fit.

	Analyzes the input signal spectra and time-series statistics to define search
	ranges for rotational frequency, eccentricity, sensor angles, and offsets.

	Parameters
	----------
	signal : npt.NDArray[np.float64]
		2D array of raw sensor signals (samples x channels).
	freq : npt.NDArray[np.float64]
		Frequency vector corresponding to the spectra.
	spektra : npt.NDArray[np.complex128]
		Complex RFFT spectra of the sensor signals.
	angle_tol : float, default=5
		Half-width of the angular search window around the estimated sensor
		positions, in degrees.

	Returns
	-------
	bounds : scipy.optimize.Bounds
		The box constraints for the differential evolution algorithm.
	x0 : npt.NDArray[np.float64]
		An initial guess vector for the optimization parameters.
	"""
	# params = [fRot, betragExzentrizitaet, winkelSensor1, offsetSensor1, ..., winkelSensorN, offsetSensorN]
	df = np.mean(np.diff(freq))	# Frequenzlinienabstand ermitteln

	fRot, winkel = _initial_guess(freq, spektra) # geschätze Rotationsfrequenz und Sensor-Winkelpositionen (bezogen auf Lage der Exzentrizität beim Start der Messung)

	maxExz = 1.1 * np.max(np.abs(np.ptp(signal, axis=0)) / 2)	# die Exzentrizität kann nicht größer als die maximal gemessene Magnitude sein, +10% Sicherheit

	perzentile = np.percentile(signal, [25, 75, 50], axis=0)	# der statische Offset jedes Kanals liegt in der IQR; außerdem Median für Startwert des Lösers

	lBounds = np.asarray([fRot-df, 0], dtype=np.float64)		# untere Grenze: ermittelte Rotationsfrequenz minus Frequenzlinienabstand; keine Exzentrizität
	uBounds = np.asarray([fRot+df, maxExz], dtype=np.float64)	# obere Grenze: fRot plus Frequenzlinienabstand; maximal mögliche Exzentrizität
	x0 = np.asarray([fRot, maxExz], dtype=np.float64)			# Startwert des Lösers

	if signal.shape[1] == 1:	# abfangen, dass bei nur einem Messkanal kein Winkel bestimmt werden kann
		winkel = [np.pi]
		angle_tol = np.pi
	else:
		angle_tol = np.deg2rad(angle_tol)

	for a in range(signal.shape[1]):								# für jeden Kanal dessen Grenzen hinzufügen:
		lBounds = np.append(lBounds, np.asarray([winkel[a] - angle_tol, perzentile[0, a]], dtype=np.float64))	#	untere Grenze
		uBounds = np.append(uBounds, np.asarray([winkel[a] + angle_tol, perzentile[1, a]], dtype=np.float64))	#	obere Grenze
		x0 = np.append(x0, [winkel[a], perzentile[2, a]])						# Startwert des Lösers

	return sop.Bounds(lb=lBounds, ub=uBounds), x0


@nb.njit(parallel=True, cache=True)
def _fitness_exzenter(params: npt.NDArray[np.float64],
					  time_values: npt.NDArray[np.float64],
					  signal: npt.NDArray[np.float64],
					  radius: float,
) -> np.float64:
	"""
	Calculates the sum of squared errors using the exact eccentric geometry model.

	This model accounts for the non-linear relationship between eccentricity
	and radial displacement observed by a stationary sensor.

	Parameters
	----------
	params : npt.NDArray[np.float64]
		Optimization parameters: `[fRot, exz, angle_s1, offset_s1, ...]`.
	time_values : npt.NDArray[np.float64]
		1D array of time stamps.
	signal : npt.NDArray[np.float64]
		2D array of measured displacements (samples x channels).
	radius : float
		Nominal radius of the rotating object.

	Returns
	-------
	numpy.float64
		The total sum of squared residuals across all sensors and time steps.
	"""
	# params = [fRot, betragExzentrizitaet, winkelSensor1, offsetSensor1, ..., winkelSensorN, offsetSensorN]

	kreisFreq = 2 * np.pi * params[0]
	exz = params[1]
	sqRadius = np.square(radius)

	error_ssq = np.float64(0)
	for p in nb.prange(signal.shape[1]):
		idxP = np.int_(2 + 2 * p)

		for t in range(time_values.shape[0]):
			phi = kreisFreq * time_values[t]  # Winkelwert aus Schätzparametern erstellen
			phiShifted = phi - params[idxP]	# in Drehrichtung positive Winkelposition führt zu negativer Phasenverschiebung

			k = exz * np.sin(phiShifted)
			l1 = exz * np.cos(phiShifted)
			l2 = np.sqrt(sqRadius - np.square(k))
			s = l1 + l2
			deltaS = s - radius + params[idxP + 1]

			error_ssq = error_ssq + np.square(signal[t, p] - deltaS)

	return error_ssq


@nb.njit(parallel=True, cache=True)
def _fitness_cos(params: npt.NDArray[np.float64],
				 time_values: npt.NDArray[np.float64],
				 signal: npt.NDArray[np.float64],
				 *args,
) -> np.float64:
	"""
	Calculates the sum of squared errors using a first-order cosine approximation.

	This model assumes the sensor displacement is a pure sinusoid, which is
	a valid approximation for very small eccentricities relative to the radius.

	Parameters
	----------
	params : npt.NDArray[np.float64]
		Optimization parameters: `[fRot, exz, angle_s1, offset_s1, ...]`.
	time_values : npt.NDArray[np.float64]
		1D array of time stamps.
	signal : npt.NDArray[np.float64]
		2D array of measured displacements (samples x channels).
	*args : Any
		Unused arguments (provided for compatibility with the `fit_eccentricity` signature).

	Returns
	-------
	numpy.float64
		The total sum of squared residuals across all sensors and time steps.
	"""
	# params = [fRot, betragExzentrizitaet, winkelSensor1, offsetSensor1, ..., winkelSensorN, offsetSensorN]

	kreisFreq = 2 * np.pi * params[0]
	exz = params[1]

	error_ssq = np.float64(0)
	for p in nb.prange(signal.shape[1]):
		idxP = np.int_(2 + 2 * p)

		for t in range(time_values.shape[0]):
			phi = kreisFreq * time_values[t]  # Winkelwert aus Schätzparametern erstellen
			phiShifted = phi - params[idxP] # in Drehrichtung positive Winkelposition führt zu negativer Phasenverschiebung

			deltaS = exz * np.cos(phiShifted) + params[idxP + 1]

			error_ssq = error_ssq + np.square(signal[t, p] - deltaS)

	return error_ssq