//! Particle Filter / Sequential Monte Carlo for nonlinear state estimation.
//!
//! Implements a bootstrap particle filter with:
//! - Xoshiro256++ PRNG for fast, high-quality randomness
//! - Box-Muller transform for Gaussian noise generation
//! - Systematic resampling (low-variance, O(N))
//! - Effective sample size (ESS) monitoring with adaptive resampling
//!
//! Ideal for tracking latent state in nonlinear/non-Gaussian systems:
//! stochastic volatility, jump-diffusion models, regime-switching.
//!
//! Pro tier required.

use pyo3::prelude::*;
use std::sync::Mutex;

// ===========================================================================
// Xoshiro256++ PRNG
// ===========================================================================

struct Xoshiro256PlusPlus {
    s: [u64; 4],
}

impl Xoshiro256PlusPlus {
    fn new(seed: u64) -> Self {
        // Use SplitMix64 to initialize state from a single seed
        let mut state = seed;
        let mut s = [0u64; 4];
        for slot in &mut s {
            state = state.wrapping_add(0x9E3779B97F4A7C15);
            let mut z = state;
            z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
            z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
            z ^= z >> 31;
            *slot = z;
        }
        Self { s }
    }

    #[inline]
    fn next_u64(&mut self) -> u64 {
        let result = (self.s[0].wrapping_add(self.s[3]))
            .rotate_left(23)
            .wrapping_add(self.s[0]);
        let t = self.s[1] << 17;
        self.s[2] ^= self.s[0];
        self.s[3] ^= self.s[1];
        self.s[1] ^= self.s[2];
        self.s[0] ^= self.s[3];
        self.s[2] ^= t;
        self.s[3] = self.s[3].rotate_left(45);
        result
    }

    /// Uniform f64 in [0, 1).
    #[inline]
    fn next_f64(&mut self) -> f64 {
        let bits = self.next_u64() >> 11;
        bits as f64 * (1.0 / (1u64 << 53) as f64)
    }
}

// ===========================================================================
// Box-Muller normal generation
// ===========================================================================

/// Generate a pair of standard normal random variables using Box-Muller.
#[inline]
fn box_muller(rng: &mut Xoshiro256PlusPlus) -> (f64, f64) {
    loop {
        let u1 = rng.next_f64();
        let u2 = rng.next_f64();
        if u1 > 1e-15 {
            let r = (-2.0 * u1.ln()).sqrt();
            let theta = 2.0 * std::f64::consts::PI * u2;
            return (r * theta.cos(), r * theta.sin());
        }
    }
}

/// Generate a single standard normal random variable.
#[inline]
fn randn(rng: &mut Xoshiro256PlusPlus) -> f64 {
    box_muller(rng).0
}

// ===========================================================================
// Result type
// ===========================================================================

/// State returned after each particle filter update step.
#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct PFState {
    /// Weighted mean estimate of the latent state.
    #[pyo3(get)]
    pub estimate: f64,
    /// Weighted variance of the particle distribution.
    #[pyo3(get)]
    pub variance: f64,
    /// Effective sample size (1 = degenerate, n_particles = perfect).
    #[pyo3(get)]
    pub ess: f64,
    /// Whether systematic resampling was triggered this step.
    #[pyo3(get)]
    pub resampled: bool,
}

#[pymethods]
impl PFState {
    fn __repr__(&self) -> String {
        format!(
            "PFState(estimate={:.6}, variance={:.6}, ess={:.1}, resampled={})",
            self.estimate, self.variance, self.ess, self.resampled
        )
    }
}

// ===========================================================================
// Particle Filter inner state
// ===========================================================================

/// ESS threshold for resampling: resample when ESS < n_particles * ESS_FRACTION.
const ESS_FRACTION: f64 = 0.5;

/// Maximum number of particles to prevent memory issues.
const MAX_PARTICLES: usize = 100_000;

