"""
Tests for SIMD-optimized polyphase resampling.

Ensures all SIMD implementations produce identical results and handle edge cases.
"""

import numpy as np
import pytest

from sgnl_cpu_interp import (
    correlate,
    get_simd_info,
    set_implementation,
    upsample,
    upsample_transposed,
)


@pytest.fixture
def test_data():
    """Generate deterministic test data."""
    np.random.seed(42)
    return np.random.randn(256, 32).astype(np.float32)


@pytest.fixture
def test_data_transposed():
    """Generate deterministic transposed test data."""
    np.random.seed(42)
    return np.random.randn(32, 256).astype(np.float32)


class TestSIMDInfo:
    """Test SIMD introspection functions."""

    def test_get_simd_info_returns_dict(self):
        info = get_simd_info()
        assert isinstance(info, dict)
        assert "implementation" in info
        assert "available" in info
        assert "cpu_features" in info

    def test_implementation_is_string(self):
        info = get_simd_info()
        assert isinstance(info["implementation"], str)
        assert len(info["implementation"]) > 0

    def test_available_is_list(self):
        info = get_simd_info()
        assert isinstance(info["available"], list)
        assert len(info["available"]) >= 1  # At least Scalar

    def test_scalar_always_available(self):
        info = get_simd_info()
        assert "Scalar" in info["available"]


class TestImplementationSelection:
    """Test manual implementation selection."""

    def test_select_scalar(self):
        set_implementation("Scalar")
        assert get_simd_info()["implementation"] == "Scalar"

    def test_select_invalid_raises(self):
        with pytest.raises(ValueError):
            set_implementation("NonExistentImpl")

    def test_can_switch_between_impls(self):
        available = get_simd_info()["available"]
        for impl in available:
            set_implementation(impl)
            assert get_simd_info()["implementation"] == impl


class TestUpsampleBasic:
    """Basic functionality tests for upsample."""

    def test_output_shape_factor2(self, test_data):
        result = upsample(test_data, factor=2, half_length=8)
        expected_samples = (256 - 17 + 1) * 2  # (n - kernel_len + 1) * factor
        assert result.shape == (expected_samples, 32)

    def test_output_shape_factor4(self, test_data):
        result = upsample(test_data, factor=4, half_length=8)
        expected_samples = (256 - 17 + 1) * 4
        assert result.shape == (expected_samples, 32)

    def test_output_dtype(self, test_data):
        result = upsample(test_data, factor=2)
        assert result.dtype == np.float32

    def test_1d_input(self):
        data = np.random.randn(256).astype(np.float32)
        result = upsample(data, factor=2)
        assert result.ndim == 1
        assert result.shape[0] == (256 - 17 + 1) * 2


class TestUpsampleTransposedBasic:
    """Basic functionality tests for upsample_transposed."""

    def test_output_shape_factor2(self, test_data_transposed):
        result = upsample_transposed(test_data_transposed, factor=2, half_length=8)
        expected_samples = (256 - 17 + 1) * 2
        assert result.shape == (32, expected_samples)

    def test_output_dtype(self, test_data_transposed):
        result = upsample_transposed(test_data_transposed, factor=2)
        assert result.dtype == np.float32


class TestSIMDConsistency:
    """Ensure all SIMD implementations produce identical results."""

    def test_all_impls_match_scalar_standard(self, test_data):
        """All implementations must match Scalar for standard layout."""
        available = get_simd_info()["available"]

        # Get scalar reference
        set_implementation("Scalar")
        reference = upsample(test_data, factor=2, half_length=8)

        for impl in available:
            if impl == "Scalar":
                continue
            set_implementation(impl)
            result = upsample(test_data, factor=2, half_length=8)
            np.testing.assert_allclose(
                result,
                reference,
                rtol=1e-5,
                atol=1e-6,
                err_msg=f"{impl} differs from Scalar reference",
            )

    def test_all_impls_match_scalar_transposed(self, test_data_transposed):
        """All implementations must match Scalar for transposed layout."""
        available = get_simd_info()["available"]

        # Get scalar reference
        set_implementation("Scalar")
        reference = upsample_transposed(test_data_transposed, factor=2, half_length=8)

        for impl in available:
            if impl == "Scalar":
                continue
            set_implementation(impl)
            result = upsample_transposed(test_data_transposed, factor=2, half_length=8)
            np.testing.assert_allclose(
                result,
                reference,
                rtol=1e-5,
                atol=1e-6,
                err_msg=f"{impl} differs from Scalar reference",
            )


