"""
Validation tests comparing zerostone (zpybci) against scipy.signal.

These tests verify that zerostone's signal processing primitives produce
correct results. Key considerations:

- zerostone uses f32 (single precision), scipy uses f64 (double precision)
- zerostone's IIR Butterworth filters use proper pole placement per section,
  matching scipy.signal.butter output to within f32 precision
- FIR filters should match exactly (same algorithm)
"""

import numpy as np
import pytest

# Skip all tests if scipy is not installed
scipy = pytest.importorskip("scipy")
from scipy import signal


class TestIirLowpassValidation:
    """Validate IIR lowpass filter functional behavior."""

    def test_lowpass_dc_passthrough(self):
        """Verify DC signal passes through lowpass filter."""
        import zpybci as zbci

        sample_rate = 1000.0
        cutoff = 40.0

        zs_filter = zbci.IirFilter.butterworth_lowpass(sample_rate, cutoff)

        # DC signal
        dc_signal = np.ones(1000, dtype=np.float32)
        output = zs_filter.process(dc_signal)

        # After settling, DC should pass through with gain ~1
        assert np.abs(output[-1] - 1.0) < 0.02, f"DC not passed: {output[-1]}"

    def test_lowpass_passband_preserves_signal(self):
        """Verify passband frequencies are preserved."""
        import zpybci as zbci

        sample_rate = 1000.0
        cutoff = 100.0
        test_freq = 20.0  # Well within passband

        zs_filter = zbci.IirFilter.butterworth_lowpass(sample_rate, cutoff)

        # Create test signal
        t = np.arange(0, 2, 1/sample_rate, dtype=np.float32)
        test_signal = np.sin(2 * np.pi * test_freq * t).astype(np.float32)

        # Filter
        output = zs_filter.process(test_signal)

        # Passband signal should be mostly preserved (>90% amplitude)
        input_rms = np.sqrt(np.mean(test_signal[500:] ** 2))
        output_rms = np.sqrt(np.mean(output[500:] ** 2))

        assert output_rms > input_rms * 0.9, \
            f"Passband attenuation too high: {output_rms/input_rms:.2%}"

    def test_lowpass_attenuation_stopband(self):
        """Verify stopband attenuation matches scipy."""
        import zpybci as zbci

        sample_rate = 1000.0
        cutoff = 50.0
        stopband_freq = 200.0  # Well into stopband

        # Create filters
        zs_filter = zbci.IirFilter.butterworth_lowpass(sample_rate, cutoff)
        sos = signal.butter(4, cutoff, btype='low', fs=sample_rate, output='sos')

        # Create test signal in stopband
        t = np.arange(0, 2, 1/sample_rate, dtype=np.float32)
        test_signal = np.sin(2 * np.pi * stopband_freq * t).astype(np.float32)

        # Filter
        zs_output = zs_filter.process(test_signal)
        scipy_output = signal.sosfilt(sos, test_signal.astype(np.float64)).astype(np.float32)

        # Both should attenuate the stopband frequency significantly
        # Compare RMS of outputs (both should be similar and small)
        zs_rms = np.sqrt(np.mean(zs_output[500:] ** 2))
        scipy_rms = np.sqrt(np.mean(scipy_output[500:] ** 2))

        # Both should have similar attenuation (within 20% of each other)
        assert abs(zs_rms - scipy_rms) / max(scipy_rms, 1e-10) < 0.2, \
            f"RMS mismatch: zerostone={zs_rms:.6f}, scipy={scipy_rms:.6f}"

        # Both should attenuate by at least 20 dB (factor of 10)
        input_rms = np.sqrt(np.mean(test_signal ** 2))
        assert zs_rms < input_rms * 0.1, f"Zerostone attenuation insufficient: {zs_rms/input_rms}"
        assert scipy_rms < input_rms * 0.1, f"Scipy attenuation insufficient: {scipy_rms/input_rms}"


