"""Tests for goertzel.py.

The headline guarantee is that Goertzel reproduces a DFT bin exactly, so most
of these tests pin the implementation against ``numpy.fft``.
"""

import numpy as np

from goertzel import (
    goertzel_coeff,
    goertzel,
    goertzel_power,
    goertzel_dft_bin,
    goertzel_freq,
    dtmf_tone,
    dtmf_decode,
    DTMF_KEYS,
    DTMF_LOW,
    DTMF_HIGH,
)


def test_coeff_known_values():
    # Pin the coefficient against closed-form values a reader can check by hand,
    # so a dropped factor of two or a sign error would be caught.
    assert np.isclose(goertzel_coeff(0.0), 2.0)          # 2 cos 0 = 2
    assert np.isclose(goertzel_coeff(np.pi / 3), 1.0)    # 2 cos(pi/3) = 1
    assert np.isclose(goertzel_coeff(np.pi / 2), 0.0)    # 2 cos(pi/2) = 0
    assert np.isclose(goertzel_coeff(np.pi), -2.0)       # 2 cos(pi) = -2


def test_complex_bin_matches_fft():
    rng = np.random.default_rng(0)
    for N in (16, 33, 64, 100):
        x = rng.standard_normal(N)
        X = np.fft.fft(x)
        for k in range(N):
            assert abs(goertzel_dft_bin(x, k) - X[k]) < 1e-9


def test_power_matches_fft_magnitude_squared():
    rng = np.random.default_rng(1)
    N = 80
    x = rng.standard_normal(N)
    X = np.fft.fft(x)
    for k in range(N):
        got = goertzel_power(x, 2 * np.pi * k / N)
        assert abs(got - abs(X[k]) ** 2) < 1e-6 * (abs(X[k]) ** 2 + 1.0)


def test_power_equals_abs_complex_squared():
    # The real-only power form must agree with the complex form it shortcuts.
    rng = np.random.default_rng(2)
    x = rng.standard_normal(50)
    for k in (0, 1, 7, 25):
        omega = 2 * np.pi * k / len(x)
        assert np.isclose(goertzel_power(x, omega), abs(goertzel(x, omega)) ** 2)


def test_dft_bin_rejects_out_of_range_k():
    import pytest
    x = np.arange(8.0)
    for bad_k in (-1, 8, 100):
        with pytest.raises(ValueError):
            goertzel_dft_bin(x, bad_k)


def test_off_grid_result_carries_expected_phase_twist():
    # The docstring claims goertzel(x, omega) = exp(j omega N) * DTFT(x, omega)
    # off the DFT grid. Pin that exact relationship.
    x = np.array([1.0, 2.0, 3.0, -1.0, 0.5])
    N = len(x)
    omega = 0.7  # not a multiple of 2 pi / N
    dtft = np.sum(x * np.exp(-1j * omega * np.arange(N)))
    assert np.isclose(goertzel(x, omega), np.exp(1j * omega * N) * dtft)
    # magnitude is the honest DTFT magnitude regardless of the twist
    assert np.isclose(abs(goertzel(x, omega)), abs(dtft))


def test_generalized_goertzel_is_selective_off_bin():
    # A pure tone gives large magnitude at its own frequency and almost nothing
    # a few bins away, even when the frequency is not on the DFT grid.
    fs = 8000.0
    f = 1234.5
    n = np.arange(2000)
    x = np.sin(2 * np.pi * f * n / fs)
    on = abs(goertzel_freq(x, f, fs))
    off = abs(goertzel_freq(x, f + 100.0, fs))
    assert on > 1000 * off


def test_dtmf_round_trip_all_keys():
    fs = 8000.0
    for row in DTMF_KEYS:
        for key in row:
            tone = dtmf_tone(key, fs, 0.04)   # 40 ms block
            assert dtmf_decode(tone, fs) == key


def test_dtmf_survives_noise():
    rng = np.random.default_rng(3)
    fs = 8000.0
    tone = dtmf_tone("5", fs, 0.04)
    noisy = tone + 0.5 * rng.standard_normal(len(tone))
    assert dtmf_decode(noisy, fs) == "5"


def test_dtmf_frequencies_are_the_itu_grid():
    # Guard against accidental edits to the standard keypad frequencies.
    assert DTMF_LOW == (697.0, 770.0, 852.0, 941.0)
    assert DTMF_HIGH == (1209.0, 1336.0, 1477.0, 1633.0)
    assert DTMF_KEYS[0] == ("1", "2", "3", "A")
    assert DTMF_KEYS[3] == ("*", "0", "#", "D")