"""Detection theory: thresholds, ROC curves, and CFAR.

A detector answers a yes/no question about a signal: is the target
there, is someone speaking, did the tone appear.  Every detector in
this workshop, however it was motivated, is a test statistic compared
to a threshold, and this module collects the closed forms that say how
well such a test can do: the probability of a false alarm (deciding
"signal" on noise alone) and of detection (deciding "signal" when it
is there), as functions of the threshold and the signal-to-noise
ratio.

Parameterisation, stated up front per the arc's standing rule
----------------------------------------------------------------
Signal strength enters every formula here as the **energy-to-noise
ratio**

    enr = Es / sigma^2,        Es = sum(s[m]^2)

with sigma^2 the noise variance per sample.  This is NOT the same
number as a per-sample SNR (enr = N * SNR for an N-sample window) nor
as the sinusoid eta = A^2 / (2 sigma^2) of the estimating-a-sinusoid
topic (for a whole-cycle tone, enr = N * A^2 / (2 sigma^2) = N * eta).
The tests pin every closed form against Monte Carlo, because a
detection-statistic distribution quoted from memory is exactly as
treacherous as a Cramer-Rao bound quoted from memory.

The three detectors covered
---------------------------
- Matched filter (known waveform): statistic sum(y s), Gaussian under
  both hypotheses, Pd = Q(Q^-1(Pfa) - sqrt(enr)).
- Energy detector (unknown waveform): statistic sum(y^2), chi-square
  with N degrees of freedom under H0, noncentral chi-square under H1.
- CA-CFAR energy detector (unknown waveform AND unknown noise level):
  the cell energy compared to a scaled average of N_ref reference
  cells; the ratio is F-distributed, so Pfa depends only on the
  threshold factor, never on sigma^2.  That independence is the whole
  point of CFAR.

All distributions are exact, not asymptotic, for white Gaussian noise.
"""

import numpy as np
from scipy.stats import chi2, f as f_dist, ncf, ncx2, norm


def matched_filter_pd(pfa: float, enr: float) -> float:
    """Detection probability of the matched filter at a given Pfa.

    For a known waveform s in white Gaussian noise the Neyman-Pearson
    detector correlates the data with the waveform, T = sum(y[m] s[m]),
    and T is Gaussian under both hypotheses with the same variance
    sigma^2 Es.  Normalising, the detection probability at false-alarm
    probability *pfa* is

        Pd = Q( Q^-1(pfa) - sqrt(enr) ),    enr = Es / sigma^2

    with Q the standard Gaussian tail function.  The single number
    sqrt(enr) is the detector's deflection: how many H0 standard
    deviations the statistic moves when the signal appears.  Waveform
    shape does not enter, only energy, which is the matched filter's
    signature property.

    Parameters
    ----------
    pfa : float
        False-alarm probability, in (0, 1).
    enr : float
        Energy-to-noise ratio Es / sigma^2 (dimensionless, >= 0).
        NOT a per-sample SNR: see the module docstring.

    Returns
    -------
    float
        Detection probability, in (0, 1).
    """
    if not 0.0 < pfa < 1.0:
        raise ValueError("pfa must be in (0, 1)")
    if enr < 0:
        raise ValueError("enr must be non-negative")
    return float(norm.sf(norm.isf(pfa) - np.sqrt(enr)))


def energy_detector_threshold(pfa: float, n: int, sigma2: float = 1.0) -> float:
    """Threshold on the energy statistic sum(y^2) for a target Pfa.

    Under H0 (noise only, variance *sigma2*), sum over *n* samples of
    y^2 / sigma2 is chi-square with n degrees of freedom, so the
    threshold is

        gamma = sigma2 * chi2_isf(pfa, n).

    Note the dependence on *sigma2*: this threshold is only as good as
    the noise-variance estimate behind it.  Quantifying how badly a
    wrong sigma2 breaks Pfa, and removing the dependence, is the CFAR
    section of the topic page.

    Parameters
    ----------
    pfa : float
        False-alarm probability, in (0, 1).
    n : int
        Number of samples in the detection window (>= 1).
    sigma2 : float
        Noise variance per sample.

    Returns
    -------
    float
        Threshold gamma such that P[sum(y^2) > gamma | H0] = pfa.
    """
    if not 0.0 < pfa < 1.0:
        raise ValueError("pfa must be in (0, 1)")
    if n < 1:
        raise ValueError("n must be at least 1")
    if sigma2 <= 0:
        raise ValueError("sigma2 must be strictly positive")
    return float(sigma2 * chi2.isf(pfa, n))


