"""Estimating the amplitude, phase, and frequency of a tone in noise.

Three parameters, three very different problems.  Given the frequency,
amplitude and phase fall out of a single complex inner product: the
signal model is linear in them, so the estimator is a projection.  The
frequency itself is not linear in anything, and the accuracy with which
it can be measured improves with the *cube* of the record length rather
than linearly.

The functions here are the estimators the topic page compares against
the Cramer-Rao bound, plus the bound itself.

Frequency convention
--------------------
Every function takes and returns frequencies in the same units as the
sample rate *fs* passed alongside them.  Leave *fs* at its default of
1.0 to work in cycles per sample (normalised frequency); pass a real
sample rate to work in Hz.
"""

import numpy as np


def coherent_dft(x: np.ndarray, f0: float, fs: float = 1.0,
                 window: np.ndarray | None = None) -> tuple[float, float]:
    """Estimate amplitude and phase of a tone of *known* frequency.

    This is the single-bin DFT, also called coherent or quadrature
    detection: correlate the record against a cosine and a sine at the
    known frequency and read the amplitude and phase off the resulting
    complex number.  It is the maximum likelihood estimator of amplitude
    and phase for a known frequency in white Gaussian noise, and it is
    the same operation a lock-in amplifier performs in hardware.

    Parameters
    ----------
    x : ndarray of shape (n,)
        Real input samples.
    f0 : float
        Known tone frequency, in the same units as *fs*.
    fs : float
        Sample rate.  Default 1.0 (frequency in cycles per sample).
    window : ndarray of shape (n,), optional
        Analysis window.  Without one the record is treated as
        rectangular, which is optimal for white noise when the tone
        falls exactly on a bin; a window trades a little noise
        performance for far less leakage when it does not.

    Returns
    -------
    amplitude : float
        Estimated amplitude A of ``A cos(2 pi f0 t + phase)``.
    phase : float
        Estimated phase in radians, wrapped to (-pi, pi].

    Notes
    -----
    The factor of two in the amplitude is the usual real-signal
    bookkeeping: a real cosine of amplitude A splits its energy between
    the positive and negative frequency components, so a single-sided
    correlation recovers A/2.
    """
    x = np.asarray(x, dtype=float).ravel()
    n = len(x)
    if n == 0:
        raise ValueError("x must not be empty")

    t = np.arange(n) / fs
    ref = np.exp(-2j * np.pi * f0 * t)

    if window is None:
        z = np.sum(x * ref) / n
    else:
        w = np.asarray(window, dtype=float).ravel()
        if len(w) != n:
            raise ValueError(f"window length {len(w)} != signal length {n}")
        # Normalise for unity passband gain, exactly as the coherent
        # gain correction in the MATLAB original.
        z = np.sum(x * w * ref) / np.sum(w)

    return 2.0 * np.abs(z), float(np.angle(z))


def crlb_sinusoid(A: float, sigma: float, N: int,
                  fs: float = 1.0) -> dict[str, float]:
    """Cramer-Rao lower bounds for a real tone in white Gaussian noise.

    For ``x[n] = A cos(2 pi f0 n / fs + phi) + w[n]`` with
    ``w ~ N(0, sigma^2)`` and all three of A, phi, f0 unknown, the
    variance of any unbiased estimator is bounded below by

        var(A_hat)   >= 2 sigma^2 / N
        var(phi_hat) >= 4 / (eta * N)                    (large N)
        var(f0_hat)  >= 12 fs^2 / ((2 pi)^2 eta N (N^2 - 1))

    where ``eta = A^2 / (2 sigma^2)`` is the signal-to-noise ratio.

    Note that these are the bounds for an *unknown* frequency.  If the
    frequency is known, phase is four times better determined; see
    :func:`crlb_known_frequency`.

    Parameters
    ----------
    A : float
        Tone amplitude (strictly positive).
    sigma : float
        Noise standard deviation (strictly positive).
    N : int
        Record length in samples (at least 2).
    fs : float
        Sample rate.  The frequency bound scales as fs^2 because it is
        expressed in the units of fs.

    Returns
    -------
    dict
        Keys ``'amplitude'``, ``'phase'``, ``'frequency'``, each the
        variance bound for that parameter, plus ``'snr'`` for eta.

    Notes
    -----
    Watch the parameterisation.  The frequency bound is conventionally
    quoted with a constant of 12 *in terms of the SNR eta*.  Written
    instead in terms of sigma^2 / A^2 the constant becomes 24, because
    eta carries a factor of 1/2 for a real sinusoid.  Substituting
    sigma^2 / A^2 for 1 / eta silently halves the bound; the topic page
    and the tests pin both forms against the exact Fisher information.

    The frequency and phase bounds use the standard large-N
    approximation, which drops terms oscillating at twice the tone
    frequency.  It is excellent except very close to DC or Nyquist,
    where those terms stop averaging away; see
    :func:`crlb_sinusoid_exact` for the version that makes no such
    assumption.
    """
    if A <= 0:
        raise ValueError("A must be strictly positive")
    if sigma <= 0:
        raise ValueError("sigma must be strictly positive")
    if N < 2:
        raise ValueError("N must be at least 2")

    eta = A**2 / (2.0 * sigma**2)
    return {
        'amplitude': 2.0 * sigma**2 / N,
        'phase': 4.0 / (eta * N),
        'frequency': 12.0 * fs**2 / ((2 * np.pi)**2 * eta * N * (N**2 - 1)),
        'snr': eta,
    }


