import pytest
import torch

from wav2aug.gpu import (
    Wav2Aug,
    add_babble_noise,
    add_noise,
    chunk_swap,
    freq_drop,
    invert_polarity,
    rand_amp_clip,
    rand_amp_scale,
    speed_perturb,
    time_dropout,
)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


def _waveforms(
    batch: int = 3, time: int = 256, *, dtype: torch.dtype = torch.float32
) -> torch.Tensor:
    torch.manual_seed(0)
    return torch.randn(batch, time, device=DEVICE, dtype=dtype)


def test_rand_amp_clip_inplace_preserves_shape():
    waveforms = _waveforms()
    ptr = waveforms.data_ptr()
    out = rand_amp_clip(waveforms)
    assert out.data_ptr() == ptr
    assert out.shape == (waveforms.size(0), waveforms.size(1))
    assert torch.isfinite(out).all()


def test_rand_amp_scale_inplace_preserves_shape():
    waveforms = _waveforms()
    ptr = waveforms.data_ptr()
    out = rand_amp_scale(waveforms)
    assert out.data_ptr() == ptr
    assert out.shape == (waveforms.size(0), waveforms.size(1))
    assert torch.isfinite(out).all()


def test_chunk_swap_outputs_permutation():
    batch, time = 2, 400
    base = torch.arange(batch * time, device=DEVICE, dtype=torch.float32).view(
        batch, time
    )
    reference = base.clone()
    out = chunk_swap(base)
    assert out.shape == base.shape
    assert torch.allclose(
        torch.sort(out, dim=1).values, torch.sort(reference, dim=1).values
    )


def test_freq_drop_no_nan_and_inplace():
    waveforms = _waveforms()
    ptr = waveforms.data_ptr()
    out = freq_drop(waveforms)
    assert out.data_ptr() == ptr
    assert torch.isnan(out).logical_not().all()


def test_add_noise_with_mock_loader():
    """Test add_noise with a mock NoiseLoader."""
    from unittest.mock import MagicMock

    waveforms = torch.ones(2, 128, device=DEVICE, dtype=torch.float32)
    ptr = waveforms.data_ptr()

    # Create mock loader that returns zeros
    mock_loader = MagicMock()
    mock_loader.get_batch.return_value = torch.zeros(2, 128)

    out = add_noise(waveforms, mock_loader, snr_low=0.0, snr_high=0.0)
    assert out.data_ptr() == ptr
    assert torch.isfinite(out).all()
    mock_loader.get_batch.assert_called_once_with(2, 128)


def test_add_babble_noise_identity_for_singleton_batch():
    waveforms = torch.full((1, 64), 2.0, device=DEVICE, dtype=torch.float32)
    ptr = waveforms.data_ptr()
    out = add_babble_noise(waveforms, snr_low=0.0, snr_high=0.0)
    assert out.data_ptr() == ptr
    assert torch.allclose(out, torch.full_like(out, 2.0))


def test_invert_polarity_flips_when_prob_one():
    waveforms = _waveforms(batch=2, time=32)
    original = waveforms.clone()
    out = invert_polarity(waveforms, prob=1.0)
    assert torch.allclose(out, -original)


def test_speed_perturb_adjusts_length():
    waveforms = torch.linspace(
        0, 1, steps=200, device=DEVICE, dtype=torch.float32
    ).repeat(2, 1)
    out = speed_perturb(waveforms, 16000, speeds=(50,))
    # speed=50% → ratio=2.0 → 2x samples (slower)
    expected_len = int(200 * 2.0)
    assert out.shape == (2, expected_len)


def test_time_dropout_zeroes_segments():
    waveforms = torch.ones(2, 64, device=DEVICE, dtype=torch.float32)
    lengths = torch.ones(2, device=DEVICE, dtype=torch.float32)
    ptr = waveforms.data_ptr()
    out = time_dropout(
        waveforms,
        lengths=lengths,
        chunk_count_low=1,
        chunk_count_high=1,
        chunk_size_low=2,
        chunk_size_high=2,
    )
    assert out.data_ptr() == ptr
    zeros_per_row = (out == 0).sum(dim=1)
    assert torch.all(zeros_per_row >= 2)


def test_wav2aug_runs_with_stubbed_noise(monkeypatch):
    def _noop_add_noise(waveforms, loader, **kwargs):
        return waveforms

    monkeypatch.setattr("wav2aug.gpu.wav2aug.add_noise", _noop_add_noise)

    aug = Wav2Aug(sample_rate=16_000)
    waveforms = _waveforms(batch=3, time=256)
    lengths = torch.ones(3, device=DEVICE, dtype=torch.float32)
    out_wave, out_lengths = aug(waveforms, lengths=lengths)
    assert out_wave.shape[0] == waveforms.shape[0]
    assert out_lengths.data_ptr() == lengths.data_ptr()


def test_wav2aug_top_k_limits_ops(monkeypatch):
    """Test that top_k limits the number of augmentations used."""

    def _noop_add_noise(waveforms, loader, **kwargs):
        return waveforms

    monkeypatch.setattr("wav2aug.gpu.wav2aug.add_noise", _noop_add_noise)

    # top_k=3 should only include: noise, freq_drop, time_dropout
    aug = Wav2Aug(sample_rate=16_000, top_k=3)
    assert len(aug._base_ops) == 3

    # top_k=6 should include: noise, freq_drop, time_dropout, speed_perturb, amp_clip, chunk_swap
    aug = Wav2Aug(sample_rate=16_000, top_k=6)
    assert len(aug._base_ops) == 6

    # no pass should include all 9
    aug = Wav2Aug(sample_rate=16_000)
    assert len(aug._base_ops) == 9


def test_wav2aug_top_k_invalid_raises():
    """Test that invalid top_k values raise ValueError."""
    with pytest.raises(ValueError, match="top_k must be between 1 and 9"):
        Wav2Aug(sample_rate=16_000, top_k=0)

    with pytest.raises(ValueError, match="top_k must be between 1 and 9"):
        Wav2Aug(sample_rate=16_000, top_k=10)


def test_wav2aug_noise_dtype(monkeypatch):
    """Test that noise_dtype is passed to NoiseLoader."""

    def _noop_add_noise(waveforms, loader, **kwargs):
        return waveforms

    monkeypatch.setattr("wav2aug.gpu.wav2aug.add_noise", _noop_add_noise)

    # Default should be float32
    aug = Wav2Aug(sample_rate=16_000)
    assert aug.noise_dtype == torch.float32
    assert aug._noise_loader.storage_dtype == torch.float32

    # Custom dtype should be passed through
    aug = Wav2Aug(sample_rate=16_000, noise_dtype=torch.float16)
    assert aug.noise_dtype == torch.float16
    assert aug._noise_loader.storage_dtype == torch.float16
