import pytest
import numpy as np
from openmusic_analysis.worker.preprocessing.audio import select_temporal_windows, aggregate_global_embedding

def test_select_temporal_windows():
    # Test short duration
    windows = select_temporal_windows(5.0)
    assert len(windows) == 1
    assert windows[0] == (0, 5.0)

    # Test exact 10s duration
    windows = select_temporal_windows(10.0)
    assert len(windows) == 1
    assert windows[0] == (0, 10.0)

    # Test longer duration
    windows = select_temporal_windows(15.0)
    assert len(windows) == 2
    assert windows[0] == (0, 10.0)
    assert windows[1] == (5.0, 15.0)

    # Test max windows constraint (e.g. 24 chunks limit by config)
    windows = select_temporal_windows(300.0)
    assert len(windows) <= 24

def test_aggregate_global_embedding():
    # Two identical embeddings should average out to same and remain normalized
    embs = [
        np.array([1.0, 0.0, 0.0], dtype=np.float32),
        np.array([1.0, 0.0, 0.0], dtype=np.float32),
    ]
    global_emb = aggregate_global_embedding(embs)
    np.testing.assert_allclose(global_emb, [1.0, 0.0, 0.0])

    # Orthogonal vectors should normalize to 1/sqrt(2)
    embs = [
        np.array([1.0, 0.0], dtype=np.float32),
        np.array([0.0, 1.0], dtype=np.float32),
    ]
    global_emb = aggregate_global_embedding(embs)
    np.testing.assert_allclose(global_emb, [0.70710677, 0.70710677], rtol=1e-5)
    assert np.isclose(np.linalg.norm(global_emb), 1.0)