"""Tests for the STFT module: definition, reconstruction, COLA, spectrogram."""

import numpy as np
import pytest
from scipy.signal import get_window

from stft import stft, istft, cola, spectrogram


def test_shapes_and_axes():
    fs, nperseg, noverlap = 1000.0, 256, 128
    x = np.random.default_rng(0).standard_normal(4000)
    freqs, times, Z = stft(x, fs, "hann", nperseg, noverlap)
    assert Z.shape[0] == nperseg // 2 + 1 == len(freqs)
    assert Z.shape[1] == len(times)
    # frame count: frames that fit fully, stepping by hop
    hop = nperseg - noverlap
    expected = 1 + (len(x) - nperseg) // hop
    assert Z.shape[1] == expected
    assert freqs[0] == 0.0 and np.isclose(freqs[-1], fs / 2)


def test_stft_is_windowed_dft_of_each_frame():
    """Each column is exactly the rfft of the corresponding windowed frame."""
    fs, nperseg, noverlap = 8000.0, 128, 64
    x = np.random.default_rng(1).standard_normal(1000)
    win = get_window("hann", nperseg, fftbins=True)
    _, _, Z = stft(x, fs, "hann", nperseg, noverlap)
    hop = nperseg - noverlap
    for i in range(Z.shape[1]):
        frame = x[i * hop:i * hop + nperseg] * win
        assert np.allclose(Z[:, i], np.fft.rfft(frame), atol=1e-9)


def test_pure_tone_lands_in_one_bin():
    fs, nperseg = 8000.0, 512
    f0 = fs * 32 / nperseg                       # exactly on bin 32, no leakage
    t = np.arange(8000) / fs
    x = np.sin(2 * np.pi * f0 * t)
    freqs, _, S = spectrogram(x, fs, "boxcar", nperseg, nperseg // 2)
    peak_bins = np.argmax(S, axis=0)             # peak bin per frame
    assert np.all(peak_bins == 32)
    assert np.isclose(freqs[32], f0)


@pytest.mark.parametrize("window,noverlap", [("hann", 128), ("hamming", 192),
                                             ("boxcar", 0), ("blackman", 200)])
def test_reconstruction_exact_on_interior(window, noverlap):
    """istft(stft(x)) reproduces x on the covered interior, any window."""
    fs, nperseg = 1000.0, 256
    x = np.random.default_rng(2).standard_normal(5000)
    freqs, times, Z = stft(x, fs, window, nperseg, noverlap)
    _, xr = istft(Z, fs, window, nperseg, noverlap)
    interior = slice(nperseg, len(xr) - nperseg)
    assert np.allclose(xr[interior], x[interior], atol=1e-9)


def test_cola_hann_half_overlap_holds():
    ok, const = cola("hann", 256, 128)           # the canonical COLA case
    assert ok and np.isclose(const, 1.0, atol=1e-6)


def test_cola_fails_without_overlap_for_tapered_window():
    ok, _ = cola("hann", 256, 0)                  # no overlap: gaps at the tapers
    assert not ok


def test_cola_rectangular_no_overlap_holds():
    ok, const = cola("boxcar", 200, 0)            # rectangular tiles perfectly
    assert ok and np.isclose(const, 1.0)


def test_chirp_peak_frequency_rises():
    fs, nperseg = 4000.0, 256
    dur = 2.0
    t = np.arange(int(fs * dur)) / fs
    f0, f1 = 100.0, 800.0
    x = np.sin(2 * np.pi * (f0 * t + (f1 - f0) / (2 * dur) * t ** 2))
    freqs, times, S = spectrogram(x, fs, "hann", nperseg, nperseg // 2)
    peak_freq = freqs[np.argmax(S, axis=0)]
    # monotone increasing peak frequency, spanning roughly f0 to f1
    assert peak_freq[0] < 200 and peak_freq[-1] > 650
    assert np.all(np.diff(peak_freq) >= -freqs[1])   # non-decreasing within a bin


def test_resolution_tradeoff_bin_width():
    """Longer frames give finer frequency bins (delta f = fs / nperseg)."""
    fs = 8000.0
    f_short = np.fft.rfftfreq(128, 1 / fs)
    f_long = np.fft.rfftfreq(1024, 1 / fs)
    assert np.isclose(f_short[1] - f_short[0], fs / 128)
    assert np.isclose(f_long[1] - f_long[0], fs / 1024)
    assert (f_long[1] - f_long[0]) < (f_short[1] - f_short[0])