"""Tests for stochastic_resonance.py."""

import numpy as np
import pytest
from stochastic_resonance import (double_well_system, threshold_crossing_detector,
                                   snr_vs_noise_curve)


class TestDoubleWell:
    def test_shape(self):
        x = double_well_system(1000, seed=42)
        assert x.shape == (1000,)

    def test_stays_in_one_well_without_noise(self):
        x = double_well_system(5000, A=0.1, sigma=0.0, seed=1)
        # Without noise, particle should stay in the starting well (x > 0).
        assert np.all(x > 0)

    def test_hops_with_strong_noise(self):
        x = double_well_system(10000, A=0.3, sigma=0.5, dt=0.1, seed=1)
        # With strong noise, should visit both wells.
        assert np.any(x > 0.5) and np.any(x < -0.5)

    def test_reproducible(self):
        a = double_well_system(500, seed=42)
        b = double_well_system(500, seed=42)
        assert np.array_equal(a, b)


class TestThresholdDetector:
    def test_binary_output(self):
        x = np.array([-1.0, -0.1, 0.0, 0.5, 2.0])
        y = threshold_crossing_detector(x, threshold=0.0)
        assert np.array_equal(y, [0, 0, 0, 1, 1])

    def test_custom_threshold(self):
        x = np.array([-1.0, 0.0, 1.0, 2.0])
        y = threshold_crossing_detector(x, threshold=1.5)
        assert np.array_equal(y, [0, 0, 0, 1])


class TestSNRCurve:
    def test_returns_keys(self):
        result = snr_vs_noise_curve(n=5000, sigmas=np.array([0.1, 0.3, 0.5]), seed=42)
        assert 'sigmas' in result and 'snr_db' in result
        assert 'optimal_sigma' in result and 'optimal_snr_db' in result

    def test_optimal_sigma_is_finite(self):
        result = snr_vs_noise_curve(n=5000, sigmas=np.array([0.1, 0.3, 0.5, 0.7]), seed=42)
        assert result['optimal_sigma'] > 0
        assert np.isfinite(result['optimal_snr_db'])

    def test_non_monotonic(self):
        """The SR signature: SNR rises to an interior peak, then falls again.

        A monotonically rising or falling curve must fail this test.
        """
        sigmas = np.array([0.05, 0.15, 0.25, 0.5, 1.0, 2.0])
        result = snr_vs_noise_curve(n=20000, sigmas=sigmas, seed=42)
        snr = result['snr_db']
        peak = int(np.argmax(snr))
        assert 0 < peak < len(sigmas) - 1, "SNR peak must be interior, not at an endpoint"
        assert snr[peak] > snr[0] + 3, "SNR should rise clearly (>3 dB) from low noise to the peak"
        assert snr[peak] > snr[-1] + 3, "SNR should fall clearly (>3 dB) from the peak to high noise"

    def test_noise_increment_scales_with_sqrt_dt(self):
        """Euler-Maruyama regression: per-step noise std must be sigma*sqrt(dt).

        The original implementation cancelled the sqrt(dt) factors, making
        the increment independent of dt.
        """
        sigma = 0.25
        for dt in (0.1, 0.01):
            x = double_well_system(20000, A=0.0, f_sig=0.005, sigma=sigma,
                                   dt=dt, seed=3)
            # Remove the deterministic drift contribution by looking at
            # increments around the stable well x=1: drift is locally ~0,
            # so increment std approximates the noise term.
            dx = np.diff(x[np.abs(x - 1.0) < 0.2])
            measured = np.std(dx)
            expected = sigma * np.sqrt(dt)
            assert abs(measured - expected) / expected < 0.25, (
                f"dt={dt}: increment std {measured:.4f}, expected ~{expected:.4f}")