"""Unit tests for the gammatone auditory filter bank."""

import numpy as np
import pytest

from gammatone import (
    erb,
    hz_to_erb_rate,
    erb_rate_to_hz,
    erb_space,
    gammatone_sos,
    make_filterbank,
    gammatone_filterbank,
    cochleagram,
)


# ---------------------------------------------------------------------------
# ERB scale
# ---------------------------------------------------------------------------

class TestERB:
    def test_known_value_at_1khz(self):
        # Glasberg & Moore (1990): 24.7 * (0.00437 * 1000 + 1)
        assert erb(1000) == pytest.approx(132.639, abs=1e-3)

    def test_monotonic_increasing(self):
        f = np.linspace(100, 8000, 50)
        assert np.all(np.diff(erb(f)) > 0)

    def test_vectorised(self):
        assert erb([1000, 2000]).shape == (2,)


class TestERBRate:
    def test_roundtrip(self):
        f = np.array([100, 500, 1000, 4000, 8000.0])
        np.testing.assert_allclose(erb_rate_to_hz(hz_to_erb_rate(f)), f, rtol=1e-9)

    def test_erb_space_endpoints(self):
        cfs = erb_space(100, 8000, 16)
        assert cfs[0] == pytest.approx(100.0)
        assert cfs[-1] == pytest.approx(8000.0)

    def test_erb_space_denser_at_low_frequencies(self):
        cfs = erb_space(100, 8000, 16)
        gaps = np.diff(cfs)
        # ERB spacing widens with frequency: later gaps exceed earlier ones.
        assert gaps[0] < gaps[-1]

    def test_erb_space_ascending(self):
        cfs = erb_space(50, 10000, 32)
        assert np.all(np.diff(cfs) > 0)


# ---------------------------------------------------------------------------
# Single-channel design
# ---------------------------------------------------------------------------

class TestGammatoneSOS:
    def test_factors_into_four_biquads(self):
        sos = gammatone_sos(1000, 16000)
        assert sos.shape == (4, 6)

    def test_peak_at_centre_frequency(self):
        from scipy import signal
        fs, fc = 16000, 1000
        sos = gammatone_sos(fc, fs)
        w, H = signal.sosfreqz(sos, worN=8192, fs=fs)
        peak = w[np.argmax(np.abs(H))]
        # Auditory filter peaks within one ERB of its design frequency.
        assert abs(peak - fc) < erb(fc)

    def test_rejects_frequency_above_nyquist(self):
        with pytest.raises(ValueError):
            gammatone_sos(9000, 16000)

    def test_rejects_nonpositive_frequency(self):
        with pytest.raises(ValueError):
            gammatone_sos(0, 16000)

    def test_rejects_unsupported_order(self):
        # SciPy's IIR gammatone is fixed at order 4; anything else is refused
        # rather than silently ignored.
        with pytest.raises(ValueError):
            gammatone_sos(1000, 16000, order=3)


# ---------------------------------------------------------------------------
# Filter bank
# ---------------------------------------------------------------------------

class TestFilterbank:
    def test_bank_size(self):
        cfs = erb_space(100, 7000, 24)
        bank = make_filterbank(16000, cfs)
        assert len(bank) == 24
        assert all(sos.shape == (4, 6) for sos in bank)

    def test_output_shape(self):
        cfs = erb_space(100, 7000, 12)
        x = np.random.randn(2000)
        out = gammatone_filterbank(x, 16000, cfs)
        assert out.shape == (12, 2000)

    def test_rejects_2d_input(self):
        cfs = erb_space(100, 7000, 4)
        with pytest.raises(ValueError):
            gammatone_filterbank(np.zeros((2, 100)), 16000, cfs)

    def test_tone_excites_nearest_channel(self):
        fs = 16000
        cfs = erb_space(100, 7000, 32)
        t = np.arange(fs) / fs
        f0 = 1000.0
        x = np.sin(2 * np.pi * f0 * t)
        out = gammatone_filterbank(x, fs, cfs)
        # Steady-state energy peaks in the channel nearest the tone.
        energy = np.mean(out[:, fs // 2:] ** 2, axis=1)
        winner = cfs[np.argmax(energy)]
        nearest = cfs[np.argmin(np.abs(cfs - f0))]
        assert winner == pytest.approx(nearest)


# ---------------------------------------------------------------------------
# Cochleagram
# ---------------------------------------------------------------------------

class TestCochleagram:
    def test_shapes_match(self):
        fs = 16000
        cfs = erb_space(100, 7000, 20)
        x = np.random.randn(fs)
        coch, times = cochleagram(x, fs, cfs, frame_ms=20, hop_ms=10)
        assert coch.shape[0] == 20
        assert coch.shape[1] == times.shape[0]

    def test_db_ceiling_is_zero(self):
        fs = 16000
        cfs = erb_space(100, 7000, 20)
        x = np.random.randn(fs)
        coch, _ = cochleagram(x, fs, cfs)
        # Normalised to the per-utterance maximum, so the ceiling is 0 dB.
        assert coch.max() == pytest.approx(0.0, abs=1e-6)