//! Kelly criterion sizing for binary prediction markets.
//!
//! All functions are `#[inline]` for zero-overhead inlining at call sites.
//! No heap allocations on the hot path — pure register-level arithmetic.

use pyo3::prelude::*;

/// Kelly fraction for buying YES at `market_price` when true probability is `prob`.
///
/// Formula: f* = (p - P) / (1 - P)
/// Returns 0.0 if no edge (prob <= price) or invalid inputs.
/// Result is the fraction of bankroll to risk.
#[pyfunction]
#[inline]
pub fn kelly(prob: f64, market_price: f64) -> f64 {
    // Guard: NaN/Inf inputs, price must be in (0, 1), prob must be in [0, 1]
    if prob.is_nan() || market_price.is_nan()
        || prob.is_infinite() || market_price.is_infinite()
        || prob < 0.0 || prob > 1.0
        || market_price <= 0.0 || market_price >= 1.0
        || prob <= market_price
    {
        return 0.0;
    }
    (prob - market_price) / (1.0 - market_price)
}

/// Kelly fraction for the NO side — buying NO at `1 - market_price`.
///
/// Formula: f* = (P - p) / P
/// Use when you believe the market is overpriced (prob < market_price).
#[pyfunction]
#[inline]
pub fn kelly_no(prob: f64, market_price: f64) -> f64 {
    if prob.is_nan() || market_price.is_nan()
        || prob.is_infinite() || market_price.is_infinite()
        || prob < 0.0 || prob > 1.0
        || market_price <= 0.0 || market_price >= 1.0
        || prob >= market_price
    {
        return 0.0;
    }
    (market_price - prob) / market_price
}

/// Fractional Kelly — multiplies raw Kelly by a fraction (typically 0.25–0.5)
/// to reduce variance at the cost of slightly lower expected growth.
#[pyfunction]
#[inline]
pub fn fractional_kelly(prob: f64, market_price: f64, fraction: f64) -> f64 {
    kelly(prob, market_price) * fraction.clamp(0.0, 1.0)
}

/// Compute the actual contract size from Kelly fraction.
///
/// `bankroll`: total capital available.
/// `fraction`: Kelly scaling factor (1.0 = full Kelly, 0.5 = half Kelly).
/// `max_size`: hard cap on position size (from risk config).
///
/// Returns number of contracts to buy at `market_price`.
#[pyfunction]
#[inline]
pub fn kelly_size(
    prob: f64,
    market_price: f64,
    bankroll: f64,
    fraction: f64,
    max_size: f64,
) -> f64 {
    if bankroll <= 0.0 || market_price <= 0.0 {
        return 0.0;
    }
    let f = fractional_kelly(prob, market_price, fraction);
    let risk_amount = f * bankroll;
    // contracts = dollars_to_risk / price_per_contract
    let contracts = risk_amount / market_price;
    contracts.min(max_size).max(0.0)
}

/// Independent multi-outcome Kelly: compute Kelly fractions for N binary outcomes.
///
/// Each outcome is treated independently. If total allocation exceeds 1.0,
/// fractions are proportionally scaled down to sum to at most `max_total`.
///
/// Returns a Vec of Kelly fractions (one per outcome).
#[pyfunction]
pub fn multi_kelly(probs: Vec<f64>, prices: Vec<f64>, max_total: f64) -> Vec<f64> {
    let n = probs.len().min(prices.len());
    let mut fractions = Vec::with_capacity(n);
    let mut total = 0.0f64;

    for i in 0..n {
        let f = kelly(probs[i], prices[i]);
        fractions.push(f);
        total += f;
    }

    // Scale down if over-allocated
    if total > max_total && total > 0.0 {
        let scale = max_total / total;
        for f in fractions.iter_mut() {
            *f *= scale;
        }
    }

    fractions
}

/// Liquidity-adjusted Kelly: reduces sizing when available liquidity is thin.
///
/// Uses square-root scaling: as your target size approaches available liquidity,
/// you face increasing market impact. The sqrt dampens aggressive sizing.
///
/// `available_liquidity`: total shares available at or near the target price.
#[pyfunction]
#[inline]
pub fn liquidity_adjusted_kelly(
    prob: f64,
    market_price: f64,
    bankroll: f64,
    fraction: f64,
    available_liquidity: f64,
    max_size: f64,
) -> f64 {
    let raw_size = kelly_size(prob, market_price, bankroll, fraction, max_size);
    if raw_size <= 0.0 || available_liquidity <= 0.0 {
        return 0.0;
    }
    let ratio = (available_liquidity / raw_size).min(1.0);
    // sqrt scaling: gentle reduction as we approach liquidity limit
    (raw_size * ratio.sqrt()).min(max_size).min(available_liquidity)
}

