//! Logit-space pricing framework for prediction markets.
//!
//! Based on Dalen 2025, "Toward Black-Scholes for Prediction Markets".
//! Operates in logit space for proper boundary compression, better Greeks,
//! and a tradable "belief volatility" parameter analogous to implied vol.

use pyo3::prelude::*;

// ---------------------------------------------------------------------------
// Private math helpers
// ---------------------------------------------------------------------------

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

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

#[inline]
fn sigmoid_impl(x: f64) -> f64 {
    // Numerically stable for both branches
    if x >= 0.0 {
        let ex = (-x).exp();
        1.0 / (1.0 + ex)
    } else {
        let ex = x.exp();
        ex / (1.0 + ex)
    }
}

#[inline]
fn sigmoid_prime_impl(x: f64) -> f64 {
    let p = sigmoid_impl(x);
    p * (1.0 - p)
}

#[inline]
fn sigmoid_double_prime_impl(x: f64) -> f64 {
    let p = sigmoid_impl(x);
    p * (1.0 - p) * (1.0 - 2.0 * p)
}

// ---------------------------------------------------------------------------
// Result types
// ---------------------------------------------------------------------------

#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct LogitGreeks {
    pub delta: f64,
    pub gamma: f64,
    pub theta: f64,
    pub belief_vega: f64,
}

#[pymethods]
impl LogitGreeks {
    fn __repr__(&self) -> String {
        format!(
            "LogitGreeks(delta={:.6}, gamma={:.6}, theta={:.6}, belief_vega={:.6})",
            self.delta, self.gamma, self.theta, self.belief_vega,
        )
    }
}

#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct LogitPnlAttribution {
    pub directional: f64,
    pub curvature: f64,
    pub belief_vega_pnl: f64,
    pub correlation_vega_pnl: f64,
    pub residual: f64,
}

#[pymethods]
impl LogitPnlAttribution {
    fn __repr__(&self) -> String {
        format!(
            "LogitPnlAttribution(dir={:.6}, curv={:.6}, vega={:.6}, corr={:.6}, res={:.6})",
            self.directional, self.curvature, self.belief_vega_pnl,
            self.correlation_vega_pnl, self.residual,
        )
    }
}

// ---------------------------------------------------------------------------
// Exposed functions
// ---------------------------------------------------------------------------

fn logit_fn_impl(p: f64) -> f64 {
    logit_impl(p)
}

#[pyfunction]
fn logit(p: f64) -> f64 {
    logit_fn_impl(p)
}

fn sigmoid_fn_impl(x: f64) -> f64 {
    sigmoid_impl(x)
}

#[pyfunction]
fn sigmoid(x: f64) -> f64 {
    sigmoid_fn_impl(x)
}

#[pyfunction]
fn sigmoid_prime(x: f64) -> f64 {
    sigmoid_prime_impl(x)
}

#[pyfunction]
fn sigmoid_double_prime(x: f64) -> f64 {
    sigmoid_double_prime_impl(x)
}

fn logit_reservation_price_impl(
    mid: f64, inventory: f64, gamma: f64, belief_vol: f64, t: f64,
) -> f64 {
    if !mid.is_finite() || !inventory.is_finite() || !gamma.is_finite()
        || !belief_vol.is_finite() || !t.is_finite()
    {
        return mid.clamp(0.01, 0.99);
    }
    let mid_c = mid.clamp(1e-15, 1.0 - 1e-15);
    let logit_mid = logit_impl(mid_c);
    let adjustment = inventory * gamma * belief_vol * belief_vol * t;
    let logit_res = logit_mid - adjustment;
    sigmoid_impl(logit_res).clamp(0.01, 0.99)
}

#[pyfunction]
#[pyo3(signature = (mid, inventory, gamma, belief_vol, t))]
fn logit_reservation_price(mid: f64, inventory: f64, gamma: f64, belief_vol: f64, t: f64) -> f64 {
    logit_reservation_price_impl(mid, inventory, gamma, belief_vol, t)
}

