tests/test_audio_analyzers.py

from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest
from pydantic import ValidationError

from openmusic_analysis.analyzers.clap import (
    ClapGlobalAudioAnalyzer,
    ClapTemporalAudioAnalyzer,
    _trajectory_summary,
)
from openmusic_analysis.audio import AnalysisContext
from openmusic_analysis.domain import TemporalEmbeddingResult
from openmusic_analysis.settings import GlobalAudioConfig, TemporalAudioConfig

from .conftest import FakeAudioEncoder, FakeDecoder


def context(tmp_path: Path, waveform: np.ndarray) -> AnalysisContext:
    source = tmp_path / "audio.wav"
    source.write_bytes(b"valid")
    return AnalysisContext(source, FakeDecoder(waveform))


@pytest.mark.asyncio
async def test_global_is_deterministic_finite_normalized_and_multiwindow(tmp_path):
    waveform = np.linspace(-1, 1, 100, dtype=np.float32)
    encoder = FakeAudioEncoder()
    config = GlobalAudioConfig(
        sample_rate=10, window_seconds=2, target_windows=4, minimum_audio_seconds=1
    )
    analyzer = ClapGlobalAudioAnalyzer(encoder, config)
    first = await analyzer.analyze(context(tmp_path, waveform))
    second = await analyzer.analyze(context(tmp_path, waveform))
    first_vector = np.asarray(first.embedding)
    second_vector = np.asarray(second.embedding)
    assert first.dimension == 4
    assert np.isfinite(first_vector).all()
    assert np.linalg.norm(first_vector) == pytest.approx(1.0, abs=1e-6)
    assert first_vector == pytest.approx(second_vector, abs=1e-7)
    assert first.analysis["windows_used"] == 4
    assert len(encoder.calls[0]) == 4


@pytest.mark.asyncio
async def test_global_short_track_uses_one_padded_window(tmp_path):
    encoder = FakeAudioEncoder()
    config = GlobalAudioConfig(
        sample_rate=10, window_seconds=2, target_windows=4, minimum_audio_seconds=0.2
    )
    result = await ClapGlobalAudioAnalyzer(encoder, config).analyze(
        context(tmp_path, np.arange(5, dtype=np.float32))
    )
    assert result.analysis["windows_used"] == 1
    assert encoder.calls[0][0].size == 20


@pytest.mark.asyncio
async def test_temporal_segments_are_ordered_and_deterministic(tmp_path):
    waveform = np.linspace(-1, 1, 95, dtype=np.float32)
    config = TemporalAudioConfig(
        sample_rate=10,
        window_seconds=2,
        hop_seconds=2,
        max_segments=20,
        minimum_audio_seconds=1,
    )
    analyzer = ClapTemporalAudioAnalyzer(FakeAudioEncoder(), config)
    first = await analyzer.analyze(context(tmp_path, waveform))
    second = await analyzer.analyze(context(tmp_path, waveform))
    starts = [segment.start_ms for segment in first.segments]
    assert starts == sorted(starts)
    assert starts == [segment.start_ms for segment in second.segments]
    assert first.model_dump() == second.model_dump()
    assert all(segment.start_ms < segment.end_ms for segment in first.segments)
    assert first.summary.number_of_segments == len(first.segments)
    assert np.isfinite(first.summary.mean_adjacent_distance)
    assert np.isfinite(first.summary.max_adjacent_distance)
    assert np.isfinite(first.summary.trajectory_variance)


@pytest.mark.asyncio
async def test_temporal_short_track_has_one_segment(tmp_path):
    config = TemporalAudioConfig(
        sample_rate=10,
        window_seconds=2,
        hop_seconds=1,
        max_segments=4,
        minimum_audio_seconds=0.2,
    )
    result = await ClapTemporalAudioAnalyzer(FakeAudioEncoder(), config).analyze(
        context(tmp_path, np.arange(5, dtype=np.float32))
    )
    assert len(result.segments) == 1
    assert result.segments[0].start_ms == 0
    assert result.segments[0].end_ms == 500
    assert result.summary.largest_transition_index is None
    assert result.summary.number_of_segments == len(result.segments)


