import numpy as np

from localizer.config import GammatoneConfig
from localizer.gammatone import GammatoneFilterbank, erb_space


def test_band_count_and_range():
    cf = erb_space(700, 18000, 28)
    assert cf.size == 28
    assert np.all(np.diff(cf) > 0)          # ascending
    assert cf[0] >= 700 - 1
    assert cf[-1] <= 18000 + 1
    assert cf[-1] > 15000                    # top band near the upper edge


def test_unit_gain_at_center_frequency():
    sr = 44100
    fb = GammatoneFilterbank(GammatoneConfig(), sr)
    t = np.arange(sr) / sr
    for idx in (3, 12, 20, 27):
        f = fb.center_freqs[idx]
        tone = np.sin(2 * np.pi * f * t)
        rms = np.sqrt(np.mean(fb.filter(tone) ** 2, axis=1))
        # The tone should peak in its own band, at ~unit gain (rms of a unit sine ≈ 0.707).
        assert int(np.argmax(rms)) == idx
        assert 0.6 < rms[idx] < 0.8


def test_bands_are_bandpass():
    sr = 44100
    fb = GammatoneFilterbank(GammatoneConfig(), sr)
    # DC should be strongly attenuated in every band.
    dc = np.ones(4096)
    rms = np.sqrt(np.mean(fb.filter(dc) ** 2, axis=1))
    assert np.all(rms < 0.1)