"""Tests for multirate.py.

These modules live under basics/ and were previously uncovered (pytest.ini used
testpaths=topics), which is how a broken polyphase_decimate shipped. The first
test below is the regression pin for that bug: polyphase decimation must equal
the direct method lfilter(h, 1, x)[::M].
"""

import numpy as np
import pytest
from scipy.signal import lfilter, firwin

from multirate import (
    design_multirate_filter, decimate, interpolate, resample_rational,
    polyphase_decimate, cic_decimate, cic_response, cic_compensator,
)


# ----------------------------------------------------------------- polyphase

@pytest.mark.parametrize("M", [2, 3, 4, 5])
@pytest.mark.parametrize("n", [60, 61, 63, 100])
def test_polyphase_matches_direct(M, n):
    """The regression pin: polyphase == filter-then-downsample, exactly."""
    rng = np.random.default_rng(M * 100 + n)
    x = rng.standard_normal(n)
    h = firwin(4 * M + 1, 0.8 / M)               # a few taps per phase
    direct = lfilter(h, 1.0, x)[::M]
    poly = polyphase_decimate(x, h, M)
    assert poly.shape == direct.shape
    assert np.allclose(poly, direct, atol=1e-9)


def test_polyphase_output_length():
    x = np.arange(50.0)
    h = firwin(13, 0.3)
    assert len(polyphase_decimate(x, h, 4)) == (50 + 3) // 4   # ceil(50/4) = 13


def test_polyphase_short_filter_padding():
    # h length not a multiple of M still works (internal zero-pad of h)
    x = np.random.default_rng(0).standard_normal(40)
    h = firwin(7, 0.3)                            # 7 taps, M = 3 -> padded to 9
    direct = lfilter(h, 1.0, x)[::3]
    assert np.allclose(polyphase_decimate(x, h, 3), direct, atol=1e-9)


# ----------------------------------------------------------------- decimate

def test_decimate_length():
    x = np.zeros(100)
    assert len(decimate(x, 4)) == 25
    assert len(decimate(np.zeros(101), 4)) == 26      # ceil(101/4)


def test_decimate_attenuates_above_new_nyquist():
    fs, M = 8000.0, 4                              # new Nyquist = 1000 Hz
    t = np.arange(4000) / fs
    low = np.sin(2 * np.pi * 200 * t)             # well below 1000 Hz
    high = np.sin(2 * np.pi * 3000 * t)           # above 1000 Hz: must be removed
    rms = lambda z: np.sqrt(np.mean(z ** 2))
    # ignore the filter transient at the start
    assert rms(decimate(high, M)[50:]) < 0.2 * rms(decimate(low, M)[50:])


# --------------------------------------------------------------- interpolate

def test_interpolate_length():
    assert len(interpolate(np.zeros(30), 3)) == 90


def test_interpolate_preserves_dc():
    x = np.ones(50)
    y = interpolate(x, 4)
    assert np.isclose(np.mean(y[40:-40]), 1.0, atol=0.05)   # DC restored by gain L


# ----------------------------------------------------------- rational resample

def test_resample_rational_length():
    x = np.random.default_rng(1).standard_normal(48)
    L, M = 3, 2
    y = resample_rational(x, L, M)
    assert len(y) == (len(x) * L + M - 1) // M      # ceil(N*L / M)


def test_resample_rational_rate():
    # 3/2 conversion of a 60-sample block gives ~90 samples
    assert len(resample_rational(np.zeros(60), 3, 2)) == 90


# ----------------------------------------------------------------- CIC

def test_cic_one_stage_is_decimated_running_sum():
    # a single-stage CIC decimating by M is a length-M running sum, decimated:
    # interior outputs of a constant-1 input equal M.
    x = np.ones(60)
    y = cic_decimate(x, 5, N_stages=1)
    assert len(y) == 12                            # ceil(60/5)
    assert np.allclose(y[1:], 5.0)                 # each = sum of 5 ones


def test_cic_length_multi_stage():
    y = cic_decimate(np.random.default_rng(2).standard_normal(100), 4, N_stages=3)
    assert len(y) == 25


def test_cic_dc_gain():
    # DC gain of an N-stage CIC decimating by M is M**N
    M, N = 4, 3
    y = cic_decimate(np.ones(400), M, N_stages=N)
    assert np.isclose(y[-1], M ** N)               # steady-state interior sample


