//! Automated market making math for prediction markets.
//!
//! Implements Avellaneda-Stoikov style inventory-skewed quoting.
//! All functions are `#[inline]` for zero-overhead inlining at call sites.

use pyo3::prelude::*;

/// Avellaneda-Stoikov reservation price: inventory-skewed fair value.
///
/// Formula: r = mid - inventory * gamma * volatility^2 * time_horizon
///
/// - `mid`: current mid price
/// - `inventory`: net inventory (positive = long, negative = short)
/// - `gamma`: risk aversion parameter (higher = more aggressive skew)
/// - `volatility`: estimated price volatility (std dev of returns)
/// - `time_horizon`: time horizon in arbitrary units (e.g., 1.0 = full session)
///
/// Returns mid price if any input is NaN/Inf.
#[pyfunction]
#[inline]
pub fn reservation_price(
    mid: f64,
    inventory: f64,
    gamma: f64,
    volatility: f64,
    time_horizon: f64,
) -> f64 {
    if !mid.is_finite()
        || !inventory.is_finite()
        || !gamma.is_finite()
        || !volatility.is_finite()
        || !time_horizon.is_finite()
    {
        return if mid.is_finite() { mid } else { 0.0 };
    }
    mid - inventory * gamma * volatility * volatility * time_horizon
}

/// Optimal bid-ask spread width (Avellaneda-Stoikov).
///
/// Formula: spread = gamma * sigma^2 * T + (2/gamma) * ln(1 + gamma/kappa)
///
/// - `volatility`: estimated price volatility
/// - `inventory`: net inventory (used indirectly via gamma)
/// - `gamma`: risk aversion parameter
/// - `kappa`: order arrival intensity (higher = more frequent fills)
/// - `time_horizon`: time horizon
///
/// Returns 0.0 if any input is NaN/Inf or spread would be negative.
#[pyfunction]
#[inline]
pub fn optimal_spread(
    volatility: f64,
    inventory: f64,
    gamma: f64,
    kappa: f64,
    time_horizon: f64,
) -> f64 {
    if !volatility.is_finite()
        || !inventory.is_finite()
        || !gamma.is_finite()
        || !kappa.is_finite()
        || !time_horizon.is_finite()
        || gamma <= 0.0
        || kappa <= 0.0
    {
        return 0.0;
    }
    let variance_component = gamma * volatility * volatility * time_horizon;
    let arrival_component = (2.0 / gamma) * (1.0 + gamma / kappa).ln();
    let spread = variance_component + arrival_component;
    spread.max(0.0)
}

/// Competitive spread: blend model spread with observed market spread.
///
/// - `base_spread`: model-derived optimal spread
/// - `book_spread`: observed bid-ask spread from orderbook
/// - `book_imbalance`: orderbook imbalance in [-1, 1] (>0 = buy pressure)
/// - `aggression`: 0.0 = use model spread, 1.0 = use book spread
///
/// Returns clamped to [0.001, 1.0].
#[pyfunction]
#[inline]
pub fn competitive_spread(
    base_spread: f64,
    book_spread: f64,
    book_imbalance: f64,
    aggression: f64,
) -> f64 {
    if !base_spread.is_finite()
        || !book_spread.is_finite()
        || !book_imbalance.is_finite()
        || !aggression.is_finite()
    {
        return if base_spread.is_finite() {
            base_spread.clamp(0.001, 1.0)
        } else {
            0.04
        };
    }
    let agg = aggression.clamp(0.0, 1.0);
    let imb = book_imbalance.clamp(-1.0, 1.0);

    // Blend model and market spread
    let blended = (1.0 - agg) * base_spread + agg * book_spread;

    // Imbalance adjustment: tighten on the strong side, widen on the weak side
    let adjustment = 1.0 - imb.abs() * 0.1;
    let spread = blended * adjustment;

    spread.clamp(0.001, 1.0)
}