class TestEdgeCases:
    """Test edge cases and boundary conditions."""

    @pytest.mark.parametrize("channels", [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33])
    def test_various_channel_counts(self, channels):
        """Test non-aligned channel counts."""
        data = np.random.randn(128, channels).astype(np.float32)
        result = upsample(data, factor=2, half_length=8)
        expected_samples = (128 - 17 + 1) * 2
        assert result.shape == (expected_samples, channels)

    @pytest.mark.parametrize("half_length", [1, 2, 4, 8, 16])
    def test_various_kernel_lengths(self, half_length):
        """Test different kernel half-lengths."""
        data = np.random.randn(256, 16).astype(np.float32)
        kernel_len = 2 * half_length + 1
        result = upsample(data, factor=2, half_length=half_length)
        expected_samples = (256 - kernel_len + 1) * 2
        assert result.shape == (expected_samples, 16)

    @pytest.mark.parametrize("factor", [2, 3, 4, 5, 8])
    def test_various_factors(self, factor):
        """Test different upsampling factors."""
        data = np.random.randn(128, 16).astype(np.float32)
        result = upsample(data, factor=factor, half_length=8)
        expected_samples = (128 - 17 + 1) * factor
        assert result.shape == (expected_samples, 16)

    def test_minimum_input_size(self):
        """Test with minimum valid input size."""
        half_length = 8
        kernel_len = 2 * half_length + 1
        data = np.random.randn(kernel_len, 4).astype(np.float32)
        result = upsample(data, factor=2, half_length=half_length)
        assert result.shape == (2, 4)  # (kernel_len - kernel_len + 1) * 2 = 2

    def test_input_too_small_raises(self):
        """Test that too-small input raises ValueError."""
        data = np.random.randn(10, 4).astype(np.float32)  # 10 < 17 (kernel_len)
        with pytest.raises(ValueError, match="at least"):
            upsample(data, factor=2, half_length=8)


class TestInputValidation:
    """Test input validation."""

    def test_factor_less_than_2_raises(self, test_data):
        with pytest.raises(ValueError, match="factor"):
            upsample(test_data, factor=1)

    def test_half_length_less_than_1_raises(self, test_data):
        with pytest.raises(ValueError, match="half_length"):
            upsample(test_data, factor=2, half_length=0)

    def test_3d_input_raises(self):
        data = np.random.randn(10, 10, 10).astype(np.float32)
        with pytest.raises(ValueError, match="1D or 2D"):
            upsample(data, factor=2)

    def test_non_contiguous_works(self, test_data):
        """Non-contiguous input should be handled (copied internally)."""
        non_contig = test_data[::2, :]  # Non-contiguous view
        assert not non_contig.flags["C_CONTIGUOUS"]
        result = upsample(non_contig, factor=2)
        assert result.shape[1] == test_data.shape[1]


class TestDtypeConversion:
    """Test automatic dtype conversion."""

    def test_float64_converted(self):
        data = np.random.randn(128, 16).astype(np.float64)
        result = upsample(data, factor=2)
        assert result.dtype == np.float32

    def test_int32_converted(self):
        data = np.random.randint(-100, 100, (128, 16)).astype(np.int32)
        result = upsample(data, factor=2)
        assert result.dtype == np.float32


class TestUpsampleTransposedValidation:
    """Test input validation for upsample_transposed."""

    def test_factor_less_than_2_raises(self, test_data_transposed):
        with pytest.raises(ValueError, match="factor"):
            upsample_transposed(test_data_transposed, factor=1)

    def test_half_length_less_than_1_raises(self, test_data_transposed):
        with pytest.raises(ValueError, match="half_length"):
            upsample_transposed(test_data_transposed, factor=2, half_length=0)

    def test_non_2d_input_raises(self):
        data = np.random.randn(100).astype(np.float32)
        with pytest.raises(ValueError, match="2D"):
            upsample_transposed(data, factor=2)

    def test_input_too_small_raises(self):
        data = np.random.randn(4, 10).astype(np.float32)
        with pytest.raises(ValueError, match="at least"):
            upsample_transposed(data, factor=2, half_length=8)

    def test_float64_converted(self):
        data = np.random.randn(4, 128).astype(np.float64)
        result = upsample_transposed(data, factor=2)
        assert result.dtype == np.float32

    def test_non_contiguous_works(self):
        data = np.random.randn(8, 256).astype(np.float32)
        non_contig = data[:, ::2]
        assert not non_contig.flags["C_CONTIGUOUS"]
        result = upsample_transposed(non_contig, factor=2)
        assert result.dtype == np.float32


