//! Breeden-Litzenberger cross-asset probability extraction.
//!
//! Extracts risk-neutral probability densities from options prices using
//! the second derivative of call prices with respect to strike, then maps
//! these to prediction market edge signals.

use pyo3::prelude::*;

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

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

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

#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct RiskNeutralDensity {
    pub strikes: Vec<f64>,
    pub densities: Vec<f64>,
    pub cdf: Vec<f64>,
    pub mean: f64,
    pub variance: f64,
}

#[pymethods]
impl RiskNeutralDensity {
    fn __repr__(&self) -> String {
        format!(
            "RiskNeutralDensity(n={}, mean={:.4}, var={:.4})",
            self.strikes.len(), self.mean, self.variance,
        )
    }
}

#[pyclass(frozen, get_all)]
#[derive(Debug, Clone)]
pub struct CrossAssetEdge {
    pub option_implied_prob: f64,
    pub market_price: f64,
    pub raw_edge: f64,
    pub adjusted_edge: f64,
    pub confidence: f64,
}

#[pymethods]
impl CrossAssetEdge {
    fn __repr__(&self) -> String {
        format!(
            "CrossAssetEdge(implied={:.4}, market={:.4}, edge={:.4}, conf={:.4})",
            self.option_implied_prob, self.market_price, self.raw_edge, self.confidence,
        )
    }
}

// ---------------------------------------------------------------------------
// Core density extraction
// ---------------------------------------------------------------------------

fn risk_neutral_density_impl(
    strikes: &[f64], call_prices: &[f64], r: f64, t: f64,
) -> RiskNeutralDensity {
    let n = strikes.len().min(call_prices.len());
    if n < 3 {
        return RiskNeutralDensity {
            strikes: strikes.to_vec(),
            densities: vec![0.0; n],
            cdf: vec![0.0; n],
            mean: 0.0,
            variance: 0.0,
        };
    }

    let df = (-r * t).exp(); // discount factor

    // Second derivative via finite differences (interior points)
    let mut raw_density = vec![0.0; n];
    for i in 1..n - 1 {
        let dk = strikes[i + 1] - strikes[i - 1];
        if dk > 1e-15 {
            let d2c = call_prices[i + 1] - 2.0 * call_prices[i] + call_prices[i - 1];
            let dk_sq = ((strikes[i + 1] - strikes[i]) * (strikes[i] - strikes[i - 1])).max(1e-15);
            raw_density[i] = (df * d2c / dk_sq).max(0.0);
        }
    }

    // Normalize: trapezoidal rule to integrate to 1.0
    let mut total = 0.0;
    for i in 0..n - 1 {
        let dk = strikes[i + 1] - strikes[i];
        total += 0.5 * (raw_density[i] + raw_density[i + 1]) * dk;
    }

    let mut densities = vec![0.0; n];
    if total > 1e-15 {
        for i in 0..n {
            densities[i] = finite_or_zero(raw_density[i] / total);
        }
    }

    // CDF via cumulative trapezoidal
    let mut cdf = vec![0.0; n];
    for i in 1..n {
        let dk = strikes[i] - strikes[i - 1];
        cdf[i] = cdf[i - 1] + 0.5 * (densities[i] + densities[i - 1]) * dk;
    }
    // Ensure CDF ends at 1.0
    if let Some(last) = cdf.last_mut() {
        if *last > 0.0 {
            let scale = 1.0 / *last;
            for c in cdf.iter_mut() {
                *c *= scale;
            }
        }
    }

    // Mean and variance
    let mut mean_val = 0.0;
    let mut second_moment = 0.0;
    for i in 0..n - 1 {
        let dk = strikes[i + 1] - strikes[i];
        let mid_k = 0.5 * (strikes[i] + strikes[i + 1]);
        let mid_d = 0.5 * (densities[i] + densities[i + 1]);
        mean_val += mid_k * mid_d * dk;
        second_moment += mid_k * mid_k * mid_d * dk;
    }
    let variance = (second_moment - mean_val * mean_val).max(0.0);

    RiskNeutralDensity {
        strikes: strikes[..n].to_vec(),
        densities,
        cdf,
        mean: finite_or_zero(mean_val),
        variance: finite_or_zero(variance),
    }
}

#[pyfunction]
#[pyo3(signature = (strikes, call_prices, r=0.0, t=1.0))]
fn risk_neutral_density(strikes: Vec<f64>, call_prices: Vec<f64>, r: f64, t: f64) -> RiskNeutralDensity {
    risk_neutral_density_impl(&strikes, &call_prices, r, t)
}