def energy_detector_pd(pfa: float, enr: float, n: int) -> float:
    """Detection probability of the energy detector at a given Pfa.

    With a deterministic signal of energy Es present, the normalised
    energy statistic sum(y^2)/sigma^2 over *n* samples is noncentral
    chi-square with n degrees of freedom and noncentrality
    lambda = enr = Es/sigma^2, so

        Pd = Qchi2'( chi2_isf(pfa, n); n, enr ).

    Compare :func:`matched_filter_pd` at the same pfa and enr: the
    energy detector is always worse, and the gap grows with n.  The
    deflection works out to enr / sqrt(2 n) against the matched
    filter's sqrt(enr): not knowing the waveform costs a factor
    ~sqrt(n) in required energy at low SNR.  Both statements are
    exercised in the tests and on the topic page.

    Parameters
    ----------
    pfa : float
        False-alarm probability, in (0, 1).
    enr : float
        Energy-to-noise ratio Es / sigma^2 (>= 0).
    n : int
        Number of samples in the detection window (>= 1).

    Returns
    -------
    float
        Detection probability, in (0, 1).
    """
    if not 0.0 < pfa < 1.0:
        raise ValueError("pfa must be in (0, 1)")
    if enr < 0:
        raise ValueError("enr must be non-negative")
    if n < 1:
        raise ValueError("n must be at least 1")
    if enr == 0:
        return float(pfa)
    return float(ncx2.sf(chi2.isf(pfa, n), n, enr))


def incoherent_pd(pfa: float, lam: float) -> float:
    """Detection probability of a magnitude (incoherent) detector.

    The lock-in / single-bin-DFT setting: two demodulated quadratures,
    each Gaussian with per-quadrature variance sigma_z^2, and the
    detector thresholds the magnitude |z|.  Under H0 the magnitude is
    Rayleigh (|z|^2 / sigma_z^2 is chi-square with 2 degrees of
    freedom), giving the closed threshold relation

        Pfa = exp( -gamma^2 / (2 sigma_z^2) )

    and under H1 (a tone of amplitude A at the reference frequency)
    |z|^2 / sigma_z^2 is noncentral chi-square with 2 degrees of
    freedom and noncentrality lam = A^2 / sigma_z^2, so

        Pd = Qchi2'( chi2_isf(pfa, 2); 2, lam )

    (equivalently the Marcum Q function Q1(sqrt(lam), sqrt(-2 ln pfa))).

    Parameters
    ----------
    pfa : float
        False-alarm probability, in (0, 1).
    lam : float
        Noncentrality A^2 / sigma_z^2, with A the tone amplitude seen
        by the quadrature pair and sigma_z the per-quadrature noise
        standard deviation (>= 0).

    Returns
    -------
    float
        Detection probability, in (0, 1).
    """
    if not 0.0 < pfa < 1.0:
        raise ValueError("pfa must be in (0, 1)")
    if lam < 0:
        raise ValueError("lam must be non-negative")
    if lam == 0:
        return float(pfa)
    return float(ncx2.sf(chi2.isf(pfa, 2), 2, lam))


