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

"""
Calculates runout, eccentricity, and true shape profiles using a multi-probe method.

This module processes sensor signals to separate spindle eccentricity from the
true radius variations of a rotating workpiece. It uses Fourier transformations
and linear least-squares equations evaluated in the frequency domain, then transforms
the separated signals back into the angular domain.
"""

from typing import Literal

import numba as nb
import numpy as np
from numpy import typing as npt

from .resample import *
from .rfft_corrected import *

__all__ = ["runout_multiprobe"]

def runout_multiprobe(sensor_sig: npt.NDArray[np.float64],
					  sensor_pos: npt.NDArray[np.float64],
					  f_sampling: float,
					  f_rotation: float,
					  n_min: int = 2,
					  n_max: int = 100,
					  tukey_alpha: float = 0.02,
					  coordinate_system: Literal['co', 'cntr'] = 'co',
) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64]]:
	"""
	Calculate eccentricity, radius variations, and the angle vector from multiprobe sensor signals.

	This function resamples time-domain sensor signals into the angular domain, computes
	their Fourier spectra, and applies a multi-probe mathematical separation on a defined
	range of harmonic frequencies. The separated spectra are then transformed back into
	the angular domain via inverse real Fourier transforms.

	Parameters
	----------
	sensor_sig : npt.NDArray[np.float64]
		The time-domain signals recorded by the sensors.
	sensor_pos : npt.NDArray[np.float64]
		The angular positions of the sensors in radians.
	f_sampling : float
		The sampling frequency of the sensor measurements in Hz.
	f_rotation : float
		The rotational frequency of the shaft in Hz.
	n_min : int, default: 2
		The minimum harmonic order to include in the multi-probe calculation.
	n_max : int, default: 100
		The maximum harmonic order to include in the multi-probe calculation.
	tukey_alpha : float, default: 0.02
		The shape parameter (alpha) for the Tukey window applied during resampling.
	coordinate_system : {'co', 'cntr'}, default: 'co'
		The orientation of the coordinate system. With ‘co’, the rotation angle
		and the angle on the surface point in the same direction; in other words,
		a rotation of the shaft by +dPhi causes the sensors on the surface to point
		in the direction of -dPhi relative to their previous position.

	Returns
	-------
	eccentricity : npt.NDArray[np.float64]
		An array representing the x and y components of the shaft's radial error
		motion in the angular domain.
	radius_phi : npt.NDArray[np.float64]
		An array representing the true radius variations (workpiece shape) in the
		angular domain.
	phi : npt.NDArray[np.float64]
		The angular position vector corresponding to the output arrays.
	"""
	# Koordinatensystem ggf. anpassen
	if coordinate_system == 'co':
		sensor_sig = np.flipud(sensor_sig)

	# Zeitvektor erstellen und resamplen
	zeit = np.arange(sensor_sig.shape[0]) / f_sampling
	y_array_resampled, x_array_resampled, f_s_resampled, steps = resample_to_harmonics(zeit, sensor_sig, f_sampling, f_rotation, tukey_alpha)

	# Spektren berechnen; x-Einheit wird "pro Umdrehung"
	spektren, freq = rfft_tukey_corrected(y_array_resampled, f_s_resampled/f_rotation, tukey=0.0)

	# Fourierkoeffizienten der Harmonischen im gewünschten Bereich markieren
	maske = np.zeros(freq.size, dtype=np.bool_)
	maske[n_min*steps:(n_max+1)*steps:steps] = True

	# Multi-probe method anwenden; nur für markierte Harmonische rechnen
	ergebnis_harm = _multiprobe_solve(np.arange(n_min, n_max+1, dtype=np.int_),
									  spektren[maske, :],
									  sensor_pos)

	# Ergebnis zu vollem Spektrum erweitern; alle anderen Fourierkoeffizienten sind 0
	ergebnis = np.zeros((3, freq.size), dtype=np.complex128)
	ergebnis[:, maske] = ergebnis_harm

	# Ergebnis in Winkelbereich transformieren
	eccentricity = irfft_corrected(ergebnis[:2, :], axis=1).T
	radius_phi = irfft_corrected(ergebnis[2, :])

	# Winkellinien erstellen
	phi = x_array_resampled * f_rotation * 2 * np.pi

	# ggf. Anpassung des Koordinatensystems rückgängig machen
	if coordinate_system == 'co':
		eccentricity = np.flipud(eccentricity)

	return eccentricity, radius_phi, phi


@nb.njit(cache=True, parallel=True)
def _multiprobe_solve(harmonics_order: npt.NDArray[np.int_],
					  harmonics_fourier_coefficient: npt.NDArray[np.complex128],
					  positions: npt.NDArray[np.float64],
) -> npt.NDArray[np.complex128]:
	"""
	Solve the linear least-squares system to separate eccentricity and radius harmonics.

	This internal function uses Numba for JIT compilation and parallelization. It iterates
	through the specified harmonic orders and constructs a position matrix for each. It then
	solves the linear system to distribute the measured deviations into eccentricity and
	true radius components.

	Parameters
	----------
	harmonics_order : npt.NDArray[np.int_]
		A 1D array containing the specific harmonic orders to be solved.
	harmonics_fourier_coefficient : npt.NDArray[np.complex128]
		A 2D array representing the complex Fourier coefficients for the requested harmonics.
	positions : npt.NDArray[np.float64]
		The angular positions of the measurement sensors in radians.

	Returns
	-------
	npt.NDArray[np.complex128]
		A 2D array containing the separated Fourier components. Row 0 represents axis movement
		in x, Row 1 represents axis movement in y, and Row 2 contains the radius variations.
	"""
	# initialisieren: Ergebnismatrix und frequenzunabhängigen komplexen Positionsvektor
	ergebnis = np.empty((2, harmonics_fourier_coefficient.shape[0]), dtype=np.complex128)
	posVektor = np.exp(1j * positions)

	for f in nb.prange(harmonics_fourier_coefficient.shape[0]):
		posMatrix = np.column_stack((posVektor,	# eigene `posMatrix` für jeden Thread, sonst race conditions
									 np.exp(1j * np.float64(harmonics_order[f]) * positions)))

		# bei Abweichungen mit Rotationsfrequenz: Rundlaufabweichung vollständig dem Radius zuschreiben
		if harmonics_order[f] == 1:	# Datentyp np.int_ ermöglicht sichere Übereinstimmung
			posMatrix[:, 0] =  np.zeros(positions.size, dtype=np.complex128)

		# Fourierkoeffizienten für Radius und Exzentrizität trennen
		ergebnis[:, f], _, _, _ = np.linalg.lstsq(posMatrix, harmonics_fourier_coefficient[f, :].T)

	return np.vstack((ergebnis[0, :], 1j * ergebnis[0, :], ergebnis[1, :]))