struct PFInner {
    /// Number of particles.
    n_particles: usize,
    /// Particle state values.
    particles: Vec<f64>,
    /// Normalized particle weights (sum to 1).
    weights: Vec<f64>,
    /// Standard deviation of process noise (state transition).
    process_noise: f64,
    /// Standard deviation of measurement noise (observation likelihood).
    measurement_noise: f64,
    /// PRNG state.
    rng: Xoshiro256PlusPlus,
    /// Scratch buffer for resampling (avoid allocation).
    scratch: Vec<f64>,
}

impl PFInner {
    fn new(
        n_particles: usize,
        initial_state: f64,
        process_noise: f64,
        measurement_noise: f64,
        seed: u64,
    ) -> Self {
        let mut rng = Xoshiro256PlusPlus::new(seed);

        // Initialize particles around the initial state with process noise spread
        let mut particles = Vec::with_capacity(n_particles);
        for _ in 0..n_particles {
            particles.push(initial_state + process_noise * randn(&mut rng));
        }

        let uniform_w = 1.0 / n_particles as f64;
        let weights = vec![uniform_w; n_particles];
        let scratch = vec![0.0; n_particles];

        Self {
            n_particles,
            particles,
            weights,
            process_noise,
            measurement_noise,
            rng,
            scratch,
        }
    }

    /// Propagate all particles through the state transition model.
    /// Simple random walk: x_t = x_{t-1} + w, w ~ N(0, process_noise^2).
    fn propagate(&mut self) {
        for p in self.particles.iter_mut() {
            *p += self.process_noise * randn(&mut self.rng);
        }
    }

    /// Compute unnormalized log-weights based on observation likelihood.
    /// Assumes Gaussian observation model: p(z | x) = N(z; x, measurement_noise^2).
    fn compute_log_weights(&self, observation: f64) -> Vec<f64> {
        let var = self.measurement_noise * self.measurement_noise;
        let inv_2var = if var > 1e-15 { 1.0 / (2.0 * var) } else { 1e15 };
        self.particles
            .iter()
            .map(|&x| {
                let diff = observation - x;
                -diff * diff * inv_2var
            })
            .collect()
    }

    /// Normalize log-weights to weights using log-sum-exp trick.
    fn normalize_log_weights(&mut self, log_weights: &[f64]) {
        // Find max for numerical stability
        let max_lw = log_weights
            .iter()
            .cloned()
            .fold(f64::NEG_INFINITY, f64::max);

        if !max_lw.is_finite() {
            // All weights are -inf: reset to uniform
            let uniform = 1.0 / self.n_particles as f64;
            for w in self.weights.iter_mut() {
                *w = uniform;
            }
            return;
        }

        // exp(log_w - max) and sum
        let mut sum = 0.0;
        for (i, &lw) in log_weights.iter().enumerate() {
            let w = (lw - max_lw).exp();
            self.weights[i] = w;
            sum += w;
        }

        // Normalize
        if sum > 0.0 {
            for w in self.weights.iter_mut() {
                *w /= sum;
            }
        } else {
            let uniform = 1.0 / self.n_particles as f64;
            for w in self.weights.iter_mut() {
                *w = uniform;
            }
        }
    }

    /// Compute effective sample size: ESS = 1 / sum(w_i^2).
    fn effective_sample_size(&self) -> f64 {
        let sum_sq: f64 = self.weights.iter().map(|w| w * w).sum();
        if sum_sq > 0.0 {
            1.0 / sum_sq
        } else {
            0.0
        }
    }

    /// Systematic resampling (low-variance, O(N)).
    /// Resamples particles in-place using the scratch buffer.
    fn systematic_resample(&mut self) {
        let n = self.n_particles;
        let u0 = self.rng.next_f64() / n as f64;

        let mut cumsum = 0.0;
        let mut j = 0;

        for i in 0..n {
            let u = u0 + i as f64 / n as f64;
            while j < n - 1 && cumsum + self.weights[j] < u {
                cumsum += self.weights[j];
                j += 1;
            }
            self.scratch[i] = self.particles[j];
        }

        // Copy back
        self.particles.copy_from_slice(&self.scratch);

        // Reset weights to uniform
        let uniform = 1.0 / n as f64;
        for w in self.weights.iter_mut() {
            *w = uniform;
        }
    }

