//! Belief-Volatility Surface Engine.
//!
//! Streaming estimator that decomposes prediction market price changes into
//! diffusion and jump components using online EM, producing a "belief volatility"
//! parameter analogous to implied vol in options markets.

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

// ---------------------------------------------------------------------------
// Private math helpers (self-contained, no cross-module deps)
// ---------------------------------------------------------------------------

#[inline]
fn logit(p: f64) -> f64 {
    let p = p.clamp(1e-15, 1.0 - 1e-15);
    (p / (1.0 - p)).ln()
}

#[inline]
fn finite_or_zero(v: f64) -> f64 {
    if v.is_finite() { v } else { 0.0 }
}

/// Standard normal PDF.
#[inline]
fn normal_pdf(x: f64, mean: f64, var: f64) -> f64 {
    if var <= 0.0 { return 0.0; }
    let diff = x - mean;
    let exponent = -0.5 * diff * diff / var;
    (1.0 / (2.0 * std::f64::consts::PI * var).sqrt()) * exponent.exp()
}

// ---------------------------------------------------------------------------
// Inner state
// ---------------------------------------------------------------------------

struct Observation {
    logit_price: f64,
    _timestamp: f64,
}

struct BeliefVolInner {
    observations: VecDeque<Observation>,
    prev_logit: Option<f64>,
    ewma_logit_var: f64,
    diffusion_var: f64,
    jump_var: f64,
    jump_prob: f64,
    decay: f64,
    min_observations: usize,
    count: usize,
    max_observations: usize,
}

// ---------------------------------------------------------------------------
// BeliefVolSurface pyclass
// ---------------------------------------------------------------------------

#[pyclass]
pub struct BeliefVolSurface {
    inner: Mutex<BeliefVolInner>,
}

#[pymethods]
impl BeliefVolSurface {
    #[new]
    #[pyo3(signature = (decay=0.99, min_observations=20))]
    fn new(decay: f64, min_observations: usize) -> Self {
        Self {
            inner: Mutex::new(BeliefVolInner {
                observations: VecDeque::with_capacity(1000),
                prev_logit: None,
                ewma_logit_var: 0.01,
                diffusion_var: 0.01,
                jump_var: 0.1,
                jump_prob: 0.05,
                decay: decay.clamp(0.5, 0.9999),
                min_observations,
                count: 0,
                max_observations: 1000,
            }),
        }
    }

    /// Update with a new observation. Returns current belief vol estimate.
    #[pyo3(signature = (price, timestamp, bid=0.0, ask=0.0, volume=0.0))]
    fn update(&self, price: f64, timestamp: f64, bid: f64, ask: f64, volume: f64) -> f64 {
        let _volume = volume;
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());

        let logit_price = logit(price);

        // Store observation
        if inner.observations.len() >= inner.max_observations {
            inner.observations.pop_front();
        }
        inner.observations.push_back(Observation {
            logit_price,
            _timestamp: timestamp,
        });
        inner.count += 1;

        // Need at least one previous observation for delta
        let Some(prev_logit) = inner.prev_logit else {
            inner.prev_logit = Some(logit_price);
            return inner.diffusion_var.sqrt();
        };

        let delta_x = logit_price - prev_logit;
        inner.prev_logit = Some(logit_price);

        // Estimate microstructure noise from bid-ask spread
        let noise_var = if bid > 0.0 && ask > 0.0 && ask > bid {
            let logit_ask = logit(ask);
            let logit_bid = logit(bid);
            let spread = logit_ask - logit_bid;
            spread * spread / 12.0 // uniform noise approximation
        } else {
            0.0
        };

        // Total observed variance minus noise
        let observed_var = (delta_x * delta_x - noise_var).max(0.0);

        if inner.count < inner.min_observations {
            // Accumulate EWMA until we have enough observations
            inner.ewma_logit_var = inner.decay * inner.ewma_logit_var
                + (1.0 - inner.decay) * observed_var;
            inner.diffusion_var = inner.ewma_logit_var;
            return inner.diffusion_var.sqrt();
        }

        // Online EM step: separate diffusion and jump components
        let diff_var = inner.diffusion_var.max(1e-10);
        let jmp_var = inner.jump_var.max(1e-10);
        let jp = inner.jump_prob.clamp(0.001, 0.999);