def cfar_factor(pfa: float, n_block: int, n_ref: int) -> float:
    """Threshold factor alpha for a cell-averaging CFAR energy detector.

    The detector compares the energy of the cell under test (a block
    of *n_block* real Gaussian samples) against alpha times the MEAN
    energy of *n_ref* reference blocks:

        E_cut  >  alpha * (E_1 + ... + E_nref) / n_ref   =>  detect.

    Under H0 every block energy is sigma^2 times a chi-square, and the
    ratio statistic E_cut / mean(E_ref) is F-distributed with
    (n_block, n_ref * n_block) degrees of freedom; sigma^2 cancels.
    Hence

        alpha = F_isf(pfa; n_block, n_ref * n_block)

    and the false-alarm probability is exactly *pfa* at every noise
    level, which is the constant-false-alarm-rate property.  For
    n_block = 2 this reduces to the classic radar CA-CFAR result for
    exponentially distributed (square-law) cells,
    Pfa = (1 + alpha/n_ref)^(-n_ref), which the tests verify.

    If the block energies are computed with the block mean removed
    (as the embedded implementation does, to kill the ADC's DC
    pedestal), each block loses one degree of freedom: pass
    n_block - 1 for n_block.

    Parameters
    ----------
    pfa : float
        False-alarm probability, in (0, 1).
    n_block : int
        Chi-square degrees of freedom per block: the number of samples
        per block (>= 1), or samples minus one if the block mean is
        subtracted.
    n_ref : int
        Number of reference blocks averaged (>= 1).

    Returns
    -------
    float
        Threshold factor alpha (dimensionless).
    """
    if not 0.0 < pfa < 1.0:
        raise ValueError("pfa must be in (0, 1)")
    if n_block < 1 or n_ref < 1:
        raise ValueError("n_block and n_ref must be at least 1")
    return float(f_dist.isf(pfa, n_block, n_ref * n_block))


def cfar_pd(pfa: float, enr: float, n_block: int, n_ref: int) -> float:
    """Detection probability of the CA-CFAR energy detector.

    With a deterministic signal of energy Es in the cell under test
    (and none in the reference blocks), the ratio statistic becomes
    noncentral-F with (n_block, n_ref * n_block) degrees of freedom
    and noncentrality enr = Es / sigma^2:

        Pd = QF'( alpha; n_block, n_ref * n_block, enr )

    with alpha from :func:`cfar_factor` at the same pfa.  Comparing
    against :func:`energy_detector_pd` at the same pfa and enr gives
    the **CFAR loss**: the price, in required signal energy, of having
    to estimate the noise level from n_ref finite reference blocks
    instead of knowing it.  The loss shrinks toward zero as n_ref
    grows; the topic page plots it.

    Parameters
    ----------
    pfa : float
        False-alarm probability, in (0, 1).
    enr : float
        Energy-to-noise ratio Es / sigma^2 of the signal in the cell
        under test (>= 0).
    n_block : int
        Chi-square degrees of freedom per block (see
        :func:`cfar_factor` for the mean-removed case).
    n_ref : int
        Number of reference blocks averaged (>= 1).

    Returns
    -------
    float
        Detection probability, in (0, 1).
    """
    if not 0.0 < pfa < 1.0:
        raise ValueError("pfa must be in (0, 1)")
    if enr < 0:
        raise ValueError("enr must be non-negative")
    if n_block < 1 or n_ref < 1:
        raise ValueError("n_block and n_ref must be at least 1")
    alpha = f_dist.isf(pfa, n_block, n_ref * n_block)
    if enr == 0:
        return float(pfa)
    return float(ncf.sf(alpha, n_block, n_ref * n_block, enr))