fn implied_probability_above_impl(
    strikes: &[f64], call_prices: &[f64], threshold: f64, r: f64, t: f64,
) -> f64 {
    let rnd = risk_neutral_density_impl(strikes, call_prices, r, t);
    let n = rnd.strikes.len();
    if n < 2 { return 0.0; }

    // Integrate density above threshold
    let mut prob = 0.0;
    for i in 0..n - 1 {
        let k0 = rnd.strikes[i];
        let k1 = rnd.strikes[i + 1];
        if k1 <= threshold { continue; }
        let effective_k0 = k0.max(threshold);
        let dk = k1 - effective_k0;
        if dk <= 0.0 { continue; }

        // Linear interpolation of density at effective_k0
        let frac = if (k1 - k0).abs() > 1e-15 { (effective_k0 - k0) / (k1 - k0) } else { 0.0 };
        let d0 = rnd.densities[i] + frac * (rnd.densities[i + 1] - rnd.densities[i]);
        let d1 = rnd.densities[i + 1];
        prob += 0.5 * (d0 + d1) * dk;
    }
    finite_or_zero(prob.clamp(0.0, 1.0))
}

#[pyfunction]
#[pyo3(signature = (strikes, call_prices, threshold, r=0.0, t=1.0))]
fn implied_probability_above(strikes: Vec<f64>, call_prices: Vec<f64>, threshold: f64, r: f64, t: f64) -> f64 {
    implied_probability_above_impl(&strikes, &call_prices, threshold, r, t)
}

#[pyfunction]
#[pyo3(signature = (strikes, call_prices, threshold, r=0.0, t=1.0))]
fn implied_probability_below(strikes: Vec<f64>, call_prices: Vec<f64>, threshold: f64, r: f64, t: f64) -> f64 {
    1.0 - implied_probability_above_impl(&strikes, &call_prices, threshold, r, t)
}

fn implied_probability_between_impl(
    strikes: &[f64], call_prices: &[f64], lower: f64, upper: f64, r: f64, t: f64,
) -> f64 {
    let above_lower = implied_probability_above_impl(strikes, call_prices, lower, r, t);
    let above_upper = implied_probability_above_impl(strikes, call_prices, upper, r, t);
    finite_or_zero((above_lower - above_upper).clamp(0.0, 1.0))
}

#[pyfunction]
#[pyo3(signature = (strikes, call_prices, lower, upper, r=0.0, t=1.0))]
fn implied_probability_between(strikes: Vec<f64>, call_prices: Vec<f64>, lower: f64, upper: f64, r: f64, t: f64) -> f64 {
    implied_probability_between_impl(&strikes, &call_prices, lower, upper, r, t)
}

fn cross_asset_edge_impl(
    option_prob: f64, market_price: f64, transaction_cost: f64,
) -> CrossAssetEdge {
    let op = option_prob.clamp(0.0, 1.0);
    let mp = market_price.clamp(0.0, 1.0);
    let tc = transaction_cost.clamp(0.0, 1.0);

    let raw_edge = op - mp;
    let adjusted_edge = if raw_edge > 0.0 {
        (raw_edge - tc).max(0.0)
    } else {
        (raw_edge + tc).min(0.0)
    };

    // Confidence: based on magnitude of edge relative to transaction cost
    let confidence = if tc > 1e-10 {
        (raw_edge.abs() / tc).min(1.0)
    } else {
        if raw_edge.abs() > 1e-10 { 1.0 } else { 0.0 }
    };

    CrossAssetEdge {
        option_implied_prob: op,
        market_price: mp,
        raw_edge: finite_or_zero(raw_edge),
        adjusted_edge: finite_or_zero(adjusted_edge),
        confidence: finite_or_zero(confidence),
    }
}

#[pyfunction]
#[pyo3(signature = (option_prob, market_price, transaction_cost=0.0))]
fn cross_asset_edge(option_prob: f64, market_price: f64, transaction_cost: f64) -> CrossAssetEdge {
    cross_asset_edge_impl(option_prob, market_price, transaction_cost)
}

fn risk_premium_adjustment_impl(
    option_prob: f64, market_price: f64, historical_ratio: f64,
) -> f64 {
    let op = option_prob.clamp(0.0, 1.0);
    let mp = market_price.clamp(0.0, 1.0);
    let hr = historical_ratio.clamp(0.0, 2.0);
    // Blend: adjusted_prob = option_prob * historical_ratio + market_price * (1 - historical_ratio)
    // When historical_ratio = 1.0, use options entirely
    // When historical_ratio = 0.0, use market entirely
    let adjusted = op * hr + mp * (1.0 - hr);
    finite_or_zero(adjusted.clamp(0.0, 1.0))
}