        // E-step: posterior probability this observation is a jump
        let p_diff = (1.0 - jp) * normal_pdf(delta_x, 0.0, diff_var);
        let p_jump = jp * normal_pdf(delta_x, 0.0, diff_var + jmp_var);
        let total = p_diff + p_jump;
        let gamma_jump = if total > 1e-30 { p_jump / total } else { jp };

        // M-step: update parameters with EWMA decay
        let d = inner.decay;
        inner.diffusion_var = d * diff_var
            + (1.0 - d) * (1.0 - gamma_jump) * observed_var;
        inner.diffusion_var = inner.diffusion_var.max(1e-10);

        inner.jump_var = d * jmp_var
            + (1.0 - d) * gamma_jump * observed_var;
        inner.jump_var = inner.jump_var.max(1e-10);

        inner.jump_prob = d * jp + (1.0 - d) * gamma_jump;
        inner.jump_prob = inner.jump_prob.clamp(0.001, 0.999);

        inner.diffusion_var.sqrt()
    }

    /// Interpolated vol at a given time-to-resolution and logit moneyness.
    #[pyo3(signature = (time_to_resolution, logit_moneyness=0.0))]
    fn surface_point(&self, time_to_resolution: f64, logit_moneyness: f64) -> f64 {
        let _logit_moneyness = logit_moneyness;
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        let base_vol = inner.diffusion_var.sqrt();
        // Scale by sqrt(T/T_ref) where T_ref = 1.0
        let t = time_to_resolution.max(1e-6);
        finite_or_zero(base_vol * t.sqrt())
    }

    /// EM-estimated jump probability.
    fn jump_probability(&self) -> f64 {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        inner.jump_prob
    }

    /// Latest belief vol estimate.
    fn current_vol(&self) -> f64 {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        inner.diffusion_var.sqrt()
    }

    /// Reset all state.
    fn reset(&self) {
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        inner.observations.clear();
        inner.prev_logit = None;
        inner.ewma_logit_var = 0.01;
        inner.diffusion_var = 0.01;
        inner.jump_var = 0.1;
        inner.jump_prob = 0.05;
        inner.count = 0;
    }

    fn __repr__(&self) -> String {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        format!(
            "BeliefVolSurface(vol={:.6}, jump_prob={:.4}, n={})",
            inner.diffusion_var.sqrt(), inner.jump_prob, inner.count,
        )
    }
}

// ---------------------------------------------------------------------------
// Registration
// ---------------------------------------------------------------------------

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

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

    fn make_surface() -> BeliefVolSurface {
        BeliefVolSurface::new(0.99, 5)
    }

    #[test]
    fn initial_state() {
        let s = make_surface();
        assert!(s.current_vol() > 0.0);
        assert!(s.jump_probability() > 0.0);
    }

    #[test]
    fn converges_with_data() {
        let s = make_surface();
        // Feed stable prices → vol should stay small
        for i in 0..50 {
            let price = 0.50 + 0.001 * (i as f64 % 3.0 - 1.0);
            s.update(price, i as f64, 0.49, 0.51, 100.0);
        }
        assert!(s.current_vol() < 1.0);
    }

    #[test]
    fn jump_detection() {
        let s = make_surface();
        // Feed stable prices then a big jump
        for i in 0..30 {
            s.update(0.50, i as f64, 0.49, 0.51, 100.0);
        }
        let prob_before = s.jump_probability();
        // Big price jump
        s.update(0.80, 30.0, 0.79, 0.81, 100.0);
        let prob_after = s.jump_probability();
        assert!(prob_after >= prob_before * 0.5, "Jump prob should not collapse after a jump");
    }

    #[test]
    fn surface_point_positive() {
        let s = make_surface();
        for i in 0..10 {
            s.update(0.50 + 0.01 * i as f64, i as f64, 0.0, 0.0, 0.0);
        }
        assert!(s.surface_point(1.0, 0.0) > 0.0);
    }

    #[test]
    fn reset_works() {
        let s = make_surface();
        for i in 0..20 {
            s.update(0.50, i as f64, 0.0, 0.0, 0.0);
        }
        s.reset();
        assert!(s.current_vol() > 0.0); // back to initial
    }
}
