import numpy as np
import pytest

from bandsaunter import dsp


def test_decimator_is_stateful_across_blocks():
    """Block-by-block filtering must equal one-shot filtering exactly."""
    rng = np.random.default_rng(0)
    x = (rng.standard_normal(20000) + 1j * rng.standard_normal(20000)).astype(np.complex64)
    d1 = dsp.FIRDecimator(8)
    blocked = np.concatenate([d1(x[:7000]), d1(x[7000:13000]), d1(x[13000:])])
    d2 = dsp.FIRDecimator(8)
    assert np.allclose(blocked, d2(x), atol=1e-9)


def test_decimation_chain_factors():
    assert dsp.design_decimation(128) == [8, 8, 2]
    assert dsp.design_decimation(1) == []
    chain = dsp.DecimationChain(64)
    out = chain(np.ones(6400, dtype=np.complex64))
    assert out.size == 100


def test_mixer_keeps_phase_continuous():
    fs = 48000.0
    m = dsp.Mixer(1000.0, fs)
    a = m(np.ones(1000, dtype=np.complex64))
    b = m(np.ones(1000, dtype=np.complex64))
    joined = np.concatenate([a, b])
    one_shot = dsp.frequency_shift(np.ones(2000, dtype=np.complex64), 1000.0, fs)
    assert np.allclose(joined, one_shot, atol=1e-4)


def test_noise_floor_curve_ignores_signals():
    """A carrier must not raise the floor it is being measured against."""
    psd = np.full(1024, -80.0)
    psd[500:504] = -20.0
    floor = dsp.noise_floor_curve(psd)
    assert floor[502] == pytest.approx(-80.0, abs=1.0)
    assert (psd - floor)[502] > 55.0


def test_peak_hold_beats_averaging_on_bursts():
    fs = 2_048_000
    x = np.zeros(102400, dtype=np.complex64)
    t = np.arange(102400) / fs
    x[:8000] = 0.5 * np.exp(2j * np.pi * 100_000 * t[:8000])
    rng = np.random.default_rng(1)
    x += 0.01 * (rng.standard_normal(102400) + 1j * rng.standard_normal(102400))
    _, p_avg = dsp.welch_psd(x, 1024, combine="mean")
    _, p_max = dsp.welch_psd(x, 1024, combine="max")
    avg = (dsp.db(p_avg) - dsp.noise_floor_curve(dsp.db(p_avg))).max()
    peak = (dsp.db(p_max) - dsp.noise_floor_curve(dsp.db(p_max))).max()
    assert peak > avg


def test_occupied_bandwidth_and_flatness():
    psd = np.zeros(1024)
    psd[500:524] = 1.0
    bw, offset = dsp.occupied_bandwidth(psd, 100.0)
    assert bw == pytest.approx(2400.0, rel=0.15)
    assert dsp.spectral_flatness(np.ones(256)) == pytest.approx(1.0)
    tone = np.full(256, 1e-9)
    tone[128] = 1.0
    assert dsp.spectral_flatness(tone) < 0.01


def test_decimator_matches_a_direct_filter_reference():
    """The polyphase form must equal filtering then discarding samples.

    It replaced an lfilter that computed every output and threw most away;
    that was slow enough to stop captures keeping up with real time, which
    shows up as recordings that play too fast.
    """
    from scipy import signal as sps
    rng = np.random.default_rng(4)
    x = (rng.standard_normal(30000) + 1j * rng.standard_normal(30000)).astype(np.complex64)
    for factor in (2, 4, 8):
        d = dsp.FIRDecimator(factor)
        got = d(x)
        ref = sps.lfilter(d.taps.astype(np.float64), [1.0],
                          x.astype(np.complex128))[::factor]
        assert np.allclose(got, ref[:got.size], atol=1e-4), f"factor {factor}"


def test_quarter_rate_mixer_matches_an_explicit_reference():
    """The trig-free shortcut must be exact, and stay phase-continuous."""
    rng = np.random.default_rng(5)
    n = 9000
    x = (rng.standard_normal(n) + 1j * rng.standard_normal(n)).astype(np.complex64)
    fs = 2_048_000.0
    mixer = dsp.Mixer(fs / 4, fs)
    assert mixer._is_quarter_rate
    got = np.concatenate([mixer(x[:3000]), mixer(x[3000:5000]), mixer(x[5000:])])
    ref = x * np.exp(-2j * np.pi * (fs / 4) * np.arange(n) / fs)
    assert np.allclose(got, ref, atol=1e-5)


def test_decimation_is_fast_enough_for_realtime():
    """A 50 ms block must decimate in well under 50 ms of CPU."""
    import time
    sr, n = 2_048_000, 102_400
    x = (np.random.default_rng(6).standard_normal(n)
         + 1j * np.random.default_rng(7).standard_normal(n)).astype(np.complex64)
    chain = dsp.DecimationChain(64)
    chain(x)
    t0 = time.perf_counter()
    for _ in range(5):
        chain(x)
    per_block = (time.perf_counter() - t0) / 5
    budget = n / sr
    assert per_block < 0.35 * budget, \
        f"{per_block*1000:.1f} ms per {budget*1000:.0f} ms block"