"""
Stochastic resonance: when noise improves signal detection.

In a nonlinear threshold system, adding the *right* amount of noise can
push a subthreshold signal above the detection threshold, improving the
output signal-to-noise ratio.  This counterintuitive phenomenon is well-
established in neurobiology and has engineered analogues in dither and
sigma-delta conversion.

This module simulates the canonical double-well system and a threshold
crossing detector to demonstrate the non-monotonic SNR-vs-noise curve.

References
----------
* Benzi, R., Sutera, A., & Vulpiani, A. (1981). The mechanism of stochastic
  resonance. *J. Phys. A*, 14(11), L453-L457.
* Gammaitoni, L., Hänggi, P., Jung, P., & Marchesoni, F. (1998). Stochastic
  resonance. *Rev. Mod. Phys.*, 70(1), 223-287.
"""

import numpy as np
from scipy.signal import welch


def double_well_system(n: int, A: float = 0.1, f_sig: float = 0.01,
                       sigma: float = 0.3, dt: float = 0.1,
                       seed: int | None = None) -> np.ndarray:
    """Simulate an overdamped particle in a symmetric double-well potential
    with a weak periodic forcing and additive Gaussian noise.

    The Langevin equation is
        dx/dt = x - x^3 + A*cos(2*pi*f_sig*t) + sigma*xi(t)
    where xi(t) is white Gaussian noise with <xi(t)xi(t')> = delta(t-t').

    Without noise (sigma=0) and with A < sqrt(4/27) ≈ 0.385, the particle
    stays in one well.  With the right noise, it hops between wells in
    sync with the periodic forcing.

    Parameters
    ----------
    n : int
        Number of time steps.
    A : float
        Amplitude of the periodic forcing (subthreshold if A < 0.385).
    f_sig : float
        Forcing frequency (normalised to 1/dt).
    sigma : float
        Noise intensity.
    dt : float
        Time step.
    seed : int or None

    Returns
    -------
    x : ndarray of shape (n,)
        Particle position over time.
    """
    rng = np.random.default_rng(seed)
    x = np.zeros(n)
    x[0] = 1.0  # start in right well
    t = np.arange(n) * dt

    for i in range(n - 1):
        forcing = A * np.cos(2 * np.pi * f_sig * t[i])
        drift = x[i] - x[i]**3 + forcing
        # Euler-Maruyama: the white-noise increment scales with sqrt(dt).
        dx = drift * dt + sigma * np.sqrt(dt) * rng.standard_normal()
        # Clip to prevent Euler blow-up at high sigma.
        x[i+1] = np.clip(x[i] + dx, -5.0, 5.0)

    return x


def threshold_crossing_detector(x: np.ndarray, threshold: float = 0.0) -> np.ndarray:
    """Simple threshold detector: output 1 when x > threshold, 0 otherwise.

    Parameters
    ----------
    x : ndarray
        Input signal.
    threshold : float
        Detection threshold.

    Returns
    -------
    y : ndarray
        Binary detection output.
    """
    return (x > threshold).astype(float)


def snr_vs_noise_curve(n: int = 50000, A: float = 0.1, f_sig: float = 0.005,
                       sigmas: np.ndarray | None = None,
                       dt: float = 0.1, seed: int = 42) -> dict:
    """Compute the SNR-vs-noise-intensity curve.

    For each sigma in *sigmas*, simulates the double-well system, applies
    a threshold detector, and measures the output SNR at the forcing
    frequency.  The resulting curve should be non-monotonic: a peak at
    the optimal noise level.

    Parameters
    ----------
    n : int
        Number of time steps per simulation.
    A : float
        Forcing amplitude.
    f_sig : float
        Forcing frequency.
    sigmas : ndarray or None
        Noise intensities to sweep.  Defaults to logspace(-1.5, 0.5, 20).
    dt : float
        Time step.
    seed : int
        Base RNG seed (incremented per sigma).

    Returns
    -------
    result : dict
        Keys: 'sigmas' (swept values), 'snr_db' (output SNR at each sigma),
        'optimal_sigma' (sigma giving maximum SNR).
    """
    if sigmas is None:
        sigmas = np.logspace(-1.5, 1.0, 25)

    snr_vals = np.zeros(len(sigmas))

    for j, sigma in enumerate(sigmas):
        x = double_well_system(n, A=A, f_sig=f_sig, sigma=sigma, dt=dt, seed=seed + j)
        y = threshold_crossing_detector(x)

        # SNR at the forcing frequency.
        fs = 1.0 / dt
        f, psd = welch(y, fs, nperseg=min(4096, n // 4))
        idx_sig = np.argmin(np.abs(f - f_sig))
        # Signal power: integrate around the forcing frequency.
        half_width = max(1, len(f) // 200)
        i0 = max(1, idx_sig - half_width)
        i1 = min(len(f) - 1, idx_sig + half_width + 1)
        p_sig = np.sum(psd[i0:i1])
        p_noise = np.sum(psd[1:]) - p_sig
        snr_vals[j] = 10 * np.log10(p_sig / p_noise) if p_noise > 0 else 0.0

    best_idx = np.argmax(snr_vals)
    return {
        'sigmas': sigmas,
        'snr_db': snr_vals,
        'optimal_sigma': float(sigmas[best_idx]),
        'optimal_snr_db': float(snr_vals[best_idx]),
    }