#[pyfunction]
#[pyo3(signature = (option_prob, market_price, historical_ratio=0.5))]
fn risk_premium_adjustment(option_prob: f64, market_price: f64, historical_ratio: f64) -> f64 {
    risk_premium_adjustment_impl(option_prob, market_price, historical_ratio)
}

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

pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<RiskNeutralDensity>()?;
    m.add_class::<CrossAssetEdge>()?;
    m.add_function(wrap_pyfunction!(risk_neutral_density, m)?)?;
    m.add_function(wrap_pyfunction!(implied_probability_above, m)?)?;
    m.add_function(wrap_pyfunction!(implied_probability_below, m)?)?;
    m.add_function(wrap_pyfunction!(implied_probability_between, m)?)?;
    m.add_function(wrap_pyfunction!(cross_asset_edge, m)?)?;
    m.add_function(wrap_pyfunction!(risk_premium_adjustment, m)?)?;
    Ok(())
}

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

    fn bs_call_price(s: f64, k: f64, r: f64, t: f64, sigma: f64) -> f64 {
        // Simplified Black-Scholes call
        let d1 = ((s / k).ln() + (r + 0.5 * sigma * sigma) * t) / (sigma * t.sqrt());
        let d2 = d1 - sigma * t.sqrt();
        s * normal_cdf(d1) - k * (-r * t).exp() * normal_cdf(d2)
    }

    fn normal_cdf(x: f64) -> f64 {
        const A1: f64 = 0.254829592;
        const A2: f64 = -0.284496736;
        const A3: f64 = 1.421413741;
        const A4: f64 = -1.453152027;
        const A5: f64 = 1.061405429;
        const P: f64 = 0.3275911;
        let sign = if x < 0.0 { -1.0 } else { 1.0 };
        let x_abs = x.abs() / std::f64::consts::SQRT_2;
        let t = 1.0 / (1.0 + P * x_abs);
        let y = 1.0 - (((((A5 * t + A4) * t) + A3) * t + A2) * t + A1) * t * (-x_abs * x_abs).exp();
        0.5 * (1.0 + sign * y)
    }

    #[test]
    fn density_from_bs_integrates_to_one() {
        let s = 100.0;
        let r = 0.05;
        let t = 1.0;
        let sigma = 0.2;
        let strikes: Vec<f64> = (50..=150).map(|k| k as f64).collect();
        let calls: Vec<f64> = strikes.iter().map(|&k| bs_call_price(s, k, r, t, sigma)).collect();

        let rnd = risk_neutral_density_impl(&strikes, &calls, r, t);
        // CDF should end at approximately 1.0
        let last_cdf = rnd.cdf.last().copied().unwrap_or(0.0);
        assert!((last_cdf - 1.0).abs() < 0.1, "CDF should end near 1.0, got {last_cdf}");
        assert!(rnd.mean > 50.0 && rnd.mean < 150.0);
    }

    #[test]
    fn above_below_sum_to_one() {
        let s = 100.0;
        let r = 0.0;
        let t = 1.0;
        let sigma = 0.3;
        let strikes: Vec<f64> = (60..=140).map(|k| k as f64).collect();
        let calls: Vec<f64> = strikes.iter().map(|&k| bs_call_price(s, k, r, t, sigma)).collect();

        let above = implied_probability_above_impl(&strikes, &calls, 100.0, r, t);
        let below = 1.0 - above;
        assert!((above + below - 1.0).abs() < 1e-10);
    }

    #[test]
    fn between_leq_above() {
        let s = 100.0;
        let r = 0.0;
        let t = 1.0;
        let sigma = 0.3;
        let strikes: Vec<f64> = (60..=140).map(|k| k as f64).collect();
        let calls: Vec<f64> = strikes.iter().map(|&k| bs_call_price(s, k, r, t, sigma)).collect();

        let above = implied_probability_above_impl(&strikes, &calls, 90.0, r, t);
        let between = implied_probability_between_impl(&strikes, &calls, 90.0, 110.0, r, t);
        assert!(between <= above + 1e-10);
    }

    #[test]
    fn cross_asset_edge_positive() {
        let e = cross_asset_edge_impl(0.7, 0.5, 0.01);
        assert!(e.raw_edge > 0.0);
        assert!(e.adjusted_edge > 0.0);
        assert!(e.confidence > 0.0);
    }

    #[test]
    fn cross_asset_edge_negative() {
        let e = cross_asset_edge_impl(0.3, 0.5, 0.01);
        assert!(e.raw_edge < 0.0);
        assert!(e.adjusted_edge < 0.0);
    }

    #[test]
    fn risk_premium_blend() {
        let r = risk_premium_adjustment_impl(0.8, 0.5, 0.5);
        assert!((r - 0.65).abs() < 1e-10); // 0.8*0.5 + 0.5*0.5 = 0.65
    }
}
