"""Tests for the MFCC module.

Each stage is pinned to a closed form or a structural property: the mel warp is
checked at its anchor point and as an exact inverse, the DCT against scipy and
its own orthonormality, the filterbank against the triangle geometry, deltas
against the slope of a known ramp, and the full pipeline by shape and by its
ability to separate a tone from noise.
"""

import numpy as np
import pytest
from scipy.fft import dct as scipy_dct

from mfcc import (
    hz_to_mel,
    mel_to_hz,
    pre_emphasis,
    frame_signal,
    power_spectrum,
    mel_filterbank,
    log_mel_spectrogram,
    dct2,
    mfcc,
    delta,
)


# --- mel warp ---------------------------------------------------------------

def test_mel_of_zero_is_zero():
    assert hz_to_mel(0.0) == 0.0


def test_mel_anchor_1000hz_is_about_1000_mel():
    # the 2595/700 form is tuned so 1000 Hz lands near 1000 mel
    assert np.isclose(hz_to_mel(1000.0), 1000.0, atol=1.5)


def test_mel_round_trip_is_exact():
    f = np.array([0.0, 100.0, 700.0, 1000.0, 4000.0, 8000.0])
    assert np.allclose(mel_to_hz(hz_to_mel(f)), f)


def test_mel_is_monotonic():
    f = np.linspace(0.0, 8000.0, 50)
    m = hz_to_mel(f)
    assert np.all(np.diff(m) > 0)


# --- pre-emphasis -----------------------------------------------------------

def test_pre_emphasis_closed_form():
    x = np.array([1.0, 2.0, 3.0, 4.0])
    y = pre_emphasis(x, coeff=0.97)
    assert y[0] == 1.0                                   # first sample untouched
    assert np.allclose(y[1:], x[1:] - 0.97 * x[:-1])


def test_pre_emphasis_kills_dc_gradually():
    # a constant gets flattened toward (1 - coeff) of its level after sample 0
    x = np.full(10, 5.0)
    y = pre_emphasis(x, 0.97)
    assert np.allclose(y[1:], 5.0 * (1.0 - 0.97))


# --- framing ----------------------------------------------------------------

def test_frame_count_and_shape():
    x = np.zeros(1000)
    frames = frame_signal(x, frame_len=400, hop=160, window="hamming")
    # frames start at 0, 160, 320, 480 (next would need samples past 1000)
    assert frames.shape == (4, 400)


def test_frame_applies_the_window():
    # for a flat unit signal, each frame equals the window itself
    from scipy.signal import get_window
    x = np.ones(400)
    win = get_window("hamming", 400, fftbins=True)
    frames = frame_signal(x, 400, 160, window="hamming")
    assert np.allclose(frames[0], win)
    assert frames[0].max() <= 1.0 + 1e-9


def test_frame_too_short_returns_empty():
    frames = frame_signal(np.zeros(100), 400, 160)
    assert frames.shape == (0, 400)


# --- power spectrum ---------------------------------------------------------

