tests/test_features_and_baselines.py

"""
Unit tests for feature extraction and baselines.

IMPORTANT: every audio signal in this file is SYNTHETIC (generated sine
waves / white noise / silence), constructed purely to validate that the
feature-math and baseline logic run correctly and behave sensibly on known
signal shapes. This is standard unit-testing practice for signal-processing
code. It is explicitly NOT a substitute for running against the real
pipecat-ai/smart-turn-data-v3.2-train dataset, and no result here is ever
reported as a dataset metric. See experiments/EXPERIMENTS.md, where every
real experiment row remains marked NOT RUN until it is run against real
data.
"""

import numpy as np
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from turn_detector.features import (
    frame_signal, rms_per_frame, zcr_per_frame, spectral_centroid,
    spectral_bandwidth, spectral_rolloff, mfcc, trailing_silence_duration,
    energy_slope, extract_features, FeatureConfig,
)
from turn_detector.audio_io import pad_or_trim, resample_linear
from turn_detector.baseline_energy import EnergySilenceBaseline, EnergySilenceConfig, tune_threshold
from turn_detector.baseline_classifier import ClassifierBaseline, build_feature_matrix
from turn_detector.splits import split_random, split_source_aware, source_overlap
from turn_detector.evaluation import classification_metrics, slice_metrics
from turn_detector.data import stratified_reservoir_sample, duration_bucket

SR = 16_000


def make_tone(freq_hz, duration_sec, sr=SR, amplitude=0.5):
    t = np.arange(int(sr * duration_sec)) / sr
    return (amplitude * np.sin(2 * np.pi * freq_hz * t)).astype(np.float32)


def make_silence(duration_sec, sr=SR):
    return np.zeros(int(sr * duration_sec), dtype=np.float32)


def make_speech_then_silence(speech_sec=1.5, silence_sec=0.5, sr=SR):
    speech = make_tone(180, speech_sec, sr, amplitude=0.3) + 0.02 * np.random.default_rng(0).standard_normal(int(sr * speech_sec)).astype(np.float32)
    silence = make_silence(silence_sec, sr)
    return np.concatenate([speech, silence]).astype(np.float32)


# ---------------------------------------------------------------------------
# frame_signal / rms / zcr
# ---------------------------------------------------------------------------

def test_frame_signal_shape():
    audio = np.arange(1000, dtype=np.float32)
    frames = frame_signal(audio, frame_len=100, hop_len=50)
    assert frames.shape[1] == 100
    assert frames.shape[0] == 1 + (1000 - 100) // 50


def test_rms_silence_near_zero():
    silence = make_silence(0.5)
    frames = frame_signal(silence, 400, 160)
    rms = rms_per_frame(frames)
    assert np.all(rms < 1e-4)


def test_rms_tone_greater_than_silence():
    tone = make_tone(220, 0.5)
    silence = make_silence(0.5)
    rms_tone = rms_per_frame(frame_signal(tone, 400, 160)).mean()
    rms_sil = rms_per_frame(frame_signal(silence, 400, 160)).mean()
    assert rms_tone > rms_sil


def test_zcr_high_freq_greater_than_low_freq():
    low = make_tone(100, 0.5)
    high = make_tone(3000, 0.5)
    zcr_low = zcr_per_frame(frame_signal(low, 400, 160)).mean()
    zcr_high = zcr_per_frame(frame_signal(high, 400, 160)).mean()
    assert zcr_high > zcr_low


# ---------------------------------------------------------------------------
# spectral features
# ---------------------------------------------------------------------------

def test_spectral_centroid_tracks_frequency():
    low = make_tone(200, 0.5)
    high = make_tone(4000, 0.5)
    c_low = spectral_centroid(frame_signal(low, 800, 400), SR).mean()
    c_high = spectral_centroid(frame_signal(high, 800, 400), SR).mean()
    assert c_high > c_low
    # centroid should be in the right ballpark (within an octave) of the tone freq
    assert 100 < c_low < 800
    assert 2000 < c_high < 8000


def test_spectral_rolloff_below_nyquist():
    tone = make_tone(1000, 0.5)
    rolloff = spectral_rolloff(frame_signal(tone, 800, 400), SR)
    assert np.all(rolloff <= SR / 2)
    assert np.all(rolloff >= 0)


def test_mfcc_shape():
    tone = make_tone(220, 0.5)
    frames = frame_signal(tone, 400, 160)
    coeffs = mfcc(frames, SR, n_mfcc=13, n_mels=26)
    assert coeffs.shape == (frames.shape[0], 13)
    assert np.all(np.isfinite(coeffs))


# ---------------------------------------------------------------------------
# silence duration / energy slope
# ---------------------------------------------------------------------------

def test_trailing_silence_duration_detects_appended_silence():
    audio = make_speech_then_silence(speech_sec=1.0, silence_sec=0.6)
    dur = trailing_silence_duration(audio, SR, silence_rms_threshold=0.05)
    # should detect roughly the appended silence duration (allow tolerance
    # for frame/hop quantization)
    assert 0.3 < dur <= 0.7


