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

"""
Provide initial parameter estimation for eccentric motion fitting.

This internal module contains a helper function to estimate the rotational
frequency and the relative angular positions of sensors from their frequency spectra.
"""

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

from ._set_position_zero import _set_position_zero

__all__ = ["_initial_guess"]


def _initial_guess(frequency_lines: npt.NDArray[np.float64],
				   spectra: npt.NDArray[np.complex128],
				   position_zero: bool = False,
) -> tuple[np.float64, npt.NDArray[np.float64] | None]:
	"""
	Estimate rotational frequency and sensor angular positions from spectra.

	Analyzes the sum of the magnitudes of the provided spectra to find the
	predominant frequency peak, which is used to estimate the rotational
	frequency. It then calculates the angular positions of the sensors based
	on the phase angles at this estimated rotational frequency.

	Parameters
	----------
	frequency_lines : numpy.ndarray of numpy.float64
		A 1D array of frequency lines corresponding to the spectra.
	spectra : numpy.ndarray of numpy.complex128
		A 1D or 2D array of complex RFFT spectra. If 2D, the columns
		should represent different sensor channels.
	position_zero : bool, default=False
		If True, shifts the calculated angular positions so that the first
		sensor's position is exactly 0 radians, making the others relative.

	Returns
	-------
	f_rotation : numpy.float64
		The estimated rotational frequency.
	angular_position : numpy.ndarray of numpy.float64 or None
		A 1D array containing the estimated angular positions of the sensors
		in radians. Returns None if the input `spectra` is only 1-dimensional.
	"""
	frequency_lines = np.squeeze(np.asanyarray(frequency_lines))  # ggf. zu ndarray konvertieren und leere Dimensionen entfernen;
	spectra = np.squeeze(np.asanyarray(spectra))  # `frequency_lines` sollte nun 1D-Array sein

	if spectra.ndim > 1:
		if spectra.shape[1] > spectra.shape[0]:  # spaltenweise Orientierung der Kanäle sicherstellen
			spectra = spectra.T
		summeSpektrum = np.sum(np.abs(spectra), axis=1)  # und Beträge aller Kanäle summieren
	else:
		summeSpektrum = np.abs(np.ravel(spectra))

	if summeSpektrum[0] > summeSpektrum[1]:
		idxMinima, _ = ssig.find_peaks(-summeSpektrum, distance=3)  # alle lokalen Minima der Spektrensumme ermitteln
		idxMinima = np.asanyarray(idxMinima, dtype=np.int_)  # Fix: Typ explizit für PyCharm-Linter
		if idxMinima.size > 0:	# Index des Maximums der Spektrensumme oberhalb der Frequenzlinie des ersten Minimums ermitteln
			idxRot = np.argmax(summeSpektrum[idxMinima[0]:]) + idxMinima[0]	# die dazugehörige Frequenz entspricht etwa der Drehfrequenz
		else:
			idxRot = np.argmax(summeSpektrum)
	else:  # bei sehr niedrigen Frequenzauflösungen gibt es ggf. kein Minimum vor dem Maximum
		idxRot = np.argmax(summeSpektrum)

	f_rotation = np.float64(frequency_lines[idxRot])

	angular_position = None
	if spectra.ndim == 2:  # wenn Winkel zwischen Kanälen ermittelt werden kann:
		phase = np.angle(spectra[idxRot, :])  # Phasengänge der Kanäle
		angular_position = -phase	# Winkel negieren, da positiver Sensor-Positionswinkel in Messung zu negativer Phasenverschiebung führt

		if position_zero:
			angular_position = _set_position_zero(angular_position)  # Position des ersten Sensors auf Null setzen, weitere Positionen relativ dazu

	return f_rotation, angular_position
