"""
test_audio_scene.py — Unit tests for the audio scene classifier (Issue #353).

All tests use synthetic audio signals generated with numpy — no microphone,
no audio files, no ROS2 runtime needed.
"""

import math

import numpy as np
import pytest

from saltybot_bringup._audio_scene import (
    SCENE_LABELS,
    AudioSceneResult,
    NearestCentroidClassifier,
    _build_centroids,
    _make_prototype,
    classify_audio,
    extract_features,
    _frame,
    _mel_filterbank,
    _dct2,
    _N_FEATURES,
    _SR_DEFAULT,
    _CLASSIFIER,
    _CENTROIDS,
)


# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────

def _sine(freq_hz: float, duration_s: float = 1.0, sr: int = _SR_DEFAULT,
          amp: float = 0.5) -> np.ndarray:
    """Generate a pure sine wave."""
    n = int(sr * duration_s)
    t = np.linspace(0.0, duration_s, n, endpoint=False)
    return (amp * np.sin(2.0 * math.pi * freq_hz * t)).astype(np.float64)


def _white_noise(duration_s: float = 1.0, sr: int = _SR_DEFAULT,
                 amp: float = 0.1, seed: int = 0) -> np.ndarray:
    """White noise."""
    rng = np.random.RandomState(seed)
    n = int(sr * duration_s)
    return (amp * rng.randn(n)).astype(np.float64)


def _silence(duration_s: float = 1.0, sr: int = _SR_DEFAULT) -> np.ndarray:
    return np.zeros(int(sr * duration_s), dtype=np.float64)


# ─────────────────────────────────────────────────────────────────────────────
# SCENE_LABELS
# ─────────────────────────────────────────────────────────────────────────────

def test_scene_labels_tuple():
    assert isinstance(SCENE_LABELS, tuple)


def test_scene_labels_count():
    assert len(SCENE_LABELS) == 4


def test_scene_labels_values():
    assert set(SCENE_LABELS) == {'indoor', 'outdoor', 'traffic', 'park'}


def test_scene_labels_order():
    assert SCENE_LABELS[0] == 'indoor'
    assert SCENE_LABELS[1] == 'outdoor'
    assert SCENE_LABELS[2] == 'traffic'
    assert SCENE_LABELS[3] == 'park'


# ─────────────────────────────────────────────────────────────────────────────
# _frame
# ─────────────────────────────────────────────────────────────────────────────

def test_frame_shape():
    x = np.arange(1000.0)
    frames = _frame(x, 400, 160)
    assert frames.ndim == 2
    assert frames.shape[1] == 400


def test_frame_n_frames():
    x = np.arange(1000.0)
    frames = _frame(x, 400, 160)
    expected = 1 + (1000 - 400) // 160
    assert frames.shape[0] == expected


def test_frame_first_frame():
    x = np.arange(100.0)
    frames = _frame(x, 10, 5)
    np.testing.assert_array_equal(frames[0], np.arange(10.0))


def test_frame_second_frame():
    x = np.arange(100.0)
    frames = _frame(x, 10, 5)
    np.testing.assert_array_equal(frames[1], np.arange(5.0, 15.0))


def test_frame_short_signal():
    """Signal shorter than one frame should still return one frame."""
    x = np.ones(10)
    frames = _frame(x, 400, 160)
    assert frames.shape[0] == 1


# ─────────────────────────────────────────────────────────────────────────────
# _mel_filterbank
# ─────────────────────────────────────────────────────────────────────────────

