import numpy as np
import pytest

from audio_wsola import wsola_time_stretch


SAMPLE_RATE = 48_000


def _voiced_fixture(seconds: float = 1.8) -> np.ndarray:
    samples = int(round(seconds * SAMPLE_RATE))
    time = np.arange(samples, dtype=np.float64) / SAMPLE_RATE
    phase = 2.0 * np.pi * (190.0 * time + 18.0 * np.square(time))
    envelope = 0.55 + 0.35 * np.sin(2.0 * np.pi * 2.3 * time)
    waveform = envelope * (
        0.30 * np.sin(phase)
        + 0.13 * np.sin(2.0 * phase + 0.2)
        + 0.07 * np.sin(3.0 * phase + 0.7)
    )
    return np.asarray(waveform, dtype=np.float32)


def _dominant_f0(waveform: np.ndarray) -> float:
    central = waveform[len(waveform) // 4 : 3 * len(waveform) // 4]
    central = np.asarray(central, dtype=np.float64) * np.hanning(central.size)
    spectrum = np.abs(np.fft.rfft(central))
    frequencies = np.fft.rfftfreq(central.size, 1.0 / SAMPLE_RATE)
    mask = (frequencies >= 150.0) & (frequencies <= 300.0)
    return float(frequencies[mask][np.argmax(spectrum[mask])])


@pytest.mark.parametrize("rate", (0.90, 0.95, 1.05))
def test_wsola_is_deterministic_finite_and_has_exact_length(rate):
    waveform = _voiced_fixture()

    first = wsola_time_stretch(
        waveform,
        rate=rate,
        sample_rate=SAMPLE_RATE,
    )
    second = wsola_time_stretch(
        waveform,
        rate=rate,
        sample_rate=SAMPLE_RATE,
    )

    assert first.dtype == np.float32
    assert first.ndim == 1
    assert first.flags.c_contiguous
    assert first.size == round(waveform.size / rate)
    assert np.isfinite(first).all()
    assert np.array_equal(first, second)


@pytest.mark.parametrize("rate", (0.90, 0.95, 1.05))
def test_wsola_preserves_voiced_pitch(rate):
    waveform = _voiced_fixture()
    stretched = wsola_time_stretch(
        waveform,
        rate=rate,
        sample_rate=SAMPLE_RATE,
    )

    source_f0 = _dominant_f0(waveform)
    stretched_f0 = _dominant_f0(stretched)
    assert abs(stretched_f0 - source_f0) / source_f0 < 0.015


def test_wsola_near_unity_is_bit_exact_copy():
    waveform = _voiced_fixture(0.5)

    stretched = wsola_time_stretch(
        waveform,
        rate=1.0005,
        sample_rate=SAMPLE_RATE,
    )

    assert stretched is not waveform
    assert np.array_equal(stretched, waveform)


def test_wsola_preserves_transient_count_and_endpoint_transients():
    rng = np.random.default_rng(31)
    waveform = np.asarray(
        0.002 * rng.standard_normal(2 * SAMPLE_RATE),
        dtype=np.float32,
    )
    transient_positions = (200, 12_000, 34_000, 60_000, waveform.size - 201)
    transient = np.asarray((0.08, 0.40, 0.80, 0.40, 0.08), dtype=np.float32)
    for position in transient_positions:
        waveform[position - 2 : position + 3] += transient

    stretched = wsola_time_stretch(
        waveform,
        rate=0.90,
        sample_rate=SAMPLE_RATE,
    )
    above_threshold = np.flatnonzero(np.abs(stretched) >= 0.30)
    groups = np.split(
        above_threshold,
        np.flatnonzero(np.diff(above_threshold) > 500) + 1,
    )

    assert len(groups) == len(transient_positions)
    assert int(groups[0][0]) < 500
    assert int(groups[-1][-1]) > stretched.size - 500
    assert all(float(np.max(np.abs(stretched[group]))) >= 0.50 for group in groups)


@pytest.mark.parametrize(
    ("audio", "rate", "sample_rate"),
    (
        (np.zeros(128, dtype=np.float32), 0.90, SAMPLE_RATE),
        (np.zeros((1024, 2), dtype=np.float32), 0.90, SAMPLE_RATE),
        (np.full(1024, np.nan, dtype=np.float32), 0.90, SAMPLE_RATE),
        (np.zeros(1024, dtype=np.float32), 0.0, SAMPLE_RATE),
        (np.zeros(1024, dtype=np.float32), 0.49, SAMPLE_RATE),
        (np.zeros(1024, dtype=np.float32), 2.01, SAMPLE_RATE),
        (np.zeros(1024, dtype=np.float32), 0.90, 0),
    ),
)
def test_wsola_rejects_invalid_inputs(audio, rate, sample_rate):
    with pytest.raises(ValueError):
        wsola_time_stretch(
            audio,
            rate=rate,
            sample_rate=sample_rate,
        )