#  Copyright (c) 2026. Jose M. Requena-Plens
"""
Frequency-response and coherence estimators (Bendat & Piersol).

Two-channel (input/output) system identification from measured signals, using
the Welch-averaged cross- and auto-spectral densities. Following Bendat &
Piersol, *Random Data: Analysis and Measurement Procedures* (4th ed., 2010):

* the **H1** estimator ``H1 = Gxy / Gxx`` (unbiased when the noise is on the
  output),
* the **H2** estimator ``H2 = Gyy / Gyx`` (unbiased when the noise is on the
  input),
* the **ordinary coherence** ``γ² = |Gxy|² / (Gxx · Gyy)`` ∈ [0, 1], the
  fraction of the output power linearly explained by the input.

For a noiseless linear time-invariant path both estimators recover the true
transfer function and the coherence is unity; additive output noise biases H2
but not H1 and pulls the coherence down to ``SNR / (1 + SNR)``, which is the
analytic oracle used to verify the implementation.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Literal

import numpy as np

if TYPE_CHECKING:
    from matplotlib.axes import Axes
    from numpy.typing import NDArray

#: Default overlap fraction between Welch segments.
_DEFAULT_OVERLAP = 0.5
#: Minimum samples for a spectral estimate.
_MIN_SAMPLES = 32


def _positive(value: float, name: str) -> float:
    scalar = float(value)
    if not np.isfinite(scalar) or scalar <= 0.0:
        raise ValueError(f"'{name}' must be a positive, finite number.")
    return scalar


def _validate_pair(
    x: "NDArray[np.float64] | list[float]", y: "NDArray[np.float64] | list[float]"
) -> tuple["NDArray[np.float64]", "NDArray[np.float64]"]:
    xa = np.asarray(x, dtype=np.float64)
    ya = np.asarray(y, dtype=np.float64)
    if xa.ndim != 1 or ya.ndim != 1:
        raise ValueError("'x' and 'y' must be one-dimensional.")
    if xa.size != ya.size:
        raise ValueError("'x' and 'y' must have the same length.")
    if xa.size < _MIN_SAMPLES:
        raise ValueError(
            f"Signals too short for a spectral estimate: {xa.size} samples."
        )
    if not (np.all(np.isfinite(xa)) and np.all(np.isfinite(ya))):
        raise ValueError("'x' and 'y' must be finite.")
    return xa, ya


def _default_nperseg(n: int, fs: float) -> int:
    nperseg = int(min(n, 2 ** int(np.ceil(np.log2(fs / 4.0)))))
    return max(nperseg, _MIN_SAMPLES)


def _spectra(
    x: "NDArray[np.float64]",
    y: "NDArray[np.float64]",
    fs: float,
    nperseg: int,
    overlap: float,
) -> tuple[
    "NDArray[np.float64]",
    "NDArray[np.complex128]",
    "NDArray[np.float64]",
    "NDArray[np.float64]",
]:
    """Return ``(freqs, Gxy, Gxx, Gyy)`` from Welch-averaged Hann segments."""
    from scipy import signal as sp_signal

    noverlap = min(int(round(overlap * nperseg)), nperseg - 1)
    kw = dict(fs=fs, window="hann", nperseg=nperseg, noverlap=noverlap, detrend=False)
    freqs, gxy = sp_signal.csd(x, y, **kw)
    _, gxx = sp_signal.welch(x, **kw)
    _, gyy = sp_signal.welch(y, **kw)
    return (
        np.asarray(freqs, dtype=np.float64),
        np.asarray(gxy, dtype=np.complex128),
        np.asarray(gxx, dtype=np.float64),
        np.asarray(gyy, dtype=np.float64),
    )


@dataclass(frozen=True)
class FrequencyResponseResult:
    """Estimated frequency response of an input/output path (Bendat & Piersol).

    :ivar frequencies: Frequency axis, in Hz.
    :ivar response: Complex frequency-response estimate ``H(f)``.
    :ivar magnitude_db: Magnitude ``20·lg|H|``, in dB.
    :ivar phase: Phase of ``H``, in radians (unwrapped).
    :ivar coherence: Ordinary coherence ``γ²(f)`` ∈ [0, 1].
    :ivar estimator: Estimator used (``'H1'`` or ``'H2'``).
    """

    frequencies: "NDArray[np.float64]"
    response: "NDArray[np.complex128]"
    magnitude_db: "NDArray[np.float64]"
    phase: "NDArray[np.float64]"
    coherence: "NDArray[np.float64]"
    estimator: str

    def plot(
        self, ax: "Axes | None" = None, **kwargs: Any
    ) -> "Axes | NDArray[Any]":
        """Plot the Bode magnitude/phase and the coherence."""
        from ._plotting import plot_frequency_response

        return plot_frequency_response(self, ax=ax, **kwargs)


def transfer_function(
    x: "NDArray[np.float64] | list[float]",
    y: "NDArray[np.float64] | list[float]",
    fs: float,
    *,
    estimator: Literal["H1", "H2"] = "H1",
    nperseg: int | None = None,
    overlap: float = _DEFAULT_OVERLAP,
) -> FrequencyResponseResult:
    """Estimate the frequency response from input ``x`` to output ``y``.

    ``H1 = Gxy / Gxx`` (the default; unbiased for output noise) or
    ``H2 = Gyy / Gyx`` (unbiased for input noise), from Welch-averaged Hann
    segments. The ordinary coherence ``γ² = |Gxy|² / (Gxx·Gyy)`` is returned
    alongside as a data-quality indicator.

    :param x: Input (reference) signal, 1-D.
    :param y: Output (response) signal, 1-D, same length as ``x``.
    :param fs: Sample rate, in Hz.
    :param estimator: ``'H1'`` (default) or ``'H2'``.
    :param nperseg: Welch segment length; ``None`` picks a length targeting
        about 4 Hz resolution (as in :mod:`~phonometry.intensity`).
    :param overlap: Segment overlap fraction (default 0.5).
    :return: A :class:`FrequencyResponseResult`.
    :raises ValueError: If the inputs or parameters are invalid.
    """
    xa, ya = _validate_pair(x, y)
    fs_v = _positive(fs, "fs")
    if estimator not in ("H1", "H2"):
        raise ValueError("'estimator' must be 'H1' or 'H2'.")
    if not 0.0 <= float(overlap) < 1.0:
        raise ValueError("'overlap' must be in [0, 1).")
    seg = _default_nperseg(xa.size, fs_v) if nperseg is None else int(nperseg)
    if seg < _MIN_SAMPLES or seg > xa.size:
        raise ValueError("'nperseg' must be between 32 and the signal length.")

    freqs, gxy, gxx, gyy = _spectra(xa, ya, fs_v, seg, float(overlap))
    if estimator == "H1":
        response = np.divide(
            gxy, gxx, out=np.zeros_like(gxy), where=gxx > 0.0
        )
    else:
        # H2 = Gyy / Gyx = Gyy / conj(Gxy).
        gyx = np.conj(gxy)
        response = np.divide(
            gyy.astype(np.complex128),
            gyx,
            out=np.zeros_like(gxy),
            # |gyx|^2 > 0 guards division by a zero denominator without a
            # float equality check or the sqrt in np.abs.
            where=gyx.real**2 + gyx.imag**2 > 0.0,
        )
    denom = gxx * gyy
    coh = np.divide(
        np.abs(gxy) ** 2, denom, out=np.zeros_like(gxx), where=denom > 0.0
    )
    coh = np.clip(coh, 0.0, 1.0)
    with np.errstate(divide="ignore"):
        magnitude_db = 20.0 * np.log10(np.abs(response))
    magnitude_db[~np.isfinite(magnitude_db)] = -np.inf
    phase = np.unwrap(np.angle(response))
    return FrequencyResponseResult(
        frequencies=freqs,
        response=response,
        magnitude_db=magnitude_db,
        phase=phase,
        coherence=coh,
        estimator=estimator,
    )


def coherence(
    x: "NDArray[np.float64] | list[float]",
    y: "NDArray[np.float64] | list[float]",
    fs: float,
    *,
    nperseg: int | None = None,
    overlap: float = _DEFAULT_OVERLAP,
) -> tuple["NDArray[np.float64]", "NDArray[np.float64]"]:
    """Ordinary coherence ``γ²(f)`` between ``x`` and ``y`` (Bendat & Piersol).

    ``γ² = |Gxy|² / (Gxx·Gyy)`` ∈ [0, 1]: unity for a noiseless linear path and
    ``SNR/(1+SNR)`` with additive output noise. Averaging over several segments
    is required for a meaningful estimate (a single segment gives ``γ² ≡ 1``).

    :param x: First signal, 1-D.
    :param y: Second signal, 1-D, same length as ``x``.
    :param fs: Sample rate, in Hz.
    :param nperseg: Welch segment length; ``None`` picks a default.
    :param overlap: Segment overlap fraction (default 0.5).
    :return: ``(frequencies, gamma_squared)``.
    :raises ValueError: If the inputs or parameters are invalid.
    """
    xa, ya = _validate_pair(x, y)
    fs_v = _positive(fs, "fs")
    if not 0.0 <= float(overlap) < 1.0:
        raise ValueError("'overlap' must be in [0, 1).")
    seg = _default_nperseg(xa.size, fs_v) if nperseg is None else int(nperseg)
    if seg < _MIN_SAMPLES or seg > xa.size:
        raise ValueError("'nperseg' must be between 32 and the signal length.")
    freqs, gxy, gxx, gyy = _spectra(xa, ya, fs_v, seg, float(overlap))
    denom = gxx * gyy
    coh = np.divide(
        np.abs(gxy) ** 2, denom, out=np.zeros_like(gxx), where=denom > 0.0
    )
    return freqs, np.clip(coh, 0.0, 1.0)