def test_power_spectrum_shape():
    frames = np.zeros((3, 400))
    p = power_spectrum(frames, n_fft=512)
    assert p.shape == (3, 512 // 2 + 1)


def test_power_spectrum_peaks_at_the_tone_bin():
    fs, n_fft = 16000.0, 512
    k = 64                                               # an exact FFT bin
    n = np.arange(n_fft)
    tone = np.cos(2 * np.pi * k * n / n_fft)
    p = power_spectrum(tone[None, :], n_fft)[0]
    assert np.argmax(p) == k


# --- mel filterbank ---------------------------------------------------------

def test_filterbank_shape():
    fb = mel_filterbank(26, n_fft=512, fs=16000.0)
    assert fb.shape == (26, 512 // 2 + 1)


def test_filterbank_is_nonnegative():
    fb = mel_filterbank(26, 512, 16000.0)
    assert np.all(fb >= 0.0)


def test_filterbank_triangles_peak_at_one():
    # each triangle reaches 1 at its centre frequency (which lands on or between
    # bins); the per-row maximum is at most 1 and close to it
    fb = mel_filterbank(26, 1024, 16000.0)              # finer bins -> closer to 1
    row_max = fb.max(axis=1)
    assert np.all(row_max <= 1.0 + 1e-9)
    assert np.all(row_max > 0.9)


def test_filterbank_centres_are_mel_spaced():
    # the argmax bin of each triangle should march up in frequency, with spacing
    # that widens (mel spacing is near-linear low, near-log high)
    fb = mel_filterbank(20, 1024, 16000.0)
    centres = fb.argmax(axis=1)
    assert np.all(np.diff(centres) > 0)                  # strictly increasing
    gaps = np.diff(centres)
    assert gaps[-1] > gaps[0]                            # wider at the top


# --- DCT --------------------------------------------------------------------

def test_dct2_matches_scipy():
    rng = np.random.default_rng(0)
    x = rng.standard_normal(26)
    assert np.allclose(dct2(x), scipy_dct(x, type=2, norm="ortho"))


def test_dct2_is_orthonormal():
    # the DCT-II matrix has orthonormal rows: applying it to the identity columns
    # and checking D D^T = I
    N = 16
    eye = np.eye(N)
    D = dct2(eye, axis=-1)                               # each row transformed
    assert np.allclose(D @ D.T, np.eye(N), atol=1e-12)


def test_dct2_along_axis_matches_per_row():
    rng = np.random.default_rng(1)
    X = rng.standard_normal((5, 26))
    out = dct2(X, axis=-1)
    for i in range(5):
        assert np.allclose(out[i], dct2(X[i]))


def test_dct2_compacts_a_smooth_signal():
    # a smooth ramp puts almost all its energy in the first few coefficients
    x = np.linspace(-1.0, 1.0, 32)
    c = dct2(x)
    head = np.sum(c[:4] ** 2)
    tail = np.sum(c[4:] ** 2)
    assert head > 50 * tail


# --- log mel + full pipeline ------------------------------------------------

def test_log_mel_shape():
    fs = 16000.0
    x = np.random.default_rng(2).standard_normal(16000)
    lm = log_mel_spectrogram(x, fs, n_filters=26, frame_len=400, hop=160,
                             n_fft=512)
    assert lm.shape[1] == 26
    assert lm.shape[0] == frame_signal(x, 400, 160).shape[0]


def test_mfcc_shape_keeps_n_mfcc():
    fs = 16000.0
    x = np.random.default_rng(3).standard_normal(16000)
    c = mfcc(x, fs, n_mfcc=13, frame_len=400, hop=160, n_fft=512)
    assert c.shape[1] == 13


def test_mfcc_silence_is_finite():
    # the log floor keeps a silent frame from producing -inf or nan
    c = mfcc(np.zeros(16000), 16000.0, n_mfcc=13)
    assert np.all(np.isfinite(c))


def test_mfcc_separates_tone_from_noise():
    fs = 16000.0
    n = np.arange(fs)
    tone = np.sin(2 * np.pi * 440.0 * n / fs)
    noise = np.random.default_rng(4).standard_normal(int(fs))
    c_tone = mfcc(tone, fs).mean(axis=0)
    c_noise = mfcc(noise, fs).mean(axis=0)
    # the mean MFCC vectors should be clearly different signals
    assert np.linalg.norm(c_tone - c_noise) > 1.0


def test_dct_decorrelates_the_bands():
    # the headline property: log mel bands are correlated, MFCCs much less so.
    # White noise already has near-independent bands and nothing to decorrelate;
    # the realistic case is a shared, drifting spectral level (here a slow
    # amplitude modulation) that makes every band rise and fall together. That
    # common offset is exactly what the DCT pulls out into the low coefficients.
    fs = 16000.0
    rng = np.random.default_rng(5)
    t = np.arange(int(4 * fs)) / fs
    env = 1.0 + 0.9 * np.sin(2 * np.pi * 2.0 * t)        # slow level drift
    x = env * rng.standard_normal(t.size)
    log_mel = log_mel_spectrogram(x, fs, n_filters=26)
    cep = dct2(log_mel, axis=-1)[:, :13]

    def offdiag_ratio(M):
        c = np.corrcoef(M, rowvar=False)
        off = c - np.diag(np.diag(c))
        return np.sqrt(np.mean(off ** 2))

    assert offdiag_ratio(cep) < offdiag_ratio(log_mel)


# --- deltas -----------------------------------------------------------------

def test_delta_of_a_ramp_is_the_slope():
    # feat[t] = slope*t (per coefficient): interior deltas equal the slope
    slope = np.array([2.0, -3.0])
    t = np.arange(20).reshape(-1, 1)
    feat = t * slope
    d = delta(feat, width=2)
    # check the interior frames (edges use replicated padding)
    assert np.allclose(d[2:-2], slope)


def test_delta_of_a_constant_is_zero():
    feat = np.ones((15, 4))
    assert np.allclose(delta(feat), 0.0)


def test_delta_preserves_shape():
    feat = np.random.default_rng(6).standard_normal((30, 13))
    assert delta(feat).shape == (30, 13)


def test_delta_rejects_bad_width():
    with pytest.raises(ValueError):
        delta(np.ones((5, 3)), width=0)