fn logit_optimal_spread_impl(
    belief_vol: f64, inventory: f64, gamma: f64, kappa: f64, t: f64,
) -> f64 {
    if !belief_vol.is_finite() || !inventory.is_finite() || !gamma.is_finite()
        || !kappa.is_finite() || !t.is_finite() || gamma <= 0.0 || kappa <= 0.0
    {
        return 0.0;
    }
    // Logit-space spread = γ*σ_b²*T + (2/γ)*ln(1 + γ/κ)
    let variance_component = gamma * belief_vol * belief_vol * t;
    let arrival_component = (2.0 / gamma) * (1.0 + gamma / kappa).ln();
    let logit_spread = variance_component + arrival_component;
    // Convert to probability space: multiply by S'(x) = p(1-p)
    // Use mid=0.5 as reference point for conversion
    let p = 0.5_f64; // neutral reference
    let prob_spread = logit_spread * p * (1.0 - p);
    finite_or_zero(prob_spread.clamp(0.0, 1.0))
}

#[pyfunction]
#[pyo3(signature = (belief_vol, inventory, gamma, kappa, t))]
fn logit_optimal_spread(belief_vol: f64, inventory: f64, gamma: f64, kappa: f64, t: f64) -> f64 {
    logit_optimal_spread_impl(belief_vol, inventory, gamma, kappa, t)
}

fn logit_greeks_impl(
    price: f64, size: f64, is_yes: bool, t_hours: f64, belief_vol: f64,
) -> LogitGreeks {
    let p = price.clamp(1e-15, 1.0 - 1e-15);
    let sign = if is_yes { 1.0 } else { -1.0 };
    let t = t_hours.max(1e-6);
    let sigma = belief_vol.max(1e-10);
    let sp = p * (1.0 - p); // sigmoid_prime

    let delta = finite_or_zero(sign * size * sp);
    let gamma = finite_or_zero(sign * size * sp * (1.0 - 2.0 * p));
    let theta = finite_or_zero(-sign * size * sp * sigma * sigma / (2.0 * t));
    let belief_vega = finite_or_zero(sign * size * sp * 2.0 * sigma * t);

    LogitGreeks { delta, gamma, theta, belief_vega }
}

#[pyfunction]
#[pyo3(signature = (price, size, is_yes, t_hours=24.0, belief_vol=0.2))]
fn logit_greeks(price: f64, size: f64, is_yes: bool, t_hours: f64, belief_vol: f64) -> LogitGreeks {
    logit_greeks_impl(price, size, is_yes, t_hours, belief_vol)
}

fn logit_pnl_decompose_impl(
    p_before: f64, p_after: f64, size: f64, is_yes: bool,
    vol_before: f64, vol_after: f64, corr_before: f64, corr_after: f64,
) -> LogitPnlAttribution {
    let sign = if is_yes { 1.0 } else { -1.0 };
    let pb = p_before.clamp(1e-15, 1.0 - 1e-15);
    let pa = p_after.clamp(1e-15, 1.0 - 1e-15);

    // Actual PnL
    let actual_pnl = sign * size * (pa - pb);

    // Logit-space move
    let x_before = logit_impl(pb);
    let x_after = logit_impl(pa);
    let dx = x_after - x_before;

    // Greeks at entry point
    let sp = pb * (1.0 - pb);
    let delta = sign * size * sp;
    let gamma = sign * size * sp * (1.0 - 2.0 * pb);

    // Vega: d(value)/d(sigma) ~ sign * size * sp * 2*sigma*T (use T=1 for decomposition)
    let vega = sign * size * sp * 2.0 * vol_before;
    // Correlation vega: d(value)/d(rho) ~ sign * size * sp
    let corr_vega = sign * size * sp;

    let directional = finite_or_zero(delta * dx);
    let curvature = finite_or_zero(0.5 * gamma * dx * dx);
    let belief_vega_pnl = finite_or_zero(vega * (vol_after - vol_before));
    let correlation_vega_pnl = finite_or_zero(corr_vega * (corr_after - corr_before));
    let residual = finite_or_zero(actual_pnl - directional - curvature - belief_vega_pnl - correlation_vega_pnl);

    LogitPnlAttribution { directional, curvature, belief_vega_pnl, correlation_vega_pnl, residual }
}