    /// Weighted mean of particles.
    fn state_estimate(&self) -> f64 {
        self.particles
            .iter()
            .zip(self.weights.iter())
            .map(|(x, w)| x * w)
            .sum()
    }

    /// Weighted variance of particles.
    fn state_variance(&self) -> f64 {
        let mean = self.state_estimate();
        self.particles
            .iter()
            .zip(self.weights.iter())
            .map(|(x, w)| w * (x - mean).powi(2))
            .sum()
    }

    /// Full update cycle: propagate -> weight -> (optionally) resample.
    fn update(&mut self, observation: f64) -> PFState {
        // 1. Propagate particles through state transition
        self.propagate();

        // 2. Compute weights from observation likelihood
        let log_weights = self.compute_log_weights(observation);
        self.normalize_log_weights(&log_weights);

        // 3. Compute state estimate and ESS
        let estimate = self.state_estimate();
        let variance = self.state_variance();
        let ess = self.effective_sample_size();

        // 4. Resample if ESS is too low
        let threshold = ESS_FRACTION * self.n_particles as f64;
        let resampled = ess < threshold;
        if resampled {
            self.systematic_resample();
        }

        PFState {
            estimate,
            variance,
            ess,
            resampled,
        }
    }

    /// Reset particles to a new initial state.
    fn reset(&mut self, initial_state: f64) {
        for p in self.particles.iter_mut() {
            *p = initial_state + self.process_noise * randn(&mut self.rng);
        }
        let uniform = 1.0 / self.n_particles as f64;
        for w in self.weights.iter_mut() {
            *w = uniform;
        }
    }
}

// ===========================================================================
// PyO3 interface
// ===========================================================================

/// Bootstrap Particle Filter for nonlinear state estimation.
///
/// Tracks a latent scalar state using Sequential Monte Carlo:
/// - State model: x_t = x_{t-1} + w,  w ~ N(0, process_noise^2)
/// - Obs model:   z_t = x_t + v,      v ~ N(0, measurement_noise^2)
///
/// Adaptive systematic resampling triggers when ESS drops below 50%
/// of the particle count.
///
/// Pro tier required.
#[pyclass]
pub struct ParticleFilter {
    inner: Mutex<PFInner>,
}

