import numpy as np
import pytest

from fasr.data.waveform import Waveform


class TestWaveformBasic:
    def test_create(self, sine_wave, sample_rate):
        wf = Waveform(data=sine_wave, sample_rate=sample_rate)
        assert wf.data is not None
        assert wf.sample_rate == sample_rate
        assert len(wf) == len(sine_wave)

    def test_duration(self, waveform, sample_rate):
        assert abs(waveform.duration - 1.0) < 0.01
        assert abs(waveform.duration_s - 1.0) < 0.01
        assert abs(waveform.duration_ms - 1000.0) < 10

    def test_ndim(self, waveform):
        assert waveform.ndim == 1

    def test_len(self, waveform, sample_rate):
        assert len(waveform) == sample_rate


class TestWaveformSlice:
    def test_getitem_int(self, waveform):
        sliced = waveform[0:100]
        assert isinstance(sliced, Waveform)
        assert len(sliced) == 100
        assert sliced.sample_rate == waveform.sample_rate

    def test_getitem_slice(self, waveform, sample_rate):
        half = waveform[: sample_rate // 2]
        assert len(half) == sample_rate // 2

    def test_select_by_ms(self, waveform, sample_rate):
        selected = waveform.select_by_ms(0, 500)
        expected_len = int(500 * sample_rate / 1000)
        assert len(selected) == expected_len
        assert selected.sample_rate == waveform.sample_rate

    def test_select_by_ms_none_start(self, waveform, sample_rate):
        selected = waveform.select_by_ms(None, 500)
        expected_len = int(500 * sample_rate / 1000)
        assert len(selected) == expected_len

    def test_select_by_ms_none_end(self, waveform, sample_rate):
        selected = waveform.select_by_ms(200, None)
        start_idx = int(200 * sample_rate / 1000)
        expected_len = sample_rate - start_idx
        assert len(selected) == expected_len


class TestWaveformResample:
    def test_resample_down(self, waveform):
        resampled = waveform.resample(8000)
        assert resampled.sample_rate == 8000
        assert abs(resampled.duration - 1.0) < 0.05

    def test_resample_up(self, sample_rate):
        data = np.zeros(8000, dtype=np.float32)
        wf = Waveform(data=data, sample_rate=8000)
        resampled = wf.resample(sample_rate)
        assert resampled.sample_rate == sample_rate


class TestWaveformConcat:
    def test_concatenate(self, waveform, short_waveform):
        combined = Waveform.concatenate([waveform, short_waveform])
        assert len(combined) == len(waveform) + len(short_waveform)
        assert combined.sample_rate == waveform.sample_rate

    def test_append(self, sample_rate):
        wf1 = Waveform(data=np.ones(100, dtype=np.float32), sample_rate=sample_rate)
        wf2 = Waveform(data=np.zeros(50, dtype=np.float32), sample_rate=sample_rate)
        wf1.append(wf2)
        assert len(wf1) == 150

    def test_append_different_sr_raises(self):
        wf1 = Waveform(data=np.ones(100, dtype=np.float32), sample_rate=16000)
        wf2 = Waveform(data=np.ones(100, dtype=np.float32), sample_rate=8000)
        with pytest.raises(AssertionError):
            wf1.append(wf2)


class TestWaveformTransform:
    def test_flatten(self, sample_rate):
        data = np.ones((2, 100), dtype=np.float32)
        wf = Waveform(data=data, sample_rate=sample_rate)
        wf.flatten()
        assert wf.ndim == 1
        assert len(wf) == 200

    def test_as_1d_clears_feats(self, sample_rate):
        data = np.ones((2, 100), dtype=np.float32)
        feats = np.ones((10, 80), dtype=np.float32)
        wf = Waveform(data=data, sample_rate=sample_rate, feats=feats, scores=feats)
        wf.as_1d(clear=True)
        assert wf.ndim == 1
        assert wf.feats is None
        assert wf.scores is None

    def test_as_1d_keeps_feats(self, sample_rate):
        data = np.ones((2, 100), dtype=np.float32)
        feats = np.ones((10, 80), dtype=np.float32)
        wf = Waveform(data=data, sample_rate=sample_rate, feats=feats)
        wf.as_1d(clear=False)
        assert wf.ndim == 1
        assert wf.feats is not None