def test_trailing_silence_zero_when_ends_in_speech():
    audio = make_tone(180, 0.8, amplitude=0.4)
    dur = trailing_silence_duration(audio, SR, silence_rms_threshold=0.01)
    assert dur == 0.0


def test_energy_slope_sign_on_decaying_energy():
    rms = np.array([0.5, 0.4, 0.3, 0.2, 0.1])
    assert energy_slope(rms) < 0


def test_energy_slope_sign_on_rising_energy():
    rms = np.array([0.1, 0.2, 0.3, 0.4, 0.5])
    assert energy_slope(rms) > 0


# ---------------------------------------------------------------------------
# audio_io
# ---------------------------------------------------------------------------

def test_pad_or_trim_left_pads_short_clip():
    audio = np.ones(100, dtype=np.float32)
    out = pad_or_trim(audio, 200, pad_side="left")
    assert len(out) == 200
    assert np.all(out[:100] == 0)
    assert np.all(out[100:] == 1)


def test_pad_or_trim_trims_long_clip_keeps_tail():
    audio = np.arange(200, dtype=np.float32)
    out = pad_or_trim(audio, 100, pad_side="left")
    assert len(out) == 100
    np.testing.assert_array_equal(out, audio[-100:])


def test_resample_linear_changes_length_correctly():
    audio = make_tone(200, 1.0, sr=8000)
    out = resample_linear(audio, sr_in=8000, sr_out=16000)
    assert abs(len(out) - 16000) <= 2


# ---------------------------------------------------------------------------
# full feature extraction
# ---------------------------------------------------------------------------

def test_extract_features_end_to_end_no_nan_on_normal_clip():
    audio = make_speech_then_silence(1.0, 0.5)
    feats = extract_features(audio, FeatureConfig(sr=SR))
    assert feats["too_short_for_framing"] is False
    for k, v in feats.items():
        if isinstance(v, float) and "tail" not in k:
            assert np.isfinite(v), f"{k} is not finite: {v}"


def test_extract_features_handles_very_short_clip_gracefully():
    audio = make_tone(200, 0.01)  # 10ms — shorter than default frame length
    feats = extract_features(audio, FeatureConfig(sr=SR))
    assert "duration_sec" in feats
    # Should not raise, and should flag itself as too short for framing.
    assert feats["too_short_for_framing"] is True


def test_recent_window_features_flag_unavailable_when_clip_too_short():
    audio = make_tone(200, 0.05)  # 50ms
    feats = extract_features(audio, FeatureConfig(sr=SR, recent_windows_ms=(100, 250)))
    assert feats["tail100ms_available"] is False
    assert feats["tail250ms_available"] is False
    assert np.isnan(feats["tail100ms_rms_mean"])


def test_recent_window_features_available_on_long_clip():
    audio = make_speech_then_silence(1.0, 0.5)
    feats = extract_features(audio, FeatureConfig(sr=SR, recent_windows_ms=(100, 250, 500)))
    assert feats["tail100ms_available"] is True
    assert feats["tail500ms_available"] is True


# ---------------------------------------------------------------------------
# baseline_energy
# ---------------------------------------------------------------------------

def test_energy_baseline_predicts_end_on_long_trailing_silence():
    audio = make_speech_then_silence(speech_sec=1.0, silence_sec=0.8)
    baseline = EnergySilenceBaseline(EnergySilenceConfig(silence_rms_threshold=0.05, silence_duration_threshold_sec=0.3))
    assert baseline.predict(audio) is True


def test_energy_baseline_predicts_continue_on_no_trailing_silence():
    audio = make_tone(180, 1.0, amplitude=0.4)
    baseline = EnergySilenceBaseline(EnergySilenceConfig(silence_rms_threshold=0.01, silence_duration_threshold_sec=0.3))
    assert baseline.predict(audio) is False


def test_tune_threshold_selects_reasonable_threshold_on_synthetic_data():
    rng = np.random.default_rng(1)
    data = []
    for _ in range(20):
        # END examples: long trailing silence
        sil = 0.4 + rng.uniform(0, 0.3)
        data.append((make_speech_then_silence(1.0, sil), True))
        # CONTINUE examples: short/no trailing silence
        sil_short = rng.uniform(0, 0.05)
        data.append((make_speech_then_silence(1.0, sil_short), False))

    cfg = EnergySilenceConfig(silence_rms_threshold=0.05)
    best_t, results = tune_threshold(data, cfg, candidate_thresholds_sec=np.arange(0.05, 0.5, 0.05))
    assert 0.05 <= best_t <= 0.5
    assert results[best_t]["accuracy"] > 0.8  # synthetic data is cleanly separable


# ---------------------------------------------------------------------------
# baseline_classifier
# ---------------------------------------------------------------------------

def test_build_feature_matrix_consistent_shape():
    feats = [{"a": 1.0, "b": 2.0}, {"a": 3.0, "c": 4.0}]
    X, order = build_feature_matrix(feats)
    assert X.shape == (2, 3)
    assert order == ["a", "b", "c"]
    assert np.isnan(X[0, order.index("c")])
    assert np.isnan(X[1, order.index("b")])