#[pymethods]
impl ParticleFilter {
    /// Create a new particle filter.
    ///
    /// Args:
    ///     n_particles: Number of particles (100-100000).
    ///     initial_state: Initial state estimate.
    ///     process_noise: Std dev of state transition noise.
    ///     measurement_noise: Std dev of observation noise.
    ///     seed: RNG seed for reproducibility.
    #[new]
    #[pyo3(signature = (n_particles, initial_state, process_noise, measurement_noise, seed=42))]
    fn new(
        n_particles: usize,
        initial_state: f64,
        process_noise: f64,
        measurement_noise: f64,
        seed: u64,
    ) -> PyResult<Self> {
        crate::auth::require_pro()
            .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
        if n_particles < 10 {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "n_particles must be >= 10",
            ));
        }
        if n_particles > MAX_PARTICLES {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "n_particles must be <= {}, got {}",
                MAX_PARTICLES, n_particles
            )));
        }
        if process_noise <= 0.0 || !process_noise.is_finite() {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "process_noise must be > 0 and finite",
            ));
        }
        if measurement_noise <= 0.0 || !measurement_noise.is_finite() {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "measurement_noise must be > 0 and finite",
            ));
        }
        if !initial_state.is_finite() {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "initial_state must be finite",
            ));
        }
        Ok(Self {
            inner: Mutex::new(PFInner::new(
                n_particles,
                initial_state,
                process_noise,
                measurement_noise,
                seed,
            )),
        })
    }

    /// Run one update cycle: propagate, weight, and conditionally resample.
    ///
    /// Args:
    ///     observation: Observed value.
    ///
    /// Returns:
    ///     PFState with estimate, variance, ESS, and resampled flag.
    fn update(&self, observation: f64) -> PyResult<PFState> {
        crate::auth::require_pro()
            .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
        if !observation.is_finite() {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "observation must be finite",
            ));
        }
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        Ok(inner.update(observation))
    }

    /// Get the current weighted mean state estimate.
    fn state_estimate(&self) -> f64 {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        inner.state_estimate()
    }

    /// Get the current weighted variance of the particle distribution.
    fn state_variance(&self) -> f64 {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        inner.state_variance()
    }

    /// Get the current effective sample size.
    fn effective_sample_size(&self) -> f64 {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        inner.effective_sample_size()
    }

    /// Get all particle state values.
    fn particles(&self) -> Vec<f64> {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        inner.particles.clone()
    }

    /// Get all particle weights.
    fn weights(&self) -> Vec<f64> {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        inner.weights.clone()
    }

    /// Reset all particles to a new initial state.
    ///
    /// Args:
    ///     initial_state: New center for the particle cloud.
    fn reset(&self, initial_state: f64) -> PyResult<()> {
        if !initial_state.is_finite() {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "initial_state must be finite",
            ));
        }
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        inner.reset(initial_state);
        Ok(())
    }

    /// Get the number of particles.
    fn n_particles(&self) -> usize {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        inner.n_particles
    }

    /// Compute the weighted median of the particle distribution.
    /// More robust to outliers than the mean estimate.
    fn state_median(&self) -> f64 {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());

        // Sort particles by value, keeping weights aligned
        let mut indexed: Vec<(f64, f64)> = inner
            .particles
            .iter()
            .zip(inner.weights.iter())
            .map(|(&x, &w)| (x, w))
            .collect();
        indexed.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));

        // Find weighted median: cumulative weight crosses 0.5
        let mut cumsum = 0.0;
        for &(x, w) in &indexed {
            cumsum += w;
            if cumsum >= 0.5 {
                return x;
            }
        }
        // Fallback
        indexed.last().map(|&(x, _)| x).unwrap_or(0.0)
    }

    /// Compute credible interval containing the specified probability mass.
    ///
    /// Args:
    ///     level: Probability mass (e.g. 0.95 for 95% CI). Default 0.95.
    ///
    /// Returns:
    ///     Tuple (lower, upper) bounding the credible interval.
    #[pyo3(signature = (level=0.95))]
    fn credible_interval(&self, level: f64) -> PyResult<(f64, f64)> {
        if level <= 0.0 || level >= 1.0 {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "level must be in (0, 1)",
            ));
        }
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());

        let mut indexed: Vec<(f64, f64)> = inner
            .particles
            .iter()
            .zip(inner.weights.iter())
            .map(|(&x, &w)| (x, w))
            .collect();
        indexed.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));

        let tail = (1.0 - level) / 2.0;
        let mut cumsum = 0.0;
        let mut lower = indexed[0].0;
        let mut upper = indexed.last().unwrap().0;

        for &(x, w) in &indexed {
            cumsum += w;
            if cumsum >= tail && lower == indexed[0].0 {
                lower = x;
            }
            if cumsum >= 1.0 - tail {
                upper = x;
                break;
            }
        }

        Ok((lower, upper))
    }

    fn __repr__(&self) -> String {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        format!(
            "ParticleFilter(n={}, estimate={:.6}, ess={:.1})",
            inner.n_particles,
            inner.state_estimate(),
            inner.effective_sample_size()
        )
    }
}

// ===========================================================================
// Registration
// ===========================================================================

pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<ParticleFilter>()?;
    m.add_class::<PFState>()?;
    Ok(())
}

