import numpy as np

from cymsim.decimation import DECIMATION_SPAN, StreamingDecimator, decimation_kernel, decimation_taps


def test_kernel_is_normalised_linear_phase_and_attenuates_the_stopband():
    m, fs = 54, 44100
    h = decimation_kernel(m)
    W = DECIMATION_SPAN
    assert len(h) == W * m and np.isclose(h.sum(), 1.0)
    d = np.arange(-W * m // 2, W * m // 2)
    assert np.allclose(h[1:], h[::-1][:-1])                      # symmetric about d = 0 (odd length ignored)
    f = np.array([1000.0, 15000.0, 18000.0, 22050.0, 24000.0, 30000.0])
    H = np.abs(np.exp(-2j * np.pi * np.outer(f / (fs * m), d)) @ h)
    assert H[0] > 0.999 and H[1] > 0.99 and H[2] > 0.95         # passband
    assert H[3] < 0.1 and H[4] < 1e-3 and H[5] < 1e-4           # transition and stopband


def test_polyphase_layout_matches_the_kernel():
    m = 10
    h, taps = decimation_kernel(m), decimation_taps(m)
    W = DECIMATION_SPAN
    assert taps.shape == (W, m)
    for w, kk in [(0, 0), (7, 3), (W - 1, m - 1)]:
        d = kk + (W // 2 - 1 - w) * m
        assert taps[w, kk] == h[d + W * m // 2]
    assert np.isclose(taps.sum(), 1.0)


def test_streaming_decimator_matches_direct_convolution_for_any_chunking():
    m = 7
    h = decimation_kernel(m)
    W = DECIMATION_SPAN
    rng = np.random.default_rng(0)
    x = rng.standard_normal(m * 400)
    xz = np.concatenate([np.zeros(W * m // 2), x])           # rest before t = 0
    n_final = (len(x) - W * m // 2) // m + 1
    direct = np.array([xz[j * m:j * m + W * m] @ h for j in range(n_final)])
    dec = StreamingDecimator(m)
    out, i = [], 0
    for L in [1, 3, 50, 999, 7, 300, 2000]:
        out.append(dec.push(x[i:i + L]))
        i += L
    out = np.concatenate(out)
    assert len(out) == n_final
    assert np.abs(out - direct).max() < 1e-12
    # resuming from the saved tail continues the stream exactly
    dec2 = StreamingDecimator(m, tail=dec.tail, n_seen=dec.n_seen)
    more = rng.standard_normal(m * 50)
    a = dec.push(more)
    b = dec2.push(more)
    assert len(a) == 50 and np.array_equal(a, b)