class TestCorrelateValidation:
    """Test input validation for correlate."""

    def test_filters_wrong_dimensions_raises(self):
        data = np.random.randn(4, 100).astype(np.float32)
        filters = np.random.randn(64).astype(np.float32)
        with pytest.raises(ValueError, match="2D"):
            correlate(data, filters)

    def test_groups_less_than_1_raises(self):
        data = np.random.randn(4, 100).astype(np.float32)
        filters = np.random.randn(4, 16).astype(np.float32)
        with pytest.raises(ValueError, match="groups"):
            correlate(data, filters, groups=0)

    def test_channels_not_divisible_by_groups_raises(self):
        data = np.random.randn(5, 100).astype(np.float32)
        filters = np.random.randn(6, 16).astype(np.float32)
        with pytest.raises(ValueError, match="n_channels"):
            correlate(data, filters, groups=3)

    def test_float64_data_converted(self):
        data = np.random.randn(4, 100).astype(np.float64)
        filters = np.random.randn(4, 16).astype(np.float32)
        result = correlate(data, filters)
        assert result.dtype == np.float32

    def test_float64_filters_converted(self):
        data = np.random.randn(4, 100).astype(np.float32)
        filters = np.random.randn(4, 16).astype(np.float64)
        result = correlate(data, filters)
        assert result.dtype == np.float32

    def test_non_contiguous_data_works(self):
        data = np.random.randn(4, 200).astype(np.float32)[:, ::2]
        assert not data.flags["C_CONTIGUOUS"]
        filters = np.random.randn(4, 16).astype(np.float32)
        result = correlate(data, filters)
        assert result.dtype == np.float32

    def test_non_contiguous_filters_works(self):
        data = np.random.randn(4, 100).astype(np.float32)
        filters = np.random.randn(4, 32).astype(np.float32)[:, ::2]
        assert not filters.flags["C_CONTIGUOUS"]
        result = correlate(data, filters)
        assert result.dtype == np.float32