def roc_empirical(scores_h0: np.ndarray,
                  scores_h1: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Empirical ROC curve from Monte Carlo detector scores.

    Sweeps the threshold over every distinct observed score and
    returns, at each setting, the fraction of H0 scores at or above it
    (the false-alarm rate) and the fraction of H1 scores at or above
    it (the detection rate).  This is the measurement that a closed
    Pfa/Pd formula must survive being compared against.

    Parameters
    ----------
    scores_h0 : ndarray
        Detector statistic under H0 (noise only), any shape, flattened.
    scores_h1 : ndarray
        Detector statistic under H1 (signal present), flattened.

    Returns
    -------
    pfa : ndarray
        False-alarm rates, decreasing from 1 toward 0.
    pd : ndarray
        Detection rates at the same thresholds.
    """
    s0 = np.sort(np.asarray(scores_h0, dtype=float).ravel())
    s1 = np.sort(np.asarray(scores_h1, dtype=float).ravel())
    if len(s0) == 0 or len(s1) == 0:
        raise ValueError("both score arrays must be non-empty")
    thresholds = np.unique(np.concatenate([s0, s1]))
    # P[score >= t] via searchsorted on the sorted arrays.
    pfa = 1.0 - np.searchsorted(s0, thresholds, side='left') / len(s0)
    pd = 1.0 - np.searchsorted(s1, thresholds, side='left') / len(s1)
    return pfa, pd


def block_energies(x: np.ndarray, n_block: int,
                   remove_mean: bool = False) -> np.ndarray:
    """Energy of consecutive non-overlapping blocks of a record.

    Splits *x* into floor(len(x)/n_block) whole blocks and returns
    sum(x^2) per block, optionally with each block's own mean removed
    first.  Mean removal is what an embedded energy detector does to
    ignore the ADC's DC pedestal; it costs one chi-square degree of
    freedom per block (pass n_block - 1 to the threshold functions).

    Parameters
    ----------
    x : ndarray
        Input record; trailing samples short of a block are dropped.
    n_block : int
        Samples per block (>= 2 if remove_mean, else >= 1).
    remove_mean : bool
        Subtract each block's mean before squaring.

    Returns
    -------
    ndarray
        One energy per whole block.
    """
    x = np.asarray(x, dtype=float).ravel()
    if n_block < 1:
        raise ValueError("n_block must be at least 1")
    if remove_mean and n_block < 2:
        raise ValueError("mean removal needs at least 2 samples per block")
    n_whole = len(x) // n_block
    if n_whole == 0:
        raise ValueError("record shorter than one block")
    blocks = x[:n_whole * n_block].reshape(n_whole, n_block)
    if remove_mean:
        blocks = blocks - blocks.mean(axis=1, keepdims=True)
    return np.sum(blocks**2, axis=1)


def ca_cfar(energies: np.ndarray, alpha: float, n_ref: int,
            n_guard: int = 1) -> tuple[np.ndarray, np.ndarray]:
    """Causal cell-averaging CFAR over a sequence of block energies.

    For each cell k, the reference window is the *n_ref* blocks that
    precede the *n_guard* guard blocks that precede the cell:

        [ ... ref ... | guard | cell ]      (all earlier in time)

    matching what a streaming detector on a microcontroller can do.
    The guard blocks keep a signal that straddles a block boundary
    from leaking into its own noise estimate.  Cells without a full
    reference window behind them get threshold NaN and no detection:
    a CFAR detector is blind until its noise history fills, the same
    warm-up the sliding-window outlier detectors have.

    Parameters
    ----------
    energies : ndarray
        Block energies, e.g. from :func:`block_energies`.
    alpha : float
        Threshold factor, from :func:`cfar_factor` (> 0).
    n_ref : int
        Reference blocks per decision (>= 1).
    n_guard : int
        Guard blocks between reference window and cell (>= 0).

    Returns
    -------
    detections : ndarray of bool
        True where the cell energy exceeds alpha times the reference
        mean.
    thresholds : ndarray
        The threshold alpha * mean(reference) per cell; NaN during
        warm-up.
    """
    e = np.asarray(energies, dtype=float).ravel()
    if alpha <= 0:
        raise ValueError("alpha must be strictly positive")
    if n_ref < 1 or n_guard < 0:
        raise ValueError("n_ref must be >= 1 and n_guard >= 0")
    thresholds = np.full(len(e), np.nan)
    start = n_ref + n_guard
    for k in range(start, len(e)):
        ref = e[k - n_guard - n_ref:k - n_guard]
        thresholds[k] = alpha * ref.mean()
    detections = np.zeros(len(e), dtype=bool)
    valid = ~np.isnan(thresholds)
    detections[valid] = e[valid] > thresholds[valid]
    return detections, thresholds