def crlb_known_frequency(A: float, sigma: float, N: int) -> dict[str, float]:
    """Cramer-Rao bounds for amplitude and phase when f0 is known.

    Dropping frequency from the unknowns decouples the problem: the
    signal becomes linear in its two remaining parameters, so

        var(A_hat)   >= 2 sigma^2 / N
        var(phi_hat) >= 1 / (eta * N)

    with ``eta = A^2 / (2 sigma^2)``.  The amplitude bound is unchanged
    (amplitude was already orthogonal to the other two), but the phase
    bound improves by a factor of four: when the frequency is unknown, a
    frequency error masquerades as a drifting phase, and the estimator
    pays for the confusion.  This is the bound :func:`coherent_dft`
    should be measured against.
    """
    if A <= 0 or sigma <= 0:
        raise ValueError("A and sigma must be strictly positive")
    if N < 2:
        raise ValueError("N must be at least 2")

    eta = A**2 / (2.0 * sigma**2)
    return {
        'amplitude': 2.0 * sigma**2 / N,
        'phase': 1.0 / (eta * N),
        'snr': eta,
    }


def crlb_sinusoid_exact(A: float, phi: float, f0: float, sigma: float,
                        N: int, fs: float = 1.0) -> dict[str, float]:
    """Cramer-Rao bounds from the exact Fisher information matrix.

    Makes none of the large-N approximations of :func:`crlb_sinusoid`,
    so it stays valid near DC and Nyquist, at the cost of depending on
    the true phase and frequency (which in practice you do not know).
    Its purpose is to check the closed forms, and to show where they
    break.

    Returns the same keys as :func:`crlb_sinusoid`.
    """
    if A <= 0 or sigma <= 0:
        raise ValueError("A and sigma must be strictly positive")
    if N < 3:
        raise ValueError("need at least 3 samples to bound 3 parameters")

    n = np.arange(N)
    psi = 2 * np.pi * f0 * n / fs + phi

    # Partial derivatives of the noiseless signal with respect to
    # (A, phi, f0); the Fisher information of a Gaussian with known
    # variance is their Gram matrix over sigma^2.
    d_A = np.cos(psi)
    d_phi = -A * np.sin(psi)
    d_f0 = -A * 2 * np.pi * (n / fs) * np.sin(psi)

    G = np.vstack([d_A, d_phi, d_f0])
    fisher = (G @ G.T) / sigma**2
    cov = np.linalg.inv(fisher)
    return {
        'amplitude': float(cov[0, 0]),
        'phase': float(cov[1, 1]),
        'frequency': float(cov[2, 2]),
        'snr': A**2 / (2.0 * sigma**2),
    }


def _parabolic(a: float, b: float, c: float) -> float:
    """Vertex offset of the parabola through (-1, a), (0, b), (1, c)."""
    denom = a - 2 * b + c
    if denom == 0:
        return 0.0
    return 0.5 * (a - c) / denom