// ===========================================================================
// Tests
// ===========================================================================

#[cfg(test)]
mod tests {
    use super::*;

    // -----------------------------------------------------------------------
    // PRNG tests
    // -----------------------------------------------------------------------

    #[test]
    fn test_xoshiro_deterministic() {
        let mut rng1 = Xoshiro256PlusPlus::new(42);
        let mut rng2 = Xoshiro256PlusPlus::new(42);
        for _ in 0..100 {
            assert_eq!(rng1.next_u64(), rng2.next_u64());
        }
    }

    #[test]
    fn test_xoshiro_range() {
        let mut rng = Xoshiro256PlusPlus::new(123);
        for _ in 0..10000 {
            let v = rng.next_f64();
            assert!(v >= 0.0 && v < 1.0, "out of range: {}", v);
        }
    }

    #[test]
    fn test_box_muller_statistics() {
        let mut rng = Xoshiro256PlusPlus::new(42);
        let n = 50000;
        let mut sum = 0.0;
        let mut sum_sq = 0.0;
        for _ in 0..n {
            let (z1, z2) = box_muller(&mut rng);
            sum += z1 + z2;
            sum_sq += z1 * z1 + z2 * z2;
        }
        let count = 2 * n;
        let mean = sum / count as f64;
        let var = sum_sq / count as f64 - mean * mean;

        // Mean should be near 0
        assert!(mean.abs() < 0.05, "mean = {}", mean);
        // Variance should be near 1
        assert!((var - 1.0).abs() < 0.1, "variance = {}", var);
    }

    // -----------------------------------------------------------------------
    // Systematic resampling
    // -----------------------------------------------------------------------

    #[test]
    fn test_systematic_resample_uniform() {
        let mut pf = PFInner::new(100, 0.0, 1.0, 1.0, 42);
        // Set uniform weights -- resampling should approximately preserve distribution
        let uniform = 1.0 / 100.0;
        for w in pf.weights.iter_mut() {
            *w = uniform;
        }
        let particles_before: Vec<f64> = pf.particles.clone();
        pf.systematic_resample();

        // After resampling uniform weights, particles should be a permutation/selection
        // All weights should be uniform again
        for &w in &pf.weights {
            assert!((w - uniform).abs() < 1e-12);
        }
        // Mean should be approximately preserved
        let mean_before: f64 = particles_before.iter().sum::<f64>() / 100.0;
        let mean_after: f64 = pf.particles.iter().sum::<f64>() / 100.0;
        assert!(
            (mean_before - mean_after).abs() < 1.0,
            "mean shifted too much"
        );
    }

    #[test]
    fn test_systematic_resample_concentrated() {
        let n = 1000;
        let mut pf = PFInner::new(n, 0.0, 0.1, 1.0, 42);

        // Concentrate all weight on one particle
        for w in pf.weights.iter_mut() {
            *w = 0.0;
        }
        pf.weights[0] = 1.0;
        let target = pf.particles[0];

        pf.systematic_resample();

        // All particles should now equal the high-weight particle
        for &p in &pf.particles {
            assert!(
                (p - target).abs() < 1e-15,
                "particle {} != target {}",
                p,
                target
            );
        }
    }

    // -----------------------------------------------------------------------
    // Particle filter convergence
    // -----------------------------------------------------------------------

    #[test]
    fn test_pf_constant_tracking() {
        // Track a constant value (5.0) with noisy observations
        let mut pf = PFInner::new(1000, 5.0, 0.1, 1.0, 42);

        // Observations near 5.0
        let observations = vec![5.2, 4.8, 5.1, 4.9, 5.0, 5.3, 4.7, 5.1, 5.0, 4.9,
                                5.1, 5.0, 4.9, 5.2, 4.8, 5.0, 5.1, 4.9, 5.0, 5.1];

        let mut last_state = PFState {
            estimate: 0.0,
            variance: 0.0,
            ess: 0.0,
            resampled: false,
        };

        for &obs in &observations {
            last_state = pf.update(obs);
        }

        // Estimate should be near 5.0
        assert!(
            (last_state.estimate - 5.0).abs() < 0.5,
            "estimate = {}",
            last_state.estimate
        );
        // Variance should be small
        assert!(
            last_state.variance < 1.0,
            "variance = {}",
            last_state.variance
        );
    }

