//! 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
}

/// Yield-adjusted Kelly fraction.
///
/// Adjusts standard Kelly by subtracting the risk-free opportunity cost.
/// Accounts for the fact that capital locked in a prediction market position
/// could otherwise earn the risk-free rate.
///
/// Formula: max(0, kelly(prob, price) - rfr * duration_days / 365) * fraction
#[pyfunction]
#[pyo3(signature = (prob, price, rfr, duration_days, fraction=1.0))]
#[inline]
pub fn yield_adjusted_kelly(
    prob: f64,
    price: f64,
    rfr: f64,
    duration_days: f64,
    fraction: f64,
) -> f64 {
    if prob.is_nan() || price.is_nan() || rfr.is_nan() || duration_days.is_nan()
        || prob.is_infinite() || price.is_infinite() || rfr.is_infinite() || duration_days.is_infinite()
        || prob < 0.0 || prob > 1.0
        || price <= 0.0 || price >= 1.0
        || duration_days < 0.0
    {
        return 0.0;
    }
    let raw_kelly = kelly(prob, price);
    let opp_cost = rfr * duration_days / 365.0;
    let adjusted = (raw_kelly - opp_cost).max(0.0);
    adjusted * fraction.clamp(0.0, 1.0)
}

/// Compute the opportunity cost of locking capital at a given risk-free rate.
///
/// Returns the dollar amount of risk-free interest foregone.
#[pyfunction]
#[inline]
pub fn opportunity_cost(notional: f64, rfr: f64, duration_days: f64) -> f64 {
    if notional.is_nan() || rfr.is_nan() || duration_days.is_nan()
        || notional.is_infinite() || rfr.is_infinite() || duration_days.is_infinite()
        || notional < 0.0 || duration_days < 0.0
    {
        return 0.0;
    }
    let cost = notional * rfr * duration_days / 365.0;
    if cost.is_finite() { cost } else { 0.0 }
}

/// Compute the breakeven edge required to justify a position given the risk-free rate.
///
/// If your edge is below this threshold, you'd be better off earning risk-free returns.
#[pyfunction]
#[inline]
pub fn breakeven_edge(rfr: f64, duration_days: f64) -> f64 {
    if rfr.is_nan() || duration_days.is_nan()
        || rfr.is_infinite() || duration_days.is_infinite()
        || duration_days < 0.0
    {
        return 0.0;
    }
    let be = rfr * duration_days / 365.0;
    if be.is_finite() { be } else { 0.0 }
}

/// Yield-adjusted Kelly sizing (contracts).
///
/// Combines yield-adjusted Kelly fraction with bankroll sizing.
#[pyfunction]
#[pyo3(signature = (prob, price, rfr, duration_days, bankroll, fraction=0.5, max_size=1000.0))]
#[inline]
pub fn yield_adjusted_kelly_size(
    prob: f64,
    price: f64,
    rfr: f64,
    duration_days: f64,
    bankroll: f64,
    fraction: f64,
    max_size: f64,
) -> f64 {
    if bankroll <= 0.0 || price <= 0.0 {
        return 0.0;
    }
    let f = yield_adjusted_kelly(prob, price, rfr, duration_days, fraction);
    let risk_amount = f * bankroll;
    let contracts = risk_amount / price;
    contracts.min(max_size).max(0.0)
}

/// 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)?)?;
    m.add_function(wrap_pyfunction!(yield_adjusted_kelly, m)?)?;
    m.add_function(wrap_pyfunction!(opportunity_cost, m)?)?;
    m.add_function(wrap_pyfunction!(breakeven_edge, m)?)?;
    m.add_function(wrap_pyfunction!(yield_adjusted_kelly_size, 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);
    }

    #[test]
    fn yield_adjusted_kelly_basic() {
        // With zero risk-free rate, should equal regular fractional kelly
        let ya = yield_adjusted_kelly(0.60, 0.50, 0.0, 30.0, 1.0);
        let fk = fractional_kelly(0.60, 0.50, 1.0);
        assert!((ya - fk).abs() < 1e-10);
    }

    #[test]
    fn yield_adjusted_kelly_with_rfr() {
        // High rfr should reduce kelly fraction
        let ya = yield_adjusted_kelly(0.55, 0.50, 0.05, 30.0, 1.0);
        let fk = fractional_kelly(0.55, 0.50, 1.0);
        assert!(ya < fk);
        assert!(ya >= 0.0);
    }

    #[test]
    fn yield_adjusted_kelly_no_edge_after_rfr() {
        // Tiny edge wiped out by rfr
        let ya = yield_adjusted_kelly(0.51, 0.50, 0.50, 365.0, 1.0);
        assert!((ya - 0.0).abs() < 1e-10);
    }

    #[test]
    fn yield_adjusted_kelly_invalid() {
        assert_eq!(yield_adjusted_kelly(f64::NAN, 0.5, 0.05, 30.0, 1.0), 0.0);
        assert_eq!(yield_adjusted_kelly(0.6, 0.5, 0.05, -1.0, 1.0), 0.0);
    }

    #[test]
    fn opportunity_cost_basic() {
        let oc = opportunity_cost(10000.0, 0.05, 365.0);
        assert!((oc - 500.0).abs() < 1e-6);
    }

    #[test]
    fn opportunity_cost_zero() {
        assert_eq!(opportunity_cost(0.0, 0.05, 30.0), 0.0);
        assert_eq!(opportunity_cost(1000.0, 0.0, 30.0), 0.0);
    }

    #[test]
    fn opportunity_cost_invalid() {
        assert_eq!(opportunity_cost(f64::NAN, 0.05, 30.0), 0.0);
        assert_eq!(opportunity_cost(-100.0, 0.05, 30.0), 0.0);
    }

    #[test]
    fn breakeven_edge_basic() {
        // 5% annual rate, 73 days
        let be = breakeven_edge(0.05, 73.0);
        assert!((be - 0.01).abs() < 0.001);
    }

    #[test]
    fn breakeven_edge_invalid() {
        assert_eq!(breakeven_edge(f64::NAN, 30.0), 0.0);
        assert_eq!(breakeven_edge(0.05, -1.0), 0.0);
    }

    #[test]
    fn yield_adjusted_kelly_size_basic() {
        let size = yield_adjusted_kelly_size(0.60, 0.50, 0.0, 30.0, 1000.0, 1.0, 100.0);
        assert!(size > 0.0);
        assert!(size <= 100.0);
    }

    #[test]
    fn yield_adjusted_kelly_size_zero_bankroll() {
        assert_eq!(yield_adjusted_kelly_size(0.60, 0.50, 0.05, 30.0, 0.0, 1.0, 100.0), 0.0);
    }
}