def estimate_frequency_fft(x: np.ndarray, fs: float = 1.0,
                           interpolation: str = 'quadratic') -> float:
    """Estimate tone frequency from the peak of the DFT magnitude.

    The raw peak is quantised to the bin spacing fs/N, which is a poor
    estimator: its error is bounded below by the grid, not by the noise.
    Interpolating between the peak bin and its two neighbours recovers
    most of the difference for very little arithmetic.

    Parameters
    ----------
    x : ndarray of shape (n,)
        Real input samples.
    fs : float
        Sample rate.
    interpolation : {'quadratic', 'jacobsen', 'none'}
        ``'quadratic'`` fits a parabola to the log magnitudes.
        ``'jacobsen'`` uses the complex three-bin estimator of Jacobsen
        and Kootsookos.  ``'none'`` returns the bare bin centre.

    Returns
    -------
    float
        Estimated frequency in the units of *fs*.

    Raises
    ------
    ValueError
        For an unknown *interpolation* mode or a record shorter than
        three samples.

    Notes
    -----
    **Match the interpolator to the window.**  Each of these fits an
    assumed main-lobe shape, and using the wrong one leaves a systematic
    bias that no amount of averaging removes, because the error depends
    on where the tone sits within its bin rather than on the noise.
    Jacobsen's estimator is derived for the rectangular window's
    Dirichlet lobe; the log-parabola fits the near-Gaussian lobe of a
    tapered window.  Measured over a swept fractional bin offset at
    23 dB SNR, N=256, the mismatched pairing is worse by an order of
    magnitude in both directions:

    ==========  ============  ============
    window      quadratic     jacobsen
    ==========  ============  ============
    rectangular   4.5e-04       1.5e-05
    Hann          4.7e-05       5.7e-04
    Blackman      2.7e-05       6.8e-04
    ==========  ============  ============

    (RMS error in cycles/sample; the topic page reproduces this table.)
    """
    x = np.asarray(x, dtype=float).ravel()
    n = len(x)
    if n < 3:
        raise ValueError("need at least 3 samples")
    if interpolation not in ('quadratic', 'jacobsen', 'none'):
        raise ValueError(f"unknown interpolation mode {interpolation!r}")

    X = np.fft.rfft(x)
    mag = np.abs(X)

    # Ignore DC and Nyquist: neither can host a resolvable tone, and the
    # interpolators need a bin on each side of the peak.
    if len(mag) < 3:
        raise ValueError("record too short to interpolate a peak")
    k = int(np.argmax(mag[1:-1])) + 1

    if interpolation == 'none':
        delta = 0.0
    elif interpolation == 'quadratic':
        # Log magnitudes: a Gaussian-ish main lobe is closer to a
        # parabola in log space than in linear space.
        eps = np.finfo(float).tiny
        delta = _parabolic(*(np.log(mag[k - 1:k + 2] + eps)))
    else:
        num = X[k - 1] - X[k + 1]
        den = 2 * X[k] - X[k - 1] - X[k + 1]
        delta = float(np.real(num / den)) if den != 0 else 0.0

    return (k + delta) * fs / n


def estimate_frequency_phase_diff(x: np.ndarray, fs: float = 1.0) -> float:
    """Estimate tone frequency from the phase advance between halves.

    Split the record in two, measure each half's phase at the coarse
    FFT peak frequency, and read the frequency offset off the phase
    difference: a tone that is off by df accumulates 2*pi*df*(N/2)/fs
    radians of extra phase across half a record.

    This is the same idea a phase-locked loop uses, and the same one
    behind phase-vocoder frequency refinement.  It is cheap and, at
    usable SNR, close to the Cramer-Rao bound; its weakness is that the
    phase difference must be unambiguous, so the coarse estimate has to
    be within half a bin already.

    Parameters
    ----------
    x : ndarray of shape (n,)
        Real input samples (at least 4).
    fs : float
        Sample rate.

    Returns
    -------
    float
        Estimated frequency in the units of *fs*.
    """
    x = np.asarray(x, dtype=float).ravel()
    n = len(x)
    if n < 4:
        raise ValueError("need at least 4 samples")

    half = n // 2
    f_coarse = estimate_frequency_fft(x[:half], fs=fs, interpolation='none')

    _, phase_1 = coherent_dft(x[:half], f_coarse, fs=fs)
    _, phase_2 = coherent_dft(x[half:2 * half], f_coarse, fs=fs)

    # Phase accumulated over the hop, wrapped into (-pi, pi].
    d_phase = np.angle(np.exp(1j * (phase_2 - phase_1)))
    return f_coarse + d_phase * fs / (2 * np.pi * half)