    #[test]
    fn test_pf_step_change() {
        // Track a step change: state goes from 0 to 10
        let mut pf = PFInner::new(2000, 0.0, 0.5, 1.0, 42);

        // 20 observations near 0, then 20 near 10
        for _ in 0..20 {
            pf.update(0.0);
        }
        let est_before = pf.state_estimate();
        assert!(est_before.abs() < 2.0, "before = {}", est_before);

        for _ in 0..20 {
            pf.update(10.0);
        }
        let est_after = pf.state_estimate();
        assert!(
            (est_after - 10.0).abs() < 3.0,
            "after = {}",
            est_after
        );
    }

    #[test]
    fn test_pf_ess_decreases() {
        // After weighting without resampling, ESS should decrease
        // when observation is far from particles
        let pf = PFInner::new(100, 0.0, 0.01, 0.01, 42);

        // Initial ESS should be near n_particles
        let initial_ess = pf.effective_sample_size();
        assert!(
            (initial_ess - 100.0).abs() < 1.0,
            "initial ess = {}",
            initial_ess
        );
    }

    #[test]
    fn test_pf_reset() {
        let mut pf = PFInner::new(500, 0.0, 1.0, 1.0, 42);
        for _ in 0..10 {
            pf.update(100.0);
        }
        let est = pf.state_estimate();
        assert!(est > 1.0); // should have moved toward 100

        pf.reset(0.0);
        let est_after = pf.state_estimate();
        assert!(
            est_after.abs() < 3.0,
            "after reset = {}",
            est_after
        );

        // Weights should be uniform
        let uniform = 1.0 / 500.0;
        for &w in &pf.weights {
            assert!((w - uniform).abs() < 1e-12);
        }
    }

    #[test]
    fn test_pf_state_variance() {
        let pf = PFInner::new(1000, 5.0, 0.1, 1.0, 42);
        let var = pf.state_variance();
        // Initial variance should be roughly process_noise^2
        assert!(var > 0.0, "variance should be positive");
        assert!(var < 1.0, "variance = {} too large", var);
    }

    #[test]
    fn test_pf_weights_sum_to_one() {
        let mut pf = PFInner::new(100, 0.0, 1.0, 1.0, 42);
        for _ in 0..10 {
            pf.update(1.0);
        }
        let sum: f64 = pf.weights.iter().sum();
        assert!(
            (sum - 1.0).abs() < 1e-10,
            "weights sum = {}",
            sum
        );
    }

    #[test]
    fn test_pf_deterministic_seed() {
        let mut pf1 = PFInner::new(100, 0.0, 1.0, 1.0, 42);
        let mut pf2 = PFInner::new(100, 0.0, 1.0, 1.0, 42);

        for _ in 0..5 {
            let s1 = pf1.update(1.0);
            let s2 = pf2.update(1.0);
            assert!(
                (s1.estimate - s2.estimate).abs() < 1e-12,
                "not deterministic"
            );
        }
    }

    #[test]
    fn test_pf_different_seeds() {
        let mut pf1 = PFInner::new(100, 0.0, 1.0, 1.0, 42);
        let mut pf2 = PFInner::new(100, 0.0, 1.0, 1.0, 99);

        let s1 = pf1.update(1.0);
        let s2 = pf2.update(1.0);

        // Different seeds should generally produce different results
        // (not guaranteed but extremely likely)
        assert!(
            (s1.estimate - s2.estimate).abs() > 1e-15,
            "different seeds produced identical output"
        );
    }
}