class TestIirHighpassValidation:
    """Validate IIR highpass filter against scipy.signal.butter + sosfilt."""

    def test_highpass_dc_rejection(self):
        """Verify both filters reject DC equally."""
        import zpybci as zbci

        sample_rate = 1000.0
        cutoff = 10.0

        # Create filters
        zs_filter = zbci.IirFilter.butterworth_highpass(sample_rate, cutoff)
        sos = signal.butter(4, cutoff, btype='high', fs=sample_rate, output='sos')

        # DC signal
        dc_signal = np.ones(2000, dtype=np.float32)

        # Filter
        zs_output = zs_filter.process(dc_signal)
        scipy_output = signal.sosfilt(sos, dc_signal.astype(np.float64)).astype(np.float32)

        # Both should reject DC
        assert np.abs(zs_output[-1]) < 0.05, f"Zerostone DC rejection failed: {zs_output[-1]}"
        assert np.abs(scipy_output[-1]) < 0.05, f"Scipy DC rejection failed: {scipy_output[-1]}"

    def test_highpass_passband(self):
        """Verify passband signal passes through both filters."""
        import zpybci as zbci

        sample_rate = 1000.0
        cutoff = 10.0
        passband_freq = 100.0  # Well into passband

        # Create filters
        zs_filter = zbci.IirFilter.butterworth_highpass(sample_rate, cutoff)
        sos = signal.butter(4, cutoff, btype='high', fs=sample_rate, output='sos')

        # Create test signal
        t = np.arange(0, 2, 1/sample_rate, dtype=np.float32)
        test_signal = np.sin(2 * np.pi * passband_freq * t).astype(np.float32)

        # Filter
        zs_output = zs_filter.process(test_signal)
        scipy_output = signal.sosfilt(sos, test_signal.astype(np.float64)).astype(np.float32)

        # Both should pass the signal with minimal attenuation
        zs_rms = np.sqrt(np.mean(zs_output[500:] ** 2))
        scipy_rms = np.sqrt(np.mean(scipy_output[500:] ** 2))
        input_rms = np.sqrt(np.mean(test_signal ** 2))

        # Should retain most of the signal (>90%)
        assert zs_rms > input_rms * 0.9, f"Zerostone passband attenuation: {zs_rms/input_rms}"
        assert scipy_rms > input_rms * 0.9, f"Scipy passband attenuation: {scipy_rms/input_rms}"


class TestIirBandpassValidation:
    """Validate IIR bandpass filter functional behavior."""

    def test_bandpass_center_frequency(self):
        """Verify center frequency passes through."""
        import zpybci as zbci

        sample_rate = 1000.0
        low_cutoff = 8.0
        high_cutoff = 12.0
        center_freq = np.sqrt(low_cutoff * high_cutoff)  # ~9.8 Hz

        # Create zerostone filter
        zs_filter = zbci.IirFilter.butterworth_bandpass(sample_rate, low_cutoff, high_cutoff)

        # Create test signal at center frequency
        t = np.arange(0, 3, 1/sample_rate, dtype=np.float32)
        test_signal = np.sin(2 * np.pi * center_freq * t).astype(np.float32)

        # Filter
        zs_output = zs_filter.process(test_signal)

        # Should pass the center frequency with reasonable amplitude
        zs_rms = np.sqrt(np.mean(zs_output[1500:] ** 2))
        input_rms = np.sqrt(np.mean(test_signal ** 2))

        # Should retain significant signal (>30% - narrow bandpass can attenuate)
        assert zs_rms > input_rms * 0.3, f"Center freq attenuation too high: {zs_rms/input_rms}"

    def test_bandpass_stopband_rejection(self):
        """Verify out-of-band frequencies are rejected."""
        import zpybci as zbci

        sample_rate = 1000.0
        low_cutoff = 20.0
        high_cutoff = 40.0

        # Create filters
        zs_filter = zbci.IirFilter.butterworth_bandpass(sample_rate, low_cutoff, high_cutoff)
        sos = signal.butter(4, [low_cutoff, high_cutoff], btype='band', fs=sample_rate, output='sos')

        # Test signal well outside passband (5 Hz)
        t = np.arange(0, 3, 1/sample_rate, dtype=np.float32)
        test_signal = np.sin(2 * np.pi * 5 * t).astype(np.float32)

        # Filter
        zs_output = zs_filter.process(test_signal)
        scipy_output = signal.sosfilt(sos, test_signal.astype(np.float64)).astype(np.float32)

        # Both should heavily attenuate
        zs_rms = np.sqrt(np.mean(zs_output[1000:] ** 2))
        scipy_rms = np.sqrt(np.mean(scipy_output[1000:] ** 2))
        input_rms = np.sqrt(np.mean(test_signal ** 2))

        # Should attenuate by at least 10x
        assert zs_rms < input_rms * 0.1, f"Zerostone stopband leak: {zs_rms/input_rms}"
        assert scipy_rms < input_rms * 0.1, f"Scipy stopband leak: {scipy_rms/input_rms}"


