tests/test_timeline.py

"""Whole-track scan: window plan, waveform peaks and the scan loop.

The classifier is a stub here; only the audio maths and the bookkeeping run.
"""

from __future__ import annotations

import dataclasses
from typing import List, Optional

import numpy as np
import pytest

from app.services import timeline
from app.services.feature_extractor import AudioFeatures
from app.services.inference_xai import FeatureContribution, WindowReading, XAIInferenceService

SR = 22050


class StubXai:
    """Scores a window by how loud it is, so each test can steer the result."""

    threshold = 0.43

    def __init__(self, probabilities: Optional[List[float]] = None) -> None:
        self._queue = list(probabilities or [])

    def predict_window(self, features, vocals, *, window_sec: float, top: int = 3) -> WindowReading:
        p = self._queue.pop(0) if self._queue else 0.8
        reason = FeatureContribution(
            name="spectral_flatness_mean", label="Spektral düzlük", label_en="Spectral flatness",
            category="spectral", value=0.1, z_score=-1.2, shap_value=0.4 if p >= self.threshold else -0.4,
            direction="towards_ai" if p >= self.threshold else "towards_human", description="",
        )
        return WindowReading(probability=p, is_ai=p >= self.threshold, margin=0.5, reasons=[reason])


@pytest.fixture
def fast_features(monkeypatch):
    """Skip the librosa work: the scan loop is what these tests cover."""
    monkeypatch.setattr(timeline, "extract_features_array", lambda seg, sr: object())
    monkeypatch.setattr(timeline, "analyze_vocals_array", lambda seg, sr: None)


def tone(seconds: float, level: float = 0.5) -> np.ndarray:
    t = np.arange(int(seconds * SR)) / SR
    return (level * np.sin(2 * np.pi * 220 * t)).astype(np.float32)


# ── plan ─────────────────────────────────────────────────────────────

@pytest.mark.parametrize("duration", [3.0, 29.0, 44.9, 45.0, 75.0, 214.2, 359.9])
def test_windows_tile_the_whole_track(duration: float) -> None:
    windows = timeline.plan_windows(duration)
    assert windows[0][0] == 0
    assert windows[-1][1] == pytest.approx(duration, abs=0.01)
    for (_, end), (start, _) in zip(windows, windows[1:]):
        assert end == pytest.approx(start, abs=0.01)


@pytest.mark.parametrize("duration", [45.0, 75.0, 214.2, 359.9])
def test_window_length_stays_near_thirty_seconds(duration: float) -> None:
    for start, end in timeline.plan_windows(duration):
        assert 20.0 <= end - start <= 45.0


def test_nothing_to_plan_for_empty_audio() -> None:
    assert timeline.plan_windows(0) == []


# ── peaks ────────────────────────────────────────────────────────────

def test_peaks_are_normalised_and_bounded() -> None:
    peaks = timeline.waveform_peaks(tone(10.0, 0.2), points=100)
    assert len(peaks) == 100
    assert max(peaks) == 1.0
    assert min(peaks) >= 0.0


def test_peaks_of_silence_are_zero() -> None:
    assert set(timeline.waveform_peaks(np.zeros(SR, dtype=np.float32), points=10)) == {0.0}


def test_peaks_of_a_very_short_clip_do_not_overrun() -> None:
    assert len(timeline.waveform_peaks(np.ones(5, dtype=np.float32), points=600)) == 5


# ── scan ─────────────────────────────────────────────────────────────

def test_scan_scores_every_window_and_summarises(fast_features) -> None:
    y = tone(90.0)  # three windows of 30 s
    out = timeline.scan_track(y, SR, StubXai([0.9, 0.1, 0.7]), total_sec=90.0)

    assert [s["state"] for s in out["segments"]] == ["ok", "ok", "ok"]
    assert [s["isAi"] for s in out["segments"]] == [True, False, True]
    summary = out["summary"]
    assert summary["scoredCount"] == 3
    assert summary["flaggedCount"] == 2
    assert summary["aiShare"] == pytest.approx(2 / 3, abs=0.01)
    assert summary["peakIndex"] == 0
    assert summary["maxProbability"] == 0.9
    assert out["truncated"] is False
    assert out["threshold"] == pytest.approx(0.43)
    assert len(out["peaks"]) > 0
    assert out["segments"][0]["reasons"][0]["direction"] == "towards_ai"