def test_mel_filterbank_shape():
    fbank = _mel_filterbank(16000, 512, 26)
    assert fbank.shape == (26, 512 // 2 + 1)


def test_mel_filterbank_non_negative():
    fbank = _mel_filterbank(16000, 512, 26)
    assert (fbank >= 0).all()


def test_mel_filterbank_max_one():
    fbank = _mel_filterbank(16000, 512, 26)
    assert fbank.max() <= 1.0 + 1e-9


def test_mel_filterbank_non_zero():
    fbank = _mel_filterbank(16000, 512, 26)
    assert fbank.sum() > 0


# ─────────────────────────────────────────────────────────────────────────────
# _dct2
# ─────────────────────────────────────────────────────────────────────────────

def test_dct2_shape():
    x = np.random.randn(10, 26)
    out = _dct2(x)
    assert out.shape == x.shape


def test_dct2_dc():
    """DC (constant) row should produce non-zero first column only."""
    x = np.ones((1, 8))
    out = _dct2(x)
    # First coefficient = sum; rest should be near zero for constant input
    assert abs(out[0, 0]) > abs(out[0, 1]) * 5


# ─────────────────────────────────────────────────────────────────────────────
# extract_features
# ─────────────────────────────────────────────────────────────────────────────

def test_extract_features_length():
    sig = _sine(440.0)
    feat = extract_features(sig)
    assert len(feat) == _N_FEATURES


def test_extract_features_dtype():
    sig = _sine(440.0)
    feat = extract_features(sig)
    assert feat.dtype in (np.float32, np.float64)


def test_extract_features_silence_finite():
    """Features from silence should be finite (no NaN/Inf)."""
    feat = extract_features(_silence())
    assert np.all(np.isfinite(feat))


def test_extract_features_sine_finite():
    feat = extract_features(_sine(440.0))
    assert np.all(np.isfinite(feat))


def test_extract_features_noise_finite():
    feat = extract_features(_white_noise())
    assert np.all(np.isfinite(feat))


def test_extract_features_centroid_low_freq():
    """Low-frequency sine → low spectral centroid (feature index 13)."""
    feat_lo = extract_features(_sine(80.0))
    feat_hi = extract_features(_sine(4000.0))
    assert feat_lo[13] < feat_hi[13]


def test_extract_features_centroid_high_freq():
    """High-frequency sine → high spectral centroid."""
    feat_4k = extract_features(_sine(4000.0))
    feat_200 = extract_features(_sine(200.0))
    assert feat_4k[13] > feat_200[13]


def test_extract_features_rolloff_low_freq():
    """Low-frequency sine → low spectral rolloff (feature index 14)."""
    feat_lo = extract_features(_sine(100.0))
    feat_hi = extract_features(_sine(5000.0))
    assert feat_lo[14] < feat_hi[14]


def test_extract_features_zcr_increases_with_freq():
    """Higher-frequency sine → higher ZCR (feature index 15)."""
    feat_lo = extract_features(_sine(100.0))
    feat_hi = extract_features(_sine(3000.0))
    assert feat_lo[15] < feat_hi[15]


def test_extract_features_short_clip():
    """Clips shorter than one frame should not crash."""
    sig = np.zeros(100)
    feat = extract_features(sig)
    assert len(feat) == _N_FEATURES


def test_extract_features_empty_clip():
    """Empty signal should return zero vector."""
    feat = extract_features(np.array([]))
    assert len(feat) == _N_FEATURES
    assert np.all(feat == 0.0)


# ─────────────────────────────────────────────────────────────────────────────
# _make_prototype
# ─────────────────────────────────────────────────────────────────────────────

def test_prototype_length():
    for label in SCENE_LABELS:
        sig = _make_prototype(label)
        assert len(sig) == _SR_DEFAULT


def test_prototype_finite():
    for label in SCENE_LABELS:
        sig = _make_prototype(label)
        assert np.all(np.isfinite(sig))


def test_prototype_unknown_label():
    with pytest.raises(ValueError, match='Unknown scene label'):
        _make_prototype('underwater')


def test_prototype_deterministic():
    """Same seed → same prototype."""
    rng1 = np.random.RandomState(42)
    rng2 = np.random.RandomState(42)
    sig1 = _make_prototype('indoor', rng=rng1)
    sig2 = _make_prototype('indoor', rng=rng2)
    np.testing.assert_array_equal(sig1, sig2)


# ─────────────────────────────────────────────────────────────────────────────
# _build_centroids
# ─────────────────────────────────────────────────────────────────────────────

def test_build_centroids_shape():
    C = _build_centroids()
    assert C.shape == (4, _N_FEATURES)


def test_build_centroids_finite():
    C = _build_centroids()
    assert np.all(np.isfinite(C))


def test_centroids_traffic_low_centroid():
    """Traffic prototype centroid (col 13) should be the lowest."""
    C = _build_centroids()
    idx_traffic = list(SCENE_LABELS).index('traffic')
    assert C[idx_traffic, 13] == C[:, 13].min()


def test_centroids_park_high_centroid():
    """Park prototype centroid should have the highest spectral centroid."""
    C = _build_centroids()
    idx_park = list(SCENE_LABELS).index('park')
    assert C[idx_park, 13] == C[:, 13].max()


# ─────────────────────────────────────────────────────────────────────────────
# NearestCentroidClassifier
# ─────────────────────────────────────────────────────────────────────────────

def test_classifier_predict_returns_known_label():
    feat = extract_features(_sine(440.0))
    label, conf = _CLASSIFIER.predict(feat)
    assert label in SCENE_LABELS


def test_classifier_confidence_range():
    feat = extract_features(_sine(440.0))
    _, conf = _CLASSIFIER.predict(feat)
    assert 0.0 < conf <= 1.0


def test_classifier_prototype_self_predict():
    """Each prototype signal should classify as its own class."""
    rng = np.random.RandomState(42)
    for label in SCENE_LABELS:
        sig = _make_prototype(label, rng=rng)
        feat = extract_features(sig)
        pred, _ = _CLASSIFIER.predict(feat)
        assert pred == label, (
            f'Prototype {label!r} classified as {pred!r}'
        )


def test_classifier_custom_centroids():
    """Simple 1-D centroid classifier sanity check."""
    # Two classes: 'low' centroid at 0, 'high' centroid at 10
    C = np.array([[0.0], [10.0]])
    clf = NearestCentroidClassifier(C, ('low', 'high'))
    assert clf.predict(np.array([1.0]))[0] == 'low'
    assert clf.predict(np.array([9.0]))[0] == 'high'


def test_classifier_confidence_max_at_centroid():
    """Feature equal to a centroid → maximum confidence for that class."""
    C = _CENTROIDS
    clf = NearestCentroidClassifier(C, SCENE_LABELS)
    for i, label in enumerate(SCENE_LABELS):
        pred, conf = clf.predict(C[i])
        assert pred == label
        # dist = 0 → conf = 1/(1+0) = 1.0
        assert abs(conf - 1.0) < 1e-9


# ─────────────────────────────────────────────────────────────────────────────
# classify_audio — end-to-end
# ─────────────────────────────────────────────────────────────────────────────

def test_classify_audio_returns_result():
    res = classify_audio(_sine(440.0))
    assert isinstance(res, AudioSceneResult)


def test_classify_audio_label_valid():
    res = classify_audio(_sine(440.0))
    assert res.label in SCENE_LABELS


def test_classify_audio_confidence_range():
    res = classify_audio(_sine(440.0))
    assert 0.0 < res.confidence <= 1.0


def test_classify_audio_features_length():
    res = classify_audio(_sine(440.0))
    assert len(res.features) == _N_FEATURES


def test_classify_traffic_low_freq():
    """80 Hz sine → traffic."""
    res = classify_audio(_sine(80.0))
    assert res.label == 'traffic', (
        f'Expected traffic, got {res.label!r} (conf={res.confidence:.2f})'
    )


def test_classify_indoor_440():
    """440 Hz sine → indoor."""
    res = classify_audio(_sine(440.0))
    assert res.label == 'indoor', (
        f'Expected indoor, got {res.label!r} (conf={res.confidence:.2f})'
    )


def test_classify_outdoor_mid_freq():
    """Mixed 1000+2000 Hz → outdoor."""
    sr = _SR_DEFAULT
    t = np.linspace(0.0, 1.0, sr, endpoint=False)
    sig = (
        0.4 * np.sin(2 * math.pi * 1000.0 * t)
        + 0.4 * np.sin(2 * math.pi * 2000.0 * t)
    )
    res = classify_audio(sig)
    assert res.label == 'outdoor', (
        f'Expected outdoor, got {res.label!r} (conf={res.confidence:.2f})'
    )


def test_classify_park_high_freq():
    """3200+4800 Hz tones → park."""
    sr = _SR_DEFAULT
    t = np.linspace(0.0, 1.0, sr, endpoint=False)
    sig = (
        0.4 * np.sin(2 * math.pi * 3200.0 * t)
        + 0.3 * np.sin(2 * math.pi * 4800.0 * t)
    )
    res = classify_audio(sig)
    assert res.label == 'park', (
        f'Expected park, got {res.label!r} (conf={res.confidence:.2f})'
    )


def test_classify_silence_no_crash():
    """Silence should return a valid (if low-confidence) result."""
    res = classify_audio(_silence())
    assert res.label in SCENE_LABELS
    assert np.isfinite(res.confidence)


def test_classify_audio_short_clip():
    """Short clip (200 ms) should not crash."""
    sig = _sine(440.0, duration_s=0.2)
    res = classify_audio(sig)
    assert res.label in SCENE_LABELS


def test_classify_prototype_indoor():
    rng = np.random.RandomState(42)
    sig = _make_prototype('indoor', rng=rng)
    res = classify_audio(sig)
    assert res.label == 'indoor'


def test_classify_prototype_outdoor():
    rng = np.random.RandomState(42)
    sig = _make_prototype('outdoor', rng=rng)
    res = classify_audio(sig)
    assert res.label == 'outdoor'


def test_classify_prototype_traffic():
    rng = np.random.RandomState(42)
    sig = _make_prototype('traffic', rng=rng)
    res = classify_audio(sig)
    assert res.label == 'traffic'


def test_classify_prototype_park():
    rng = np.random.RandomState(42)
    sig = _make_prototype('park', rng=rng)
    res = classify_audio(sig)
    assert res.label == 'park'