class TestFirFilterValidation:
    """Validate FIR filter against scipy.signal.lfilter."""

    def test_moving_average(self):
        """Compare moving average filter outputs."""
        import zpybci as zbci

        window_size = 5

        # Create zerostone moving average
        zs_filter = zbci.FirFilter.moving_average(window_size)

        # Create scipy equivalent (equal-weight FIR)
        scipy_coeffs = np.ones(window_size) / window_size

        # Test signal
        test_signal = np.random.randn(1000).astype(np.float32)

        # Filter
        zs_output = zs_filter.process(test_signal)
        scipy_output = signal.lfilter(scipy_coeffs, 1.0, test_signal.astype(np.float64)).astype(np.float32)

        # Should match exactly (same algorithm)
        assert np.allclose(zs_output, scipy_output, rtol=1e-5, atol=1e-6), \
            f"Max diff: {np.max(np.abs(zs_output - scipy_output))}"

    def test_custom_fir_taps(self):
        """Compare custom FIR coefficients."""
        import zpybci as zbci

        # Custom FIR coefficients (simple lowpass-ish)
        taps = np.array([0.1, 0.15, 0.25, 0.25, 0.15, 0.1], dtype=np.float32)

        # Create filters
        zs_filter = zbci.FirFilter(taps=taps.tolist())

        # Test signal
        test_signal = np.random.randn(1000).astype(np.float32)

        # Filter
        zs_output = zs_filter.process(test_signal)
        scipy_output = signal.lfilter(taps.astype(np.float64), 1.0, test_signal.astype(np.float64)).astype(np.float32)

        # Should match closely
        assert np.allclose(zs_output, scipy_output, rtol=1e-5, atol=1e-6), \
            f"Max diff: {np.max(np.abs(zs_output - scipy_output))}"

    def test_fir_impulse_response(self):
        """Verify FIR impulse response matches coefficients."""
        import zpybci as zbci

        taps = np.array([0.5, 0.3, 0.2], dtype=np.float32)

        zs_filter = zbci.FirFilter(taps=taps.tolist())

        # Impulse
        impulse = np.zeros(10, dtype=np.float32)
        impulse[0] = 1.0

        output = zs_filter.process(impulse)

        # Impulse response should equal coefficients
        assert np.allclose(output[:3], taps, rtol=1e-6), \
            f"Impulse response mismatch: {output[:3]} vs {taps}"


class TestAcCouplerValidation:
    """Validate AC coupler (high-pass) against scipy single-pole filter."""

    def test_ac_coupler_dc_removal(self):
        """Verify DC removal."""
        import zpybci as zbci

        sample_rate = 1000.0
        cutoff = 0.5  # Very low cutoff for DC removal

        # Create zerostone AC coupler
        ac = zbci.AcCoupler(sample_rate, cutoff)

        # Create scipy single-pole high-pass
        # AC coupler is typically a simple RC high-pass
        # Transfer function: H(z) = (1 - z^-1) / (1 - alpha*z^-1)
        # where alpha = exp(-2*pi*fc/fs)
        alpha = np.exp(-2 * np.pi * cutoff / sample_rate)
        b = np.array([1.0, -1.0]) * (1 + alpha) / 2
        a = np.array([1.0, -alpha])

        # DC + small AC signal
        t = np.arange(0, 5, 1/sample_rate, dtype=np.float32)
        dc_offset = 2.0
        ac_signal = 0.1 * np.sin(2 * np.pi * 10 * t)
        test_signal = (dc_offset + ac_signal).astype(np.float32)

        # Filter
        zs_output = ac.process(test_signal)
        scipy_output = signal.lfilter(b, a, test_signal.astype(np.float64)).astype(np.float32)

        # Both should remove DC
        # After settling, mean should be near zero
        assert np.abs(np.mean(zs_output[2000:])) < 0.1, \
            f"Zerostone DC not removed: mean={np.mean(zs_output[2000:])}"
        assert np.abs(np.mean(scipy_output[2000:])) < 0.1, \
            f"Scipy DC not removed: mean={np.mean(scipy_output[2000:])}"