/// Inventory-skewed bid/ask sizes.
///
/// Returns (bid_size, ask_size) adjusted for inventory.
/// When long, reduce bid size and increase ask size to reduce inventory.
///
/// - `base_size`: starting size for both sides
/// - `inventory`: net inventory (positive = long)
/// - `max_position`: maximum allowed position
/// - `fill_rate`: estimated fill rate (0-1, higher = more aggressive sizing)
/// - `skew_factor`: how aggressively to skew sizes (0 = no skew, 1 = max skew)
#[pyfunction]
#[inline]
pub fn mm_size(
    base_size: f64,
    inventory: f64,
    max_position: f64,
    fill_rate: f64,
    skew_factor: f64,
) -> (f64, f64) {
    if !base_size.is_finite()
        || !inventory.is_finite()
        || !max_position.is_finite()
        || !fill_rate.is_finite()
        || !skew_factor.is_finite()
        || base_size <= 0.0
        || max_position <= 0.0
    {
        return (0.0, 0.0);
    }

    let skew = skew_factor.clamp(0.0, 1.0);
    let fill = fill_rate.clamp(0.0, 1.0);

    // Inventory ratio: -1 (max short) to +1 (max long)
    let inv_ratio = (inventory / max_position).clamp(-1.0, 1.0);

    // Skew factor: reduce size on the side that would increase inventory
    let bid_skew = 1.0 - skew * inv_ratio.max(0.0);
    let ask_skew = 1.0 + skew * inv_ratio.min(0.0);

    // Fill rate adjustment
    let fill_adj = 1.0 + fill * 0.5;

    let bid_size = (base_size * bid_skew * fill_adj).max(0.0);
    let ask_size = (base_size * ask_skew * fill_adj).max(0.0);

    (bid_size, ask_size)
}

/// Estimate volatility from a price series using log-return standard deviation.
///
/// Returns 0.0 if fewer than 2 prices or all prices are the same.
#[pyfunction]
#[inline]
pub fn estimate_volatility(prices: Vec<f64>) -> f64 {
    if prices.len() < 2 {
        return 0.0;
    }

    // Compute log returns, filtering out invalid values
    let mut returns = Vec::with_capacity(prices.len() - 1);
    for i in 1..prices.len() {
        let prev = prices[i - 1];
        let curr = prices[i];
        if prev > 0.0 && curr > 0.0 && prev.is_finite() && curr.is_finite() {
            returns.push((curr / prev).ln());
        }
    }

    if returns.is_empty() {
        return 0.0;
    }

    let n = returns.len() as f64;
    let mean = returns.iter().sum::<f64>() / n;
    let variance = returns.iter().map(|r| (r - mean).powi(2)).sum::<f64>() / n;

    variance.sqrt()
}