def test_cic_response_matches_impulse_fft():
    # cic_response(w) must equal the actual (normalised) DTFT magnitude of the
    # unnormalised CIC impulse response ((1-z^-M)/(1-z^-1))^N sampled at those w.
    from numpy.fft import rfft, rfftfreq
    M, N = 8, 3
    imp = np.zeros(4096); imp[0] = 1.0
    # build the high-rate CIC impulse response directly (no decimation)
    b = np.array([1.0] + [0.0] * (M - 1) + [-1.0])
    a = np.array([1.0, -1.0])
    y = imp.copy()
    for _ in range(N):
        y = lfilter(b, a, y)
    H = np.abs(rfft(y))
    H /= H[0]                                       # normalise DC to 1
    w = 2 * np.pi * rfftfreq(len(y))               # HIGH-rate frequency axis
    # cic_response is written in output-rate frequency; map w_out = M * w_high
    ref = cic_response(M * w[w * M <= np.pi], M, N)
    got = H[w * M <= np.pi]
    assert np.allclose(ref, got, atol=1e-6)


def test_cic_compensator_flattens_droop():
    M, N, fc = 8, 3, 0.4
    h = cic_compensator(M, N, num_taps=9, fc=fc)
    assert len(h) == 9
    assert np.allclose(h, h[::-1])                  # linear phase (symmetric taps)
    assert np.isclose(h.sum(), 1.0)                 # unit DC gain, adds no gain

    w = np.linspace(0, fc * np.pi, 256)
    cic = cic_response(w, M, N)
    n = np.arange(len(h))
    comp_mag = np.abs(h @ np.exp(-1j * np.outer(n, w)))
    combined = cic * comp_mag

    droop_cic = cic.max() / cic.min()              # uncompensated droop (~1.22, 1.7 dB)
    ripple = combined.max() / combined.min()       # compensated ripple (~1.0)
    assert droop_cic > 1.2                         # the CIC really does droop here
    assert ripple < 1.001                          # compensator flattens to <0.01 dB


def test_cic_compensator_rejects_even_length():
    with pytest.raises(ValueError):
        cic_compensator(8, 3, num_taps=8)


def test_cic_overflow_wrap_is_harmless():
    # Hogenauer's key hardware claim, and the logic behind the VHDL on the
    # embedded page: with B_OUT = B_IN + ceil(K*log2 M) bits, the integrators
    # overflow but two's-complement wrap-around is undone exactly by the combs.
    M, K, B_IN = 8, 3, 16
    B_OUT = B_IN + int(np.ceil(K * np.log2(M)))     # = 25 bits

    def wrap(v, B):
        v &= (1 << B) - 1
        return v - (1 << B) if v & (1 << (B - 1)) else v

    def cic_rtl(x, wrapping):
        integ = [0] * K; comb_d = [0] * K; out = []
        for i, s in enumerate(x):
            val = int(s)
            for k in range(K):
                integ[k] += val
                if wrapping:
                    integ[k] = wrap(integ[k], B_OUT)
                val = integ[k]
            if i % M == 0:
                for k in range(K):
                    prev = comb_d[k]; comb_d[k] = val
                    val -= prev
                    if wrapping:
                        val = wrap(val, B_OUT)
                out.append(val)
        return np.array(out, dtype=np.int64)

    x = np.random.default_rng(0).integers(-(1 << (B_IN - 1)), 1 << (B_IN - 1), 2000)
    ideal = cic_rtl(x, wrapping=False)
    hw = cic_rtl(x, wrapping=True)

    # the integrators really do overflow the B_OUT-bit word (else the test is empty)
    exact = np.cumsum(np.cumsum(np.cumsum(x.astype(object))))
    assert int(np.ceil(np.log2(int(np.max(np.abs(exact)))))) > B_OUT
    # ...yet the wrapped hardware output is bit-identical to the ideal one
    assert np.array_equal(ideal, hw)
    # and both equal the trusted module (phase-0 decimation)
    assert np.array_equal(hw, cic_decimate(x, M, K).astype(np.int64)[:len(hw)])


# ------------------------------------------------------- filter design helper

def test_design_multirate_filter_is_lowpass():
    h = design_multirate_filter(4, num_taps=63)
    assert len(h) == 63
    assert np.isclose(h.sum(), 1.0, atol=1e-6)     # unity DC gain (firwin default)