class TestMedianFilterValidation:
    """Validate median filter against scipy.signal.medfilt."""

    def test_median_spike_removal(self):
        """Compare spike removal capability."""
        import zpybci as zbci

        window_size = 5

        # Create filters
        zs_filter = zbci.MedianFilter(window_size)

        # Signal with spike
        test_signal = np.ones(100, dtype=np.float32)
        test_signal[50] = 100.0  # Spike

        # Filter
        zs_output = zs_filter.process(test_signal)
        scipy_output = signal.medfilt(test_signal.astype(np.float64), kernel_size=window_size).astype(np.float32)

        # Both should remove the spike (after the filter window covers it)
        # Note: zerostone uses causal filter, scipy uses centered
        # So the spike position in output will differ
        # Just verify the spike is removed in the output range where both have seen it
        assert np.max(zs_output[55:]) < 2.0, \
            f"Zerostone didn't remove spike: max={np.max(zs_output[55:])}"
        assert np.max(scipy_output[52:]) < 2.0, \
            f"Scipy didn't remove spike: max={np.max(scipy_output[52:])}"

    def test_median_preserves_edges(self):
        """Verify edge preservation."""
        import zpybci as zbci

        window_size = 3

        zs_filter = zbci.MedianFilter(window_size)

        # Step signal
        test_signal = np.zeros(100, dtype=np.float32)
        test_signal[50:] = 1.0

        zs_output = zs_filter.process(test_signal)
        scipy_output = signal.medfilt(test_signal.astype(np.float64), kernel_size=window_size).astype(np.float32)

        # Both should preserve the step edge (with slight shift due to filter delay)
        # Check that we have both 0 and 1 regions in output
        assert np.min(zs_output) < 0.5, "Zerostone should have low region"
        assert np.max(zs_output) > 0.5, "Zerostone should have high region"
        assert np.min(scipy_output) < 0.5, "Scipy should have low region"
        assert np.max(scipy_output) > 0.5, "Scipy should have high region"


