import numpy as np

from pipu_rec.analysis.features import EMBEDDING_DIM, SAMPLE_RATE, analyze, estimate_tempo, segments


def metronome(bpm: float, seconds: float, freq: float = 1000.0) -> np.ndarray:
    t = np.arange(int(seconds * SAMPLE_RATE)) / SAMPLE_RATE
    period = 60.0 / bpm
    phase = t % period
    click = np.sin(2 * np.pi * freq * t) * np.exp(-phase * 60)
    return (0.5 * click).astype(np.float32)


def test_tempo_of_a_metronome() -> None:
    for bpm in (90, 128, 170):
        f = analyze(metronome(bpm, 40))
        assert f.bpm is not None
        # двойной / половинный темп допустим для щелчков, но для этих значений приор держит исходный
        assert abs(f.bpm - bpm) < 2.5, (bpm, f.bpm)


def test_no_rhythm_in_a_steady_tone() -> None:
    t = np.arange(SAMPLE_RATE * 20) / SAMPLE_RATE
    tone = (0.3 * np.sin(2 * np.pi * 440 * t)).astype(np.float32)
    f = analyze(tone)
    assert f.bpm is None or f.bpm_confidence < 0.3
    assert f.key == "A"
    assert f.embedding.shape == (EMBEDDING_DIM,)
    assert abs(float(np.linalg.norm(f.embedding)) - 1) < 1e-4


def test_louder_and_busier_means_more_energy() -> None:
    quiet = (metronome(80, 30) * 0.05).astype(np.float32)
    loud = metronome(160, 30)
    noise = np.random.default_rng(1).normal(0, 0.2, len(loud)).astype(np.float32)
    assert analyze(loud + noise).energy > analyze(quiet).energy


def test_similar_sounds_have_closer_embeddings() -> None:
    a = analyze(metronome(120, 30, 1000))
    b = analyze(metronome(124, 30, 1100))
    rng = np.random.default_rng(2)
    c = analyze(rng.normal(0, 0.3, SAMPLE_RATE * 30).astype(np.float32))
    assert float(a.embedding @ b.embedding) > float(a.embedding @ c.embedding)


def test_segments_cover_long_tracks() -> None:
    long = np.zeros(SAMPLE_RATE * 300, dtype=np.float32)
    parts = segments(long)
    assert len(parts) == 3
    assert all(len(p) == SAMPLE_RATE * 30 for p in parts)
    assert len(segments(np.zeros(SAMPLE_RATE * 60, dtype=np.float32))) == 1


def test_estimate_tempo_needs_signal() -> None:
    assert estimate_tempo(np.zeros(100, dtype=np.float32)) == (None, 0.0)