/// Register all market making functions in the Python module.
pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(reservation_price, m)?)?;
    m.add_function(wrap_pyfunction!(optimal_spread, m)?)?;
    m.add_function(wrap_pyfunction!(competitive_spread, m)?)?;
    m.add_function(wrap_pyfunction!(mm_size, m)?)?;
    m.add_function(wrap_pyfunction!(estimate_volatility, m)?)?;
    Ok(())
}

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

    #[test]
    fn reservation_price_no_inventory() {
        let r = reservation_price(0.50, 0.0, 0.5, 0.02, 1.0);
        assert!((r - 0.50).abs() < 1e-10);
    }

    #[test]
    fn reservation_price_long_inventory_skews_down() {
        let r = reservation_price(0.50, 10.0, 0.5, 0.02, 1.0);
        assert!(r < 0.50);
    }

    #[test]
    fn reservation_price_short_inventory_skews_up() {
        let r = reservation_price(0.50, -10.0, 0.5, 0.02, 1.0);
        assert!(r > 0.50);
    }

    #[test]
    fn reservation_price_nan_returns_mid() {
        let r = reservation_price(0.50, f64::NAN, 0.5, 0.02, 1.0);
        assert!((r - 0.50).abs() < 1e-10);
    }

    #[test]
    fn optimal_spread_positive() {
        let s = optimal_spread(0.02, 10.0, 0.5, 1.5, 1.0);
        assert!(s > 0.0);
    }

    #[test]
    fn optimal_spread_zero_gamma() {
        assert_eq!(optimal_spread(0.02, 10.0, 0.0, 1.5, 1.0), 0.0);
    }

    #[test]
    fn optimal_spread_zero_kappa() {
        assert_eq!(optimal_spread(0.02, 10.0, 0.5, 0.0, 1.0), 0.0);
    }

    #[test]
    fn optimal_spread_nan() {
        assert_eq!(optimal_spread(f64::NAN, 10.0, 0.5, 1.5, 1.0), 0.0);
    }

    #[test]
    fn competitive_spread_pure_model() {
        let s = competitive_spread(0.04, 0.02, 0.0, 0.0);
        assert!((s - 0.04).abs() < 1e-10);
    }

    #[test]
    fn competitive_spread_pure_book() {
        let s = competitive_spread(0.04, 0.02, 0.0, 1.0);
        assert!((s - 0.02).abs() < 1e-10);
    }

    #[test]
    fn competitive_spread_clamped() {
        let s = competitive_spread(0.0001, 0.0001, 0.0, 0.5);
        assert!(s >= 0.001);
    }

    #[test]
    fn mm_size_no_inventory() {
        let (bid, ask) = mm_size(5.0, 0.0, 100.0, 0.3, 0.5);
        // With zero inventory, both sides should be equal
        assert!((bid - ask).abs() < 1e-10);
        assert!(bid > 0.0);
    }

    #[test]
    fn mm_size_long_inventory_skews() {
        let (bid, ask) = mm_size(5.0, 50.0, 100.0, 0.3, 0.5);
        // Long inventory → bid_size < ask_size (encourage selling)
        assert!(bid < ask);
    }

    #[test]
    fn mm_size_short_inventory_skews() {
        let (bid, ask) = mm_size(5.0, -50.0, 100.0, 0.3, 0.5);
        // Short inventory → bid_size > ask_size (encourage buying)
        assert!(bid > ask);
    }

    #[test]
    fn mm_size_zero_base() {
        let (bid, ask) = mm_size(0.0, 10.0, 100.0, 0.3, 0.5);
        assert_eq!(bid, 0.0);
        assert_eq!(ask, 0.0);
    }

    #[test]
    fn mm_size_nan() {
        let (bid, ask) = mm_size(f64::NAN, 10.0, 100.0, 0.3, 0.5);
        assert_eq!(bid, 0.0);
        assert_eq!(ask, 0.0);
    }

    #[test]
    fn estimate_volatility_constant() {
        assert_eq!(estimate_volatility(vec![1.0, 1.0, 1.0, 1.0]), 0.0);
    }

    #[test]
    fn estimate_volatility_positive() {
        let vol = estimate_volatility(vec![1.0, 1.01, 0.99, 1.02, 0.98]);
        assert!(vol > 0.0);
    }

    #[test]
    fn estimate_volatility_too_few() {
        assert_eq!(estimate_volatility(vec![1.0]), 0.0);
        assert_eq!(estimate_volatility(vec![]), 0.0);
    }

    #[test]
    fn estimate_volatility_filters_invalid() {
        let vol = estimate_volatility(vec![1.0, 0.0, 1.0]); // 0.0 makes ln undefined
        // Should skip the 0.0 → 1.0 pair and the 1.0 → 0.0 pair
        assert_eq!(vol, 0.0); // no valid returns
    }

    #[test]
    fn reservation_price_formula() {
        // r = 0.50 - 10 * 0.5 * 0.02^2 * 1.0 = 0.50 - 0.002 = 0.498
        let r = reservation_price(0.50, 10.0, 0.5, 0.02, 1.0);
        assert!((r - 0.498).abs() < 1e-10);
    }

    #[test]
    fn optimal_spread_formula() {
        // gamma=0.5, sigma=0.02, T=1.0, kappa=1.5
        // variance_component = 0.5 * 0.0004 * 1.0 = 0.0002
        // arrival_component = (2/0.5) * ln(1 + 0.5/1.5) = 4 * ln(1.333..) ≈ 4 * 0.28768 ≈ 1.15073
        let s = optimal_spread(0.02, 10.0, 0.5, 1.5, 1.0);
        let expected = 0.0002 + 4.0 * (1.0 + 0.5 / 1.5_f64).ln();
        assert!((s - expected).abs() < 1e-6);
    }
}