#[pyfunction]
#[pyo3(signature = (p_before, p_after, size, is_yes, vol_before, vol_after, corr_before=0.0, corr_after=0.0))]
fn logit_pnl_decompose(
    p_before: f64, p_after: f64, size: f64, is_yes: bool,
    vol_before: f64, vol_after: f64, corr_before: f64, corr_after: f64,
) -> LogitPnlAttribution {
    logit_pnl_decompose_impl(p_before, p_after, size, is_yes, vol_before, vol_after, corr_before, corr_after)
}

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

pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<LogitGreeks>()?;
    m.add_class::<LogitPnlAttribution>()?;
    m.add_function(wrap_pyfunction!(logit, m)?)?;
    m.add_function(wrap_pyfunction!(sigmoid, m)?)?;
    m.add_function(wrap_pyfunction!(sigmoid_prime, m)?)?;
    m.add_function(wrap_pyfunction!(sigmoid_double_prime, m)?)?;
    m.add_function(wrap_pyfunction!(logit_reservation_price, m)?)?;
    m.add_function(wrap_pyfunction!(logit_optimal_spread, m)?)?;
    m.add_function(wrap_pyfunction!(logit_greeks, m)?)?;
    m.add_function(wrap_pyfunction!(logit_pnl_decompose, m)?)?;
    Ok(())
}

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

    #[test]
    fn logit_sigmoid_inverse() {
        for &p in &[0.1, 0.25, 0.5, 0.75, 0.9] {
            let x = logit_impl(p);
            let recovered = sigmoid_impl(x);
            assert!((recovered - p).abs() < 1e-12, "p={p}, recovered={recovered}");
        }
    }

    #[test]
    fn logit_known_values() {
        assert!((logit_impl(0.5) - 0.0).abs() < 1e-12);
        assert!((sigmoid_impl(0.0) - 0.5).abs() < 1e-12);
    }

    #[test]
    fn sigmoid_prime_at_half() {
        // S'(0) = 0.5 * 0.5 = 0.25
        assert!((sigmoid_prime_impl(0.0) - 0.25).abs() < 1e-12);
    }

    #[test]
    fn reservation_price_no_inventory() {
        let r = logit_reservation_price_impl(0.5, 0.0, 0.5, 0.2, 1.0);
        assert!((r - 0.5).abs() < 0.01);
    }

    #[test]
    fn reservation_price_long_skews_down() {
        let r = logit_reservation_price_impl(0.5, 10.0, 0.5, 0.2, 1.0);
        assert!(r < 0.5, "Long inventory should skew reservation price down");
    }

    #[test]
    fn greeks_delta_sign() {
        let yes = logit_greeks_impl(0.5, 10.0, true, 24.0, 0.2);
        let no = logit_greeks_impl(0.5, 10.0, false, 24.0, 0.2);
        assert!(yes.delta > 0.0);
        assert!(no.delta < 0.0);
    }

    #[test]
    fn greeks_gamma_maximal_at_half() {
        // At p=0.5, (1-2p)=0, so gamma should be zero
        let g = logit_greeks_impl(0.5, 10.0, true, 24.0, 0.2);
        assert!(g.gamma.abs() < 1e-10);
    }

    #[test]
    fn pnl_decompose_components_sum() {
        let attr = logit_pnl_decompose_impl(0.4, 0.6, 10.0, true, 0.2, 0.25, 0.0, 0.0);
        let total = attr.directional + attr.curvature + attr.belief_vega_pnl
            + attr.correlation_vega_pnl + attr.residual;
        let actual = 10.0 * (0.6 - 0.4); // yes side
        assert!((total - actual).abs() < 1e-10, "Components should sum to actual PnL");
    }
}