def test_scan_keeps_window_order_with_parallel_workers(fast_features, monkeypatch) -> None:
    monkeypatch.setattr(timeline, "SCAN_WORKERS", 4)
    out = timeline.scan_track(tone(300.0), SR, StubXai(), total_sec=300.0)
    assert [s["index"] for s in out["segments"]] == list(range(10))
    starts = [s["start"] for s in out["segments"]]
    assert starts == sorted(starts)


def test_a_silent_window_is_not_scored(fast_features) -> None:
    y = np.concatenate([tone(30.0), np.zeros(30 * SR, dtype=np.float32)])
    out = timeline.scan_track(y, SR, StubXai([0.8]), total_sec=60.0)
    assert [s["state"] for s in out["segments"]] == ["ok", "silent"]
    assert out["summary"]["scoredCount"] == 1
    assert out["summary"]["scoredSec"] == pytest.approx(30.0, abs=0.1)


def test_windows_past_the_budget_are_skipped(fast_features) -> None:
    out = timeline.scan_track(tone(90.0), SR, StubXai(), total_sec=90.0, budget_sec=-1.0)
    assert {s["state"] for s in out["segments"]} == {"skipped"}
    assert out["summary"]["scoredCount"] == 0
    assert out["summary"]["peakIndex"] is None


def test_cancelling_skips_the_remaining_windows(fast_features) -> None:
    out = timeline.scan_track(tone(90.0), SR, StubXai(), total_sec=90.0, cancelled=lambda: True)
    assert {s["state"] for s in out["segments"]} == {"skipped"}


def test_a_long_track_is_flagged_as_truncated(fast_features) -> None:
    out = timeline.scan_track(tone(30.0), SR, StubXai(), total_sec=500.0)
    assert out["truncated"] is True
    assert out["totalSec"] == 500.0
    assert out["durationSec"] == pytest.approx(30.0, abs=0.05)


def test_a_failing_window_does_not_sink_the_scan(fast_features, monkeypatch) -> None:
    calls = {"n": 0}

    def flaky(seg, sr):
        calls["n"] += 1
        if calls["n"] == 2:
            raise RuntimeError("boom")
        return object()

    monkeypatch.setattr(timeline, "extract_features_array", flaky)
    monkeypatch.setattr(timeline, "SCAN_WORKERS", 1)
    out = timeline.scan_track(tone(90.0), SR, StubXai(), total_sec=90.0)
    assert [s["state"] for s in out["segments"]] == ["ok", "failed", "ok"]
    assert out["summary"]["scoredCount"] == 2


def test_progress_reports_every_window(fast_features) -> None:
    seen: List[tuple] = []
    timeline.scan_track(tone(90.0), SR, StubXai(), total_sec=90.0, on_progress=lambda d, n: seen.append((d, n)))
    assert sorted(seen) == [(1, 3), (2, 3), (3, 3)]


def test_real_audio_features_run_on_a_short_window() -> None:
    """The librosa path, with only the classifier stubbed."""
    rng = np.random.default_rng(3)
    y = (tone(4.0, 0.3) + 0.02 * rng.standard_normal(4 * SR)).astype(np.float32)
    out = timeline.scan_track(y, SR, StubXai([0.2]), total_sec=4.0)
    assert out["segments"][0]["state"] == "ok"
    assert out["segments"][0]["probability"] == 0.2


def test_window_beat_count_is_scaled_to_the_training_clip() -> None:
    """A 30 s window has about half the beats of a 60 s clip; the model expects 60 s."""
    seen = {}

    class Scaler:
        def transform(self, x):
            seen["x"] = x.copy()
            return x

    class Model:
        def predict_proba(self, x):
            return np.array([[0.4, 0.6]])

    svc = XAIInferenceService.__new__(XAIInferenceService)
    svc.available = True
    svc.threshold = 0.43
    svc.feature_cols = ["beat_count", "tempo_bpm"]
    svc.feature_stats = {}
    svc.shap_explainer = None
    svc.scaler = Scaler()
    svc.model = Model()

    features = AudioFeatures(**{f.name: 0 for f in dataclasses.fields(AudioFeatures)})
    features.beat_count = 25
    features.tempo_bpm = 120.0

    reading = svc.predict_window(features, None, window_sec=30.0)

    assert seen["x"][0, 0] == 50.0
    assert seen["x"][0, 1] == 120.0
    assert reading is not None
    assert reading.probability == pytest.approx(0.6)
    assert reading.is_ai is True
    assert reading.reasons == []