def test_classifier_baseline_fits_and_predicts_on_synthetic_data():
    rng = np.random.default_rng(2)
    audios, labels = [], []
    for _ in range(15):
        sil = 0.4 + rng.uniform(0, 0.3)
        audios.append(make_speech_then_silence(1.0, sil)); labels.append(True)
        sil_short = rng.uniform(0, 0.05)
        audios.append(make_speech_then_silence(1.0, sil_short)); labels.append(False)

    clf = ClassifierBaseline()
    clf.fit(audios, labels)
    preds = clf.predict(audios)
    acc = float(np.mean(preds == np.array(labels)))
    assert acc > 0.7  # synthetic, cleanly separable data
    assert clf.param_count() > 0
    assert clf.model_size_bytes() > 0
    top = clf.top_features(5)
    assert len(top) == 5


# ---------------------------------------------------------------------------
# splits
# ---------------------------------------------------------------------------

def test_split_source_aware_has_no_source_overlap():
    records = [{"dataset": f"src{i % 4}", "id": i} for i in range(40)]
    train, val = split_source_aware(records, val_frac=0.25, seed=0)
    assert len(source_overlap(train, val)) == 0
    assert len(train) + len(val) == 40


def test_split_random_can_have_source_overlap():
    records = [{"dataset": f"src{i % 2}", "id": i} for i in range(40)]
    train, val = split_random(records, val_frac=0.5, seed=0)
    # With only 2 sources and a random split, overlap is expected (not
    # guaranteed on every seed, but true for this fixed seed/setup).
    assert len(train) + len(val) == 40


def test_duration_bucket_boundaries():
    assert duration_bucket(0.5) == "<1s"
    assert duration_bucket(1.5) == "1-2s"
    assert duration_bucket(3.0) == "2-4s"
    assert duration_bucket(6.0) == "4-8s"
    assert duration_bucket(10.0) == ">=8s"


# ---------------------------------------------------------------------------
# evaluation
# ---------------------------------------------------------------------------

def test_classification_metrics_perfect_predictions():
    y_true = np.array([True, False, True, False])
    y_pred = np.array([True, False, True, False])
    m = classification_metrics(y_true, y_pred)
    assert m["accuracy"] == 1.0
    assert m["false_end_rate"] == 0.0
    assert m["false_continue_rate"] == 0.0
    assert m["endpoint_latency_ms"] is None


def test_classification_metrics_false_end_rate():
    y_true = np.array([False, False, True, True])
    y_pred = np.array([True, False, True, True])  # one false END
    m = classification_metrics(y_true, y_pred)
    assert m["confusion_matrix"]["fp"] == 1
    assert m["false_end_rate"] == 0.5  # 1 of 2 actual-CONTINUE misclassified


def test_slice_metrics_excludes_small_slices():
    y_true = np.array([True] * 25 + [False] * 5)
    y_pred = np.array([True] * 25 + [True] * 5)
    slices = np.array(["big"] * 25 + ["small"] * 5)
    result = slice_metrics(y_true, y_pred, slices, min_samples=10)
    assert "big" in result["slices"]
    assert "small" not in result["slices"]
    assert result["excluded_insufficient_n"][0]["slice"] == "small"


# ---------------------------------------------------------------------------
# data.py stratified sampling (logic-only test, synthetic fake records)
# ---------------------------------------------------------------------------

def test_stratified_reservoir_sample_respects_target_size():
    rng = np.random.default_rng(3)
    fake_records = []
    for i in range(500):
        fake_records.append({
            "endpoint_bool": bool(rng.integers(0, 2)),
            "language": rng.choice(["eng", "hin", "spa"]),
            "dataset": rng.choice(["src_a", "src_b"]),
            "synthetic": bool(rng.integers(0, 2)),
            "midfiller": bool(rng.integers(0, 2)),
            "endfiller": bool(rng.integers(0, 2)),
            "duration_sec": rng.uniform(0.5, 5.0),
        })
    sample = stratified_reservoir_sample(fake_records, target_n=100, seed=42)
    assert 0 < len(sample) <= 120  # allow small overshoot from rounding across strata


def test_stratified_reservoir_sample_deterministic_given_seed():
    rng = np.random.default_rng(4)
    fake_records = [
        {"endpoint_bool": bool(rng.integers(0, 2)), "language": "eng",
         "dataset": "src_a", "synthetic": True, "midfiller": None,
         "endfiller": None, "duration_sec": 1.0}
        for _ in range(100)
    ]
    s1 = stratified_reservoir_sample(fake_records, target_n=20, seed=7)
    s2 = stratified_reservoir_sample(fake_records, target_n=20, seed=7)
    ids1 = [r.get("duration_sec") for r in s1]
    ids2 = [r.get("duration_sec") for r in s2]
    assert ids1 == ids2


if __name__ == "__main__":
    import pytest
    raise SystemExit(pytest.main([__file__, "-v"]))