@pytest.mark.asyncio
async def test_temporal_long_track_respects_max_segments_and_tail(tmp_path):
    config = TemporalAudioConfig(
        sample_rate=10,
        window_seconds=2,
        hop_seconds=1,
        max_segments=5,
        minimum_audio_seconds=1,
    )
    result = await ClapTemporalAudioAnalyzer(FakeAudioEncoder(), config).analyze(
        context(tmp_path, np.arange(300, dtype=np.float32))
    )
    assert len(result.segments) == 5
    assert result.segments[0].start_ms == 0
    assert result.segments[-1].end_ms == 30_000
    assert result.summary.number_of_segments == 5
    assert np.isfinite(result.summary.trajectory_variance)
    assert 0 <= result.summary.largest_transition_index <= len(result.segments) - 2


def test_largest_transition_index_uses_adjacent_transition_indices():
    fifteen_segments = np.asarray(
        [[1.0, 0.0]] * 14 + [[-1.0, 0.0]], dtype=np.float32
    )
    fifteen_summary = _trajectory_summary(fifteen_segments)
    assert fifteen_summary.number_of_segments == 15
    assert fifteen_summary.largest_transition_index == 13
    assert fifteen_summary.largest_transition_index <= 15 - 2

    two_summary = _trajectory_summary(
        np.asarray([[1.0, 0.0], [-1.0, 0.0]], dtype=np.float32)
    )
    assert two_summary.number_of_segments == 2
    assert two_summary.largest_transition_index == 0

    one_summary = _trajectory_summary(np.asarray([[1.0, 0.0]], dtype=np.float32))
    assert one_summary.number_of_segments == 1
    assert one_summary.largest_transition_index is None

    for summary in (fifteen_summary, two_summary, one_summary):
        assert np.isfinite(summary.mean_adjacent_distance)
        assert np.isfinite(summary.max_adjacent_distance)
        assert np.isfinite(summary.trajectory_variance)


@pytest.mark.asyncio
async def test_temporal_response_schema_enforces_count_and_index_range(tmp_path):
    config = TemporalAudioConfig(
        sample_rate=10,
        window_seconds=1,
        hop_seconds=1,
        max_segments=15,
        minimum_audio_seconds=1,
    )
    result = await ClapTemporalAudioAnalyzer(FakeAudioEncoder(), config).analyze(
        context(tmp_path, np.arange(150, dtype=np.float32))
    )
    assert len(result.segments) == 15
    assert result.summary.number_of_segments == len(result.segments)
    assert 0 <= result.summary.largest_transition_index <= 13

    invalid_index = result.model_dump()
    invalid_index["summary"]["largest_transition_index"] = 14
    with pytest.raises(ValidationError):
        TemporalEmbeddingResult.model_validate(invalid_index)

    invalid_count = result.model_dump()
    invalid_count["summary"]["number_of_segments"] = 14
    with pytest.raises(ValidationError):
        TemporalEmbeddingResult.model_validate(invalid_count)

    non_finite_summary = result.model_dump()
    non_finite_summary["summary"]["mean_adjacent_distance"] = float("nan")
    with pytest.raises(ValidationError):
        TemporalEmbeddingResult.model_validate(non_finite_summary)


@pytest.mark.asyncio
async def test_request_context_decodes_only_once_for_two_analyzers(tmp_path):
    source = tmp_path / "audio.wav"
    source.write_bytes(b"valid")
    decoder = FakeDecoder(np.arange(100, dtype=np.float32))
    analysis_context = AnalysisContext(source, decoder)
    encoder = FakeAudioEncoder()
    await ClapGlobalAudioAnalyzer(
        encoder,
        GlobalAudioConfig(sample_rate=10, window_seconds=2, minimum_audio_seconds=1),
    ).analyze(analysis_context)
    await ClapTemporalAudioAnalyzer(
        encoder,
        TemporalAudioConfig(sample_rate=10, window_seconds=2, minimum_audio_seconds=1),
    ).analyze(analysis_context)
    assert decoder.calls == 1