class TestFrequencyResponse:
    """Test frequency response characteristics against scipy."""

    def test_lowpass_frequency_response_shape(self):
        """Verify lowpass has monotonically decreasing response with no resonance."""
        import zpybci as zbci

        sample_rate = 1000.0
        cutoff = 100.0

        zs_filter = zbci.IirFilter.butterworth_lowpass(sample_rate, cutoff)

        # Measure response at several frequencies including cutoff
        test_freqs = [10, 50, 100, 150, 200, 300]
        zs_response = []

        for freq in test_freqs:
            zs_filter.reset()

            t = np.arange(0, 2, 1/sample_rate, dtype=np.float32)
            test_signal = np.sin(2 * np.pi * freq * t).astype(np.float32)
            zs_output = zs_filter.process(test_signal)

            zs_rms = np.sqrt(np.mean(zs_output[1000:] ** 2))
            input_rms = np.sqrt(np.mean(test_signal[1000:] ** 2))
            zs_response.append(zs_rms / input_rms)

        # 1. Passband: gain should be close to 1.0 (no resonance)
        assert 0.95 < zs_response[0] < 1.05, f"10 Hz gain: {zs_response[0]}"
        assert 0.90 < zs_response[1] < 1.05, f"50 Hz gain: {zs_response[1]}"

        # 2. Cutoff: gain should be ~0.707 (-3dB)
        assert 0.65 < zs_response[2] < 0.75, f"Cutoff gain: {zs_response[2]} (expect ~0.707)"

        # 3. No gain > 1.0 at any frequency (no resonance)
        for i, (freq, gain) in enumerate(zip(test_freqs, zs_response)):
            assert gain <= 1.05, f"Resonance detected at {freq} Hz: gain={gain}"

        # 4. Stopband attenuation
        assert zs_response[3] < 0.3, f"150 Hz should be attenuated: {zs_response[3]}"
        assert zs_response[4] < 0.05, f"200 Hz should be heavily attenuated: {zs_response[4]}"
        assert zs_response[5] < 0.005, f"300 Hz should be very attenuated: {zs_response[5]}"

    def test_frequency_response_matches_scipy(self):
        """Verify frequency response matches scipy across the spectrum."""
        import zpybci as zbci

        sample_rate = 1000.0
        cutoff = 100.0

        zs_filter = zbci.IirFilter.butterworth_lowpass(sample_rate, cutoff)
        sos = signal.butter(4, cutoff, btype='low', fs=sample_rate, output='sos')

        # Test at multiple frequencies
        test_freqs = np.linspace(5, 400, 50)
        max_db_error = 0.0

        for freq in test_freqs:
            zs_filter.reset()

            t = np.arange(0, 2, 1/sample_rate, dtype=np.float32)
            test_signal = np.sin(2 * np.pi * freq * t).astype(np.float32)

            zs_output = zs_filter.process(test_signal)
            scipy_output = signal.sosfilt(sos, test_signal.astype(np.float64)).astype(np.float32)

            zs_rms = np.sqrt(np.mean(zs_output[1000:] ** 2))
            scipy_rms = np.sqrt(np.mean(scipy_output[1000:] ** 2))

            if scipy_rms > 1e-6:
                db_error = abs(20 * np.log10(max(zs_rms, 1e-10)) - 20 * np.log10(max(scipy_rms, 1e-10)))
                max_db_error = max(max_db_error, db_error)

        assert max_db_error < 0.5, f"Max dB error vs scipy: {max_db_error:.3f} dB"

    def test_lowpass_matches_scipy_sos(self):
        """Compare SOS coefficients directly against scipy."""
        import zpybci as zbci

        sample_rate = 1000.0
        cutoff = 100.0

        # Get scipy SOS
        sos = signal.butter(4, cutoff, btype='low', fs=sample_rate, output='sos')

        # Get zerostone filter and measure response at cutoff
        zs_filter = zbci.IirFilter.butterworth_lowpass(sample_rate, cutoff)
        zs_filter.reset()

        t = np.arange(0, 2, 1/sample_rate, dtype=np.float32)
        test_signal = np.sin(2 * np.pi * cutoff * t).astype(np.float32)

        zs_output = zs_filter.process(test_signal)
        scipy_output = signal.sosfilt(sos, test_signal.astype(np.float64)).astype(np.float32)

        # Both should give -3dB at cutoff
        zs_rms = np.sqrt(np.mean(zs_output[1000:] ** 2))
        scipy_rms = np.sqrt(np.mean(scipy_output[1000:] ** 2))
        input_rms = np.sqrt(np.mean(test_signal[1000:] ** 2))

        zs_gain = zs_rms / input_rms
        scipy_gain = scipy_rms / input_rms

        assert abs(zs_gain - scipy_gain) < 0.02, \
            f"Gain mismatch at cutoff: zerostone={zs_gain:.4f}, scipy={scipy_gain:.4f}"


class TestNumericalPrecision:
    """Test numerical precision and edge cases."""

    def test_f32_precision_sufficient(self):
        """Verify f32 precision is sufficient for typical BCI signals."""
        import zpybci as zbci

        sample_rate = 1000.0
        cutoff = 40.0

        # Create filter
        zs_filter = zbci.IirFilter.butterworth_lowpass(sample_rate, cutoff)

        # Typical EEG amplitude range: ~10-100 microvolts
        # Simulate with normalized range
        np.random.seed(42)
        test_signal = np.random.randn(10000).astype(np.float32) * 0.0001  # ~100 uV scale

        # Filter should work without numerical issues
        output = zs_filter.process(test_signal)

        # Check for NaN/Inf
        assert not np.any(np.isnan(output)), "Filter produced NaN"
        assert not np.any(np.isinf(output)), "Filter produced Inf"

        # Output should be reasonable (not exploded)
        assert np.max(np.abs(output)) < 0.001, "Filter output unreasonably large"

    def test_long_signal_stability(self):
        """Verify filter is stable for long signals."""
        import zpybci as zbci

        sample_rate = 1000.0
        cutoff = 40.0

        zs_filter = zbci.IirFilter.butterworth_lowpass(sample_rate, cutoff)

        # 10 minutes of data at 1 kHz
        duration_seconds = 600
        signal_length = int(duration_seconds * sample_rate)

        # Process in chunks (like real-time)
        chunk_size = 1000
        max_output = 0.0

        np.random.seed(42)
        for i in range(0, signal_length, chunk_size):
            chunk = np.random.randn(chunk_size).astype(np.float32)
            output = zs_filter.process(chunk)
            max_output = max(max_output, np.max(np.abs(output)))

        # Filter should remain stable
        assert max_output < 10.0, f"Filter became unstable: max output = {max_output}"
        assert not np.isnan(max_output), "Filter produced NaN"


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