import pytest

from bandsaunter.classify import classify, extract_features
from signals import FS, make

CASES = [
    ("cw", "cw", 14.05e6), ("nfm", "nfm", 146.52e6), ("am", "am", 121.5e6),
    ("usb", "ssb", 14.2e6), ("lsb", "ssb", 7.2e6),
    ("fsk2", "fsk", 929.5e6), ("fsk4", "fsk", 460.0e6),
    ("psk4", "psk", 450.0e6), ("psk2", "psk", 437.0e6),
    ("carrier", "carrier", 446.0e6), ("noise", "unknown", 300e6),
]


@pytest.mark.parametrize("kind,family,freq", CASES)
@pytest.mark.parametrize("seed", [3, 11, 42])
def test_modulation_family(kind, family, freq, seed):
    c = classify(make(kind, seed=seed), FS, freq_hz=freq, snr_db=28)
    assert c.family == family, f"{kind} -> {c.label!r} (family {c.family})"


def test_known_systems_are_named_from_frequency():
    assert "airband" in classify(make("am"), FS, freq_hz=121.5e6,
                                 snr_db=28).label.lower()
    assert "broadcast" in classify(make("wfm"), FS, freq_hz=97.5e6,
                                   snr_db=28).label.lower()
    assert "weather" in classify(make("nfm"), FS, freq_hz=162.475e6,
                                 snr_db=28).label.lower()


def test_ctcss_tone_recovered():
    c = classify(make("nfm"), FS, freq_hz=146.52e6, snr_db=28)
    assert c.features.ctcss_hz == pytest.approx(100.0, abs=0.5)


def test_low_snr_reduces_confidence():
    strong = classify(make("nfm", snr_db=30), FS, freq_hz=146.52e6, snr_db=30)
    weak = classify(make("nfm", snr_db=5), FS, freq_hz=146.52e6, snr_db=5)
    assert weak.confidence < strong.confidence


def test_fsk_level_count():
    assert extract_features(make("fsk2"), FS, snr_db=28).freq_modes == 2
    assert extract_features(make("fsk4"), FS, snr_db=28).freq_modes == 4


def test_symbol_rate_estimate():
    f = extract_features(make("fsk4"), FS, snr_db=28)
    # The cyclostationary line lands on the baud rate or a low harmonic of it.
    assert f.baud > 1000