class TestCorrectnessStandard:
    """Test that upsample produces mathematically correct results."""

    def test_dc_signal_stays_constant(self):
        """A constant (DC) signal should remain constant after upsampling."""
        dc_value = 3.14159
        n_samples = 100
        data = np.full((n_samples, 1), dc_value, dtype=np.float32)

        result = upsample(data, factor=2, half_length=8)

        # After edge effects, the middle portion should be constant
        # Skip edges where the kernel doesn't have full support
        # Note: polyphase filters have small ripple, so use reasonable tolerance
        middle = result[20:-20, 0]
        np.testing.assert_allclose(
            middle,
            dc_value,
            rtol=1e-3,
            atol=1e-3,
            err_msg="DC signal not preserved after upsampling",
        )

    def test_sine_wave_frequency_preserved(self):
        """Upsampled sine wave should have correct frequency content."""
        fs_in = 100  # Input sample rate
        freq = 10  # 10 Hz sine wave
        duration = 1.0
        factor = 4

        t = np.arange(0, duration, 1 / fs_in)
        signal = np.sin(2 * np.pi * freq * t).astype(np.float32)

        result = upsample(signal[:, np.newaxis], factor=factor, half_length=8).squeeze()

        # The upsampled signal should still be a sine wave at the same frequency
        # Check using zero crossings (should be ~2*freq per second)
        fs_out = fs_in * factor
        # Skip edges
        result_middle = result[100:-100]
        zero_crossings = np.where(np.diff(np.sign(result_middle)))[0]

        # Expected ~2*freq*duration zero crossings (rising and falling)
        # Allow some tolerance for edge effects
        expected_crossings = 2 * freq * (len(result_middle) / fs_out)
        actual_crossings = len(zero_crossings)

        assert (
            abs(actual_crossings - expected_crossings) < 5
        ), f"Expected ~{expected_crossings:.0f} zero crossings, got {actual_crossings}"

    def test_impulse_response_symmetry(self):
        """Impulse response should be symmetric (linear phase filter)."""
        n_samples = 100
        half_length = 8
        factor = 2

        # Create impulse in the middle
        data = np.zeros((n_samples, 1), dtype=np.float32)
        data[n_samples // 2, 0] = 1.0

        result = upsample(data, factor=factor, half_length=half_length).squeeze()

        # Find the peak
        peak_idx = np.argmax(np.abs(result))

        # Check symmetry around peak (within the non-zero region)
        # The impulse response should be symmetric
        window = half_length * factor
        left = result[max(0, peak_idx - window) : peak_idx]
        right = result[peak_idx + 1 : min(len(result), peak_idx + window + 1)]

        # Reverse right side to compare with left
        min_len = min(len(left), len(right))
        if min_len > 2:
            np.testing.assert_allclose(
                left[-min_len:],
                right[:min_len][::-1],
                rtol=1e-4,
                atol=1e-5,
                err_msg="Impulse response is not symmetric",
            )

    def test_linear_ramp_interpolation(self):
        """Linear ramp should be interpolated correctly."""
        n_samples = 50
        factor = 2

        # Create linear ramp
        data = np.linspace(0, 10, n_samples, dtype=np.float32)[:, np.newaxis]

        result = upsample(data, factor=factor, half_length=4)

        # The middle portion should still be approximately linear
        # (edges will have artifacts due to kernel)
        middle_start = 20
        middle_end = -20
        result_middle = result[middle_start:middle_end, 0]

        # Fit a line and check residuals
        x = np.arange(len(result_middle))
        coeffs = np.polyfit(x, result_middle, 1)
        fitted = np.polyval(coeffs, x)
        residuals = np.abs(result_middle - fitted)

        # Residuals should be small for a linear signal
        assert (
            np.max(residuals) < 0.1
        ), f"Linear ramp not interpolated correctly, max residual: {np.max(residuals)}"

    @pytest.mark.parametrize(
        "impl",
        ["Scalar", "NEON"] if "NEON" in get_simd_info()["available"] else ["Scalar"],
    )
    def test_nyquist_no_aliasing(self, impl):
        """Signal below Nyquist should not alias after upsampling."""
        set_implementation(impl)

        fs_in = 1000
        freq = 100  # Well below Nyquist (500 Hz)
        factor = 4
        duration = 0.5  # Longer duration for more samples

        t = np.arange(0, duration, 1 / fs_in)
        signal = np.sin(2 * np.pi * freq * t).astype(np.float32)

        result = upsample(
            signal[:, np.newaxis], factor=factor, half_length=16
        ).squeeze()

        # Check that the output amplitude is close to input amplitude
        # (should be ~1.0 for a sine wave)
        # Skip edges - use percentage-based skip to handle various output sizes
        edge_skip = max(100, len(result) // 10)
        result_middle = result[edge_skip:-edge_skip]

        assert len(result_middle) > 0, "Result too short for amplitude check"

        amplitude = (np.max(result_middle) - np.min(result_middle)) / 2

        # Allow for small ripple in polyphase filter
        assert (
            0.90 < amplitude < 1.10
        ), f"Amplitude distortion: expected ~1.0, got {amplitude}"


class TestCorrectnessTransposed:
    """Test that upsample_transposed produces mathematically correct results."""

    def test_dc_signal_stays_constant(self):
        """A constant (DC) signal should remain constant after upsampling."""
        dc_value = 2.71828
        n_channels = 4
        n_samples = 100
        data = np.full((n_channels, n_samples), dc_value, dtype=np.float32)

        result = upsample_transposed(data, factor=2, half_length=8)

        # Check middle portion of each channel
        # Note: polyphase filters have small ripple, so use reasonable tolerance
        for ch in range(n_channels):
            middle = result[ch, 20:-20]
            np.testing.assert_allclose(
                middle,
                dc_value,
                rtol=1e-3,
                atol=1e-3,
                err_msg=f"DC signal not preserved in channel {ch}",
            )

    def test_matches_standard_layout(self):
        """Transposed result should match standard layout (transposed)."""
        np.random.seed(123)
        n_channels = 8
        n_samples = 128

        # Create data in both layouts
        data_standard = np.random.randn(n_samples, n_channels).astype(np.float32)
        data_transposed = data_standard.T.copy()

        result_standard = upsample(data_standard, factor=2, half_length=8)
        result_transposed = upsample_transposed(
            data_transposed, factor=2, half_length=8
        )

        # Results should match when transposed
        np.testing.assert_allclose(
            result_standard.T,
            result_transposed,
            rtol=1e-5,
            atol=1e-6,
            err_msg="Transposed layout doesn't match standard layout",
        )

    @pytest.mark.parametrize("factor", [2, 3, 4])
    def test_multichannel_independence(self, factor):
        """Each channel should be processed independently."""
        n_channels = 4
        n_samples = 100

        # Create different signals per channel
        data = np.zeros((n_channels, n_samples), dtype=np.float32)
        for ch in range(n_channels):
            data[ch, :] = ch + 1  # Different DC values

        result = upsample_transposed(data, factor=factor, half_length=8)

        # Check each channel has its own DC value
        # Note: polyphase filters have small ripple, so use reasonable tolerance
        for ch in range(n_channels):
            middle = result[ch, 20:-20]
            expected = ch + 1
            np.testing.assert_allclose(
                middle,
                expected,
                rtol=1e-3,
                atol=1e-3,
                err_msg=f"Channel {ch} interfered with other channels",
            )


class TestCorrelateBasic:
    """Basic functionality tests for correlate."""

    def test_output_shape(self):
        """Output shape should be (channels, samples - filter_len + 1)."""
        np.random.seed(42)
        channels, samples, filter_len = 32, 1000, 64
        data = np.random.randn(channels, samples).astype(np.float32)
        filters = np.random.randn(channels, filter_len).astype(np.float32)

        result = correlate(data, filters)

        expected_shape = (channels, samples - filter_len + 1)
        assert result.shape == expected_shape

    def test_output_dtype(self):
        """Output should be float32."""
        channels, samples, filter_len = 8, 100, 16
        data = np.random.randn(channels, samples).astype(np.float32)
        filters = np.random.randn(channels, filter_len).astype(np.float32)

        result = correlate(data, filters)

        assert result.dtype == np.float32

    def test_groups_must_divide_channels(self):
        """Groups must evenly divide both n_channels and n_filters."""
        data = np.random.randn(7, 100).astype(np.float32)
        filters = np.random.randn(20, 16).astype(np.float32)

        # Default groups=n_channels=7, but 20 filters not divisible by 7
        with pytest.raises(ValueError):
            correlate(data, filters)

    def test_groups_1_different_channels(self):
        """groups=1 should work with different input/output channels."""
        data = np.random.randn(1, 100).astype(np.float32)
        filters = np.random.randn(80, 16).astype(np.float32)

        result = correlate(data, filters, groups=1)
        assert result.shape == (80, 85)

    def test_data_too_short(self):
        """Data shorter than filter should raise error."""
        data = np.random.randn(8, 10).astype(np.float32)
        filters = np.random.randn(8, 20).astype(np.float32)

        with pytest.raises(ValueError):
            correlate(data, filters)

    def test_wrong_dimensions(self):
        """Non-2D inputs should raise error."""
        data_1d = np.random.randn(100).astype(np.float32)
        filters = np.random.randn(8, 16).astype(np.float32)

        with pytest.raises(ValueError):
            correlate(data_1d, filters)


class TestCorrelateCorrectness:
    """Test that correlate produces mathematically correct results."""

    def test_unit_impulse_filter(self):
        """Correlation with [1, 0, 0, ...] should return the signal shifted."""
        channels, samples = 4, 100
        filter_len = 8

        np.random.seed(42)
        data = np.random.randn(channels, samples).astype(np.float32)

        # Create impulse filter [1, 0, 0, ...]
        filters = np.zeros((channels, filter_len), dtype=np.float32)
        filters[:, 0] = 1.0

        result = correlate(data, filters)

        # Result should be data[:, :output_len] (first output_len samples)
        expected = data[:, : result.shape[1]]
        np.testing.assert_allclose(result, expected, rtol=1e-5, atol=1e-6)

    def test_constant_filter(self):
        """Correlation with constant filter should give moving sum."""
        channels, samples = 2, 50
        filter_len = 5

        # Create simple data with predictable values
        data = np.ones((channels, samples), dtype=np.float32)
        data[0, :] = 1.0
        data[1, :] = 2.0

        # Constant filter of all 1s
        filters = np.ones((channels, filter_len), dtype=np.float32)

        result = correlate(data, filters)

        # Result should be filter_len times the input value
        expected_ch0 = filter_len * 1.0  # = 5.0
        expected_ch1 = filter_len * 2.0  # = 10.0

        np.testing.assert_allclose(result[0, :], expected_ch0, rtol=1e-5)
        np.testing.assert_allclose(result[1, :], expected_ch1, rtol=1e-5)

    def test_matches_numpy_correlate(self):
        """Should match numpy's correlate for each channel."""
        channels = 4
        samples = 100
        filter_len = 16

        np.random.seed(42)
        data = np.random.randn(channels, samples).astype(np.float32)
        filters = np.random.randn(channels, filter_len).astype(np.float32)

        result = correlate(data, filters)

        # Compare with numpy's correlate (mode='valid')
        for ch in range(channels):
            expected = np.correlate(data[ch], filters[ch], mode="valid")
            np.testing.assert_allclose(
                result[ch],
                expected,
                rtol=1e-4,
                atol=1e-5,
                err_msg=f"Channel {ch} doesn't match numpy.correlate",
            )

    @pytest.mark.parametrize(
        "impl",
        ["Scalar", "NEON"] if "NEON" in get_simd_info()["available"] else ["Scalar"],
    )
    def test_all_implementations_match(self, impl):
        """All SIMD implementations should produce identical results."""
        np.random.seed(42)
        channels, samples, filter_len = 16, 500, 32
        data = np.random.randn(channels, samples).astype(np.float32)
        filters = np.random.randn(channels, filter_len).astype(np.float32)

        # Get reference from Scalar
        set_implementation("Scalar")
        reference = correlate(data, filters)

        # Test this implementation
        set_implementation(impl)
        result = correlate(data, filters)

        np.testing.assert_allclose(
            result,
            reference,
            rtol=1e-5,
            atol=1e-6,
            err_msg=f"{impl} doesn't match Scalar",
        )

    def test_channels_independent(self):
        """Each channel should be processed independently."""
        channels = 4
        samples = 100
        filter_len = 16

        # Create distinct signals and filters per channel
        np.random.seed(42)
        data = np.random.randn(channels, samples).astype(np.float32)
        filters = np.random.randn(channels, filter_len).astype(np.float32)

        # Full correlate
        result = correlate(data, filters)

        # Process each channel separately and compare
        for ch in range(channels):
            single_data = data[ch : ch + 1, :]
            single_filter = filters[ch : ch + 1, :]
            single_result = correlate(single_data, single_filter)

            np.testing.assert_allclose(
                result[ch],
                single_result[0],
                rtol=1e-4,
                err_msg=f"Channel {ch} interference detected",
            )


class TestCorrelateGroups:
    """Test correlate with groups parameter matching PyTorch conv1d."""

    def test_groups_1_matches_pytorch(self):
        """groups=1 should match PyTorch conv1d with groups=1."""
        try:
            import torch
            import torch.nn.functional as F
        except ImportError:
            pytest.skip("PyTorch not available")

        np.random.seed(42)
        data = np.random.randn(1, 1000).astype(np.float32)
        filters = np.random.randn(80, 64).astype(np.float32)

        # PyTorch reference
        data_torch = torch.from_numpy(data).unsqueeze(0)  # (1, 1, 1000)
        filters_torch = torch.from_numpy(filters).unsqueeze(1)  # (80, 1, 64)
        torch_out = F.conv1d(data_torch, filters_torch, groups=1).squeeze(0).numpy()

        # Our implementation
        our_out = correlate(data, filters, groups=1)

        assert our_out.shape == torch_out.shape
        np.testing.assert_allclose(our_out, torch_out, rtol=1e-4, atol=1e-5)

    def test_groups_n_matches_pytorch(self):
        """groups=n_channels should match PyTorch conv1d with groups=n."""
        try:
            import torch
            import torch.nn.functional as F
        except ImportError:
            pytest.skip("PyTorch not available")

        np.random.seed(42)
        n_channels = 80
        data = np.random.randn(n_channels, 1000).astype(np.float32)
        filters = np.random.randn(n_channels, 64).astype(np.float32)

        # PyTorch reference
        data_torch = torch.from_numpy(data).unsqueeze(0)
        filters_torch = torch.from_numpy(filters).unsqueeze(1)
        torch_out = (
            F.conv1d(data_torch, filters_torch, groups=n_channels).squeeze(0).numpy()
        )

        # Our implementation
        our_out = correlate(data, filters, groups=n_channels)

        assert our_out.shape == torch_out.shape
        np.testing.assert_allclose(our_out, torch_out, rtol=1e-4, atol=1e-5)

    def test_default_groups_is_n_channels(self):
        """Default groups should be n_channels (depthwise)."""
        np.random.seed(42)
        n_channels = 32
        data = np.random.randn(n_channels, 500).astype(np.float32)
        filters = np.random.randn(n_channels, 32).astype(np.float32)

        result_default = correlate(data, filters)
        result_explicit = correlate(data, filters, groups=n_channels)

        np.testing.assert_array_equal(result_default, result_explicit)


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