/// Compute the expected edge (expected value per dollar risked).
///
/// For buying YES: edge = prob * (1 - price) - (1 - prob) * price = prob - price
/// Returns the raw edge (can be negative if no edge).
#[pyfunction]
#[inline]
pub fn edge(prob: f64, market_price: f64) -> f64 {
    prob - market_price
}

/// Register all Kelly functions in the Python module.
pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(kelly, m)?)?;
    m.add_function(wrap_pyfunction!(kelly_no, m)?)?;
    m.add_function(wrap_pyfunction!(fractional_kelly, m)?)?;
    m.add_function(wrap_pyfunction!(kelly_size, m)?)?;
    m.add_function(wrap_pyfunction!(multi_kelly, m)?)?;
    m.add_function(wrap_pyfunction!(liquidity_adjusted_kelly, m)?)?;
    m.add_function(wrap_pyfunction!(edge, m)?)?;
    Ok(())
}

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

    #[test]
    fn kelly_with_edge() {
        // prob=0.60, price=0.50 → f = (0.60 - 0.50) / (1 - 0.50) = 0.20
        let f = kelly(0.60, 0.50);
        assert!((f - 0.20).abs() < 1e-10);
    }

    #[test]
    fn kelly_no_edge_returns_zero() {
        assert_eq!(kelly(0.50, 0.50), 0.0);
        assert_eq!(kelly(0.40, 0.50), 0.0);
    }

    #[test]
    fn kelly_invalid_price_returns_zero() {
        assert_eq!(kelly(0.60, 0.0), 0.0);
        assert_eq!(kelly(0.60, 1.0), 0.0);
        assert_eq!(kelly(0.60, -0.1), 0.0);
    }

    #[test]
    fn kelly_no_side() {
        // prob=0.40, price=0.50 → NO kelly = (0.50 - 0.40) / 0.50 = 0.20
        let f = kelly_no(0.40, 0.50);
        assert!((f - 0.20).abs() < 1e-10);
    }

    #[test]
    fn kelly_no_no_edge() {
        assert_eq!(kelly_no(0.50, 0.50), 0.0);
        assert_eq!(kelly_no(0.60, 0.50), 0.0);
    }

    #[test]
    fn fractional_kelly_half() {
        let full = kelly(0.70, 0.50);
        let half = fractional_kelly(0.70, 0.50, 0.5);
        assert!((half - full * 0.5).abs() < 1e-10);
    }

    #[test]
    fn kelly_size_basic() {
        // prob=0.60, price=0.50, bankroll=1000, fraction=1.0, max=100
        // kelly = 0.20, risk = 200, contracts = 200/0.50 = 400, capped at 100
        let size = kelly_size(0.60, 0.50, 1000.0, 1.0, 100.0);
        assert!((size - 100.0).abs() < 1e-10);
    }

    #[test]
    fn kelly_size_uncapped() {
        // prob=0.55, price=0.50, bankroll=100, fraction=0.5, max=1000
        // kelly = 0.10, frac = 0.05, risk = 5, contracts = 5/0.50 = 10
        let size = kelly_size(0.55, 0.50, 100.0, 0.5, 1000.0);
        assert!((size - 10.0).abs() < 1e-10);
    }

    #[test]
    fn multi_kelly_scales_down() {
        // Two outcomes both with edge, total > 1.0
        let fracs = multi_kelly(vec![0.90, 0.90], vec![0.50, 0.50], 1.0);
        assert_eq!(fracs.len(), 2);
        let total: f64 = fracs.iter().sum();
        assert!(total <= 1.0 + 1e-10);
    }

    #[test]
    fn multi_kelly_no_scaling_needed() {
        let fracs = multi_kelly(vec![0.55, 0.55], vec![0.50, 0.50], 1.0);
        let total: f64 = fracs.iter().sum();
        assert!(total <= 1.0 + 1e-10);
    }

    #[test]
    fn liquidity_adjusted_reduces_size() {
        let full = kelly_size(0.70, 0.50, 1000.0, 1.0, 1000.0);
        let adjusted = liquidity_adjusted_kelly(0.70, 0.50, 1000.0, 1.0, 50.0, 1000.0);
        assert!(adjusted < full);
        assert!(adjusted <= 50.0);
    }

    #[test]
    fn edge_positive() {
        assert!((edge(0.60, 0.50) - 0.10).abs() < 1e-10);
    }

    #[test]
    fn edge_negative() {
        assert!((edge(0.40, 0.50) - (-0.10)).abs() < 1e-10);
    }
}
