"""The Goertzel algorithm: a single DFT bin from a second-order IIR filter.

When you only need the spectral content at one frequency (or a handful), a
full FFT is wasteful. Goertzel evaluates one DFT bin with a tiny recursive
filter that costs one real multiply per input sample, and it can run
sample-by-sample as the data arrives, with no input buffer and no
power-of-two length restriction. That makes it the standard choice for tone
detection on microcontrollers, the classic example being a DTMF (telephone
keypad) decoder.

The filter behind it is a second-order IIR section (the structure derived in
the DSPIPS workshop notes, Ch. 3) tuned so its pole pair sits exactly on the
unit circle at the bin frequency. Everything here is plain NumPy and
importable; the topic page imports these functions for its demos, and
``test_goertzel.py`` checks them against ``numpy.fft``.

Reference: G. Goertzel, "An algorithm for the evaluation of finite
trigonometric series," American Mathematical Monthly 65(1), 1958.
"""

from __future__ import annotations

import numpy as np


def goertzel_coeff(omega: float) -> float:
    """The single real recurrence coefficient ``2 cos(omega)``.

    ``omega`` is the bin frequency in radians per sample. This is the only
    multiply that runs inside the per-sample loop, which is what makes the
    algorithm cheap.
    """
    return 2.0 * np.cos(omega)


def goertzel(x, omega: float):
    """Evaluate the DTFT-like sum of ``x`` at angular frequency ``omega``.

    Runs the real second-order recurrence

        s[n] = x[n] + 2 cos(omega) s[n-1] - s[n-2]

    over the ``N`` samples of ``x`` and combines the final two states into the
    complex result

        X = exp(j omega) s[N-1] - s[N-2] = exp(j omega N) sum_n x[n] exp(-j omega n).

    For ``omega = 2 pi k / N`` with integer ``k`` the leading factor
    ``exp(j omega N) = exp(j 2 pi k) = 1``, so ``X`` is exactly the DFT bin
    ``X[k]`` (see ``goertzel_dft_bin``). For any other ``omega`` this is the
    generalized Goertzel: the **magnitude** ``abs(X)`` is the correct DTFT
    magnitude at ``omega``, but the phase carries an extra ``exp(j omega N)``
    twist relative to ``sum_n x[n] exp(-j omega n)``.

    Returns the complex value ``X``.
    """
    x = np.asarray(x, dtype=float)
    coeff = goertzel_coeff(omega)
    s1 = 0.0   # s[n-1]
    s2 = 0.0   # s[n-2]
    for sample in x:
        s0 = sample + coeff * s1 - s2
        s2 = s1
        s1 = s0
    # s1 = s[N-1], s2 = s[N-2]
    return np.exp(1j * omega) * s1 - s2


def goertzel_power(x, omega: float) -> float:
    """Magnitude-squared at ``omega`` without the final complex step.

    Uses the real-only closed form

        |X|^2 = s[N-1]^2 + s[N-2]^2 - 2 cos(omega) s[N-1] s[N-2],

    so no complex arithmetic, no sine, and no square root are needed. This is
    the form a microcontroller runs when it only has to compare a tone's energy
    against a threshold.
    """
    x = np.asarray(x, dtype=float)
    coeff = goertzel_coeff(omega)
    s1 = 0.0
    s2 = 0.0
    for sample in x:
        s0 = sample + coeff * s1 - s2
        s2 = s1
        s1 = s0
    return s1 * s1 + s2 * s2 - coeff * s1 * s2


def goertzel_dft_bin(x, k: int):
    """Exact DFT bin ``X[k]`` of an ``N``-point sequence, via Goertzel.

    Equivalent to ``numpy.fft.fft(x)[k]`` (to floating-point precision) but
    computed with a single second-order filter instead of a full transform.
    ``k`` must be a valid bin index ``0 <= k < N``.
    """
    x = np.asarray(x, dtype=float)
    n = len(x)
    if not 0 <= k < n:
        raise ValueError(f"k={k} is out of range for an {n}-point DFT (need 0 <= k < {n})")
    return goertzel(x, 2.0 * np.pi * k / n)


def goertzel_freq(x, f: float, fs: float):
    """Generalized Goertzel at an arbitrary physical frequency ``f`` (Hz).

    ``omega = 2 pi f / fs`` need not land on a DFT bin, so the tone you are
    detecting can sit anywhere. Returns the complex result; take ``abs(...)``
    for the correct DTFT magnitude at ``f``. (Off a DFT bin the complex phase
    carries the ``exp(j omega N)`` twist noted in :func:`goertzel`, so use the
    magnitude, not the raw phase, for off-grid detection.)
    """
    return goertzel(x, 2.0 * np.pi * f / fs)


# --- DTMF: the canonical Goertzel application -------------------------------

# ITU-T Q.23 dual-tone multi-frequency keypad. Each key is the sum of one
# low-group and one high-group sinusoid.
DTMF_LOW = (697.0, 770.0, 852.0, 941.0)      # rows
DTMF_HIGH = (1209.0, 1336.0, 1477.0, 1633.0)  # columns
DTMF_KEYS = (
    ("1", "2", "3", "A"),
    ("4", "5", "6", "B"),
    ("7", "8", "9", "C"),
    ("*", "0", "#", "D"),
)


def dtmf_tone(key: str, fs: float, duration: float) -> np.ndarray:
    """Synthesize the dual-tone waveform for a DTMF ``key`` (unit amplitude each)."""
    for r, row in enumerate(DTMF_KEYS):
        for c, label in enumerate(row):
            if label == key:
                n = np.arange(int(round(duration * fs)))
                return (np.sin(2 * np.pi * DTMF_LOW[r] * n / fs)
                        + np.sin(2 * np.pi * DTMF_HIGH[c] * n / fs))
    raise ValueError(f"not a DTMF key: {key!r}")


def dtmf_decode(x, fs: float) -> str:
    """Decode a DTMF block to its key by Goertzel energy in the eight bands.

    Picks the strongest low-group and strongest high-group frequency and maps
    the (row, column) pair to a key. Returns the decoded character.
    """
    low_powers = [goertzel_power(x, 2 * np.pi * f / fs) for f in DTMF_LOW]
    high_powers = [goertzel_power(x, 2 * np.pi * f / fs) for f in DTMF_HIGH]
    r = int(np.argmax(low_powers))
    c = int(np.argmax(high_powers))
    return DTMF_KEYS[r][c]