//! Monte Carlo simulation for prediction market portfolios.
//!
//! Implements:
//! - xoshiro256++ PRNG (fast, high-quality)
//! - Box-Muller transform for normal distribution
//! - Cholesky decomposition for correlated draws
//! - Binary outcome simulation (Bernoulli via CDF threshold)

use pyo3::prelude::*;

// ---------------------------------------------------------------------------
// xoshiro256++ PRNG
// ---------------------------------------------------------------------------

struct Xoshiro256PlusPlus {
    s: [u64; 4],
}

impl Xoshiro256PlusPlus {
    fn new(seed: u64) -> Self {
        // Use SplitMix64 to initialize state from a single seed
        let mut state = seed;
        let mut s = [0u64; 4];
        for slot in &mut s {
            state = state.wrapping_add(0x9E3779B97F4A7C15);
            let mut z = state;
            z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
            z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
            z ^= z >> 31;
            *slot = z;
        }
        Self { s }
    }

    #[inline]
    fn next_u64(&mut self) -> u64 {
        let result = (self.s[0].wrapping_add(self.s[3]))
            .rotate_left(23)
            .wrapping_add(self.s[0]);
        let t = self.s[1] << 17;
        self.s[2] ^= self.s[0];
        self.s[3] ^= self.s[1];
        self.s[1] ^= self.s[2];
        self.s[0] ^= self.s[3];
        self.s[2] ^= t;
        self.s[3] = self.s[3].rotate_left(45);
        result
    }

    /// Generate a uniform f64 in [0, 1).
    #[inline]
    fn next_f64(&mut self) -> f64 {
        let bits = self.next_u64() >> 11;
        bits as f64 * (1.0 / (1u64 << 53) as f64)
    }
}

// ---------------------------------------------------------------------------
// Box-Muller transform
// ---------------------------------------------------------------------------

/// Generate a pair of standard normal random variables using Box-Muller.
#[inline]
fn box_muller(rng: &mut Xoshiro256PlusPlus) -> (f64, f64) {
    loop {
        let u1 = rng.next_f64();
        let u2 = rng.next_f64();
        if u1 > 1e-15 {
            let r = (-2.0 * u1.ln()).sqrt();
            let theta = 2.0 * std::f64::consts::PI * u2;
            return (r * theta.cos(), r * theta.sin());
        }
    }
}

// ---------------------------------------------------------------------------
// Standard normal CDF (rational approximation)
// ---------------------------------------------------------------------------

/// Approximation of the standard normal CDF using Abramowitz & Stegun.
#[inline]
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)
}

// ---------------------------------------------------------------------------
// Cholesky decomposition
// ---------------------------------------------------------------------------

/// Cholesky decomposition of a symmetric positive-definite matrix.
/// Returns lower-triangular matrix L such that A = L * L^T.
/// Falls back to identity if decomposition fails.
fn cholesky(matrix: &[Vec<f64>]) -> Vec<Vec<f64>> {
    let n = matrix.len();
    let mut l = vec![vec![0.0f64; n]; n];

    for i in 0..n {
        for j in 0..=i {
            let mut sum = 0.0;
            for k in 0..j {
                sum += l[i][k] * l[j][k];
            }
            if i == j {
                let diag = matrix[i][i] - sum;
                if diag <= 0.0 {
                    // Not positive definite — fallback to identity
                    tracing::warn!("cholesky: matrix not positive definite, falling back to identity");
                    let mut identity = vec![vec![0.0f64; n]; n];
                    for idx in 0..n {
                        identity[idx][idx] = 1.0;
                    }
                    return identity;
                }
                l[i][j] = diag.sqrt();
            } else {
                l[i][j] = (matrix[i][j] - sum) / l[j][j];
            }
        }
    }

    l
}

// ---------------------------------------------------------------------------
// Python types
// ---------------------------------------------------------------------------

/// A simulated position for Monte Carlo analysis.
#[pyclass]
#[derive(Debug, Clone)]
pub struct SimPosition {
    #[pyo3(get)]
    pub market_id: String,
    /// "yes" or "no"
    #[pyo3(get)]
    pub side: String,
    #[pyo3(get)]
    pub size: f64,
    #[pyo3(get)]
    pub entry_price: f64,
    #[pyo3(get)]
    pub current_price: f64,
}

#[pymethods]
impl SimPosition {
    #[new]
    #[pyo3(signature = (market_id, side, size, entry_price, current_price))]
    fn new(
        market_id: String,
        side: String,
        size: f64,
        entry_price: f64,
        current_price: f64,
    ) -> Self {
        Self {
            market_id,
            side: side.to_lowercase(),
            size,
            entry_price,
            current_price: current_price.clamp(0.001, 0.999),
        }
    }

    fn __repr__(&self) -> String {
        format!(
            "SimPosition('{}', '{}', size={:.1}, entry={:.4}, current={:.4})",
            self.market_id, self.side, self.size, self.entry_price, self.current_price
        )
    }
}

/// Results of a Monte Carlo simulation.
#[pyclass]
#[derive(Debug, Clone)]
pub struct SimulationResult {
    #[pyo3(get)]
    pub mean_pnl: f64,
    #[pyo3(get)]
    pub median_pnl: f64,
    #[pyo3(get)]
    pub std_dev: f64,
    /// Value at Risk at 95% confidence (5th percentile loss)
    #[pyo3(get)]
    pub var_95: f64,
    /// Value at Risk at 99% confidence (1st percentile loss)
    #[pyo3(get)]
    pub var_99: f64,
    /// Conditional VaR (Expected Shortfall) at 95%
    #[pyo3(get)]
    pub cvar_95: f64,
    #[pyo3(get)]
    pub max_loss: f64,
    #[pyo3(get)]
    pub max_gain: f64,
    /// Probability of positive PnL
    #[pyo3(get)]
    pub win_probability: f64,
    /// All scenario PnLs (sorted ascending)
    #[pyo3(get)]
    pub scenario_pnl: Vec<f64>,
    /// Key percentiles: [(5, val), (10, val), (25, val), (50, val), (75, val), (90, val), (95, val)]
    #[pyo3(get)]
    pub percentiles: Vec<(f64, f64)>,
}

#[pymethods]
impl SimulationResult {
    fn __repr__(&self) -> String {
        format!(
            "SimulationResult(mean={:.2}, var95={:.2}, win%={:.1}%, scenarios={})",
            self.mean_pnl,
            self.var_95,
            self.win_probability * 100.0,
            self.scenario_pnl.len()
        )
    }
}

// ---------------------------------------------------------------------------
// Monte Carlo core
// ---------------------------------------------------------------------------

/// Run a Monte Carlo simulation over a portfolio of prediction market positions.
///
/// Algorithm:
/// 1. For each scenario, draw correlated normal random variables (one per position)
/// 2. Threshold via normal_cdf to get a probability, compare to current_price for Bernoulli outcome
/// 3. Compute PnL: YES side = size * (outcome - entry_price), NO side = size * ((1 - outcome) - entry_price)
/// 4. Aggregate statistics across all scenarios
///
/// Args:
///     positions: List of SimPosition
///     scenarios: Number of Monte Carlo scenarios (default 10000)
///     correlation_matrix: Optional NxN correlation matrix as Vec<Vec<f64>>
///     seed: Optional PRNG seed for reproducibility
#[pyfunction]
#[pyo3(signature = (positions, scenarios=10000, correlation_matrix=None, seed=None))]
pub fn monte_carlo(
    positions: Vec<SimPosition>,
    scenarios: usize,
    correlation_matrix: Option<Vec<Vec<f64>>>,
    seed: Option<u64>,
) -> PyResult<SimulationResult> {
    if positions.is_empty() {
        return Ok(SimulationResult {
            mean_pnl: 0.0,
            median_pnl: 0.0,
            std_dev: 0.0,
            var_95: 0.0,
            var_99: 0.0,
            cvar_95: 0.0,
            max_loss: 0.0,
            max_gain: 0.0,
            win_probability: 0.0,
            scenario_pnl: vec![],
            percentiles: vec![],
        });
    }

    let n = positions.len();
    let scenarios = scenarios.max(1).min(1_000_000); // Cap at 1M

    // Validate and prepare correlation/Cholesky
    let chol = match correlation_matrix {
        Some(ref corr) => {
            if corr.len() != n || corr.iter().any(|row| row.len() != n) {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "correlation matrix must be {}x{}, got {}x{}",
                    n,
                    n,
                    corr.len(),
                    corr.first().map_or(0, |r| r.len())
                )));
            }
            Some(cholesky(corr))
        }
        None => None,
    };

    let mut rng = Xoshiro256PlusPlus::new(seed.unwrap_or(12345));
    let mut scenario_pnls = Vec::with_capacity(scenarios);

    for _ in 0..scenarios {
        // Generate N standard normals
        let mut normals = Vec::with_capacity(n);
        let mut i = 0;
        while i < n {
            let (z1, z2) = box_muller(&mut rng);
            normals.push(z1);
            if i + 1 < n {
                normals.push(z2);
            }
            i += 2;
        }

        // Apply Cholesky correlation if provided
        if let Some(ref l) = chol {
            let uncorrelated = normals.clone();
            for j in 0..n {
                let mut sum = 0.0;
                for k in 0..=j {
                    sum += l[j][k] * uncorrelated[k];
                }
                normals[j] = sum;
            }
        }

        // Compute scenario PnL
        let mut total_pnl = 0.0;
        for (j, pos) in positions.iter().enumerate() {
            let z = normals[j];
            // Convert normal to uniform via CDF, threshold for Bernoulli outcome
            // U < current_price → outcome=1 (with probability = current_price)
            let u = normal_cdf(z);
            let outcome = if u < pos.current_price { 1.0 } else { 0.0 };

            let pnl = if pos.side == "yes" {
                pos.size * (outcome - pos.entry_price)
            } else {
                pos.size * ((1.0 - outcome) - pos.entry_price)
            };

            total_pnl += pnl;
        }

        scenario_pnls.push(total_pnl);
    }

    // Sort for percentile computation
    scenario_pnls.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

    let s = scenarios as f64;
    let mean_pnl = scenario_pnls.iter().sum::<f64>() / s;
    let variance = scenario_pnls.iter().map(|p| (p - mean_pnl).powi(2)).sum::<f64>() / s;
    let std_dev = variance.sqrt();

    let median_idx = scenarios / 2;
    let median_pnl = scenario_pnls[median_idx];

    let var_95_idx = ((scenarios as f64) * 0.05).floor() as usize;
    let var_99_idx = ((scenarios as f64) * 0.01).floor() as usize;
    let var_95 = scenario_pnls[var_95_idx.min(scenarios - 1)];
    let var_99 = scenario_pnls[var_99_idx.min(scenarios - 1)];

    // CVaR 95: average of worst 5%
    let cvar_count = (var_95_idx + 1).max(1);
    let cvar_95 = scenario_pnls[..cvar_count].iter().sum::<f64>() / cvar_count as f64;

    let max_loss = scenario_pnls[0];
    let max_gain = scenario_pnls[scenarios - 1];

    let win_count = scenario_pnls.iter().filter(|&&p| p > 0.0).count();
    let win_probability = win_count as f64 / s;

    // Percentiles
    let pct_levels = [5.0, 10.0, 25.0, 50.0, 75.0, 90.0, 95.0];
    let percentiles: Vec<(f64, f64)> = pct_levels
        .iter()
        .map(|&pct| {
            let idx = ((scenarios as f64) * pct / 100.0).floor() as usize;
            (pct, scenario_pnls[idx.min(scenarios - 1)])
        })
        .collect();

    Ok(SimulationResult {
        mean_pnl,
        median_pnl,
        std_dev,
        var_95,
        var_99,
        cvar_95,
        max_loss,
        max_gain,
        win_probability,
        scenario_pnl: scenario_pnls,
        percentiles,
    })
}

/// Register simulation types and functions in the Python module.
pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<SimPosition>()?;
    m.add_class::<SimulationResult>()?;
    m.add_function(wrap_pyfunction!(monte_carlo, m)?)?;
    Ok(())
}

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

    #[test]
    fn xoshiro_deterministic() {
        let mut rng1 = Xoshiro256PlusPlus::new(42);
        let mut rng2 = Xoshiro256PlusPlus::new(42);
        for _ in 0..100 {
            assert_eq!(rng1.next_u64(), rng2.next_u64());
        }
    }

    #[test]
    fn xoshiro_different_seeds() {
        let mut rng1 = Xoshiro256PlusPlus::new(42);
        let mut rng2 = Xoshiro256PlusPlus::new(99);
        let same = (0..10).all(|_| rng1.next_u64() == rng2.next_u64());
        assert!(!same);
    }

    #[test]
    fn xoshiro_uniform_range() {
        let mut rng = Xoshiro256PlusPlus::new(42);
        for _ in 0..10_000 {
            let v = rng.next_f64();
            assert!((0.0..1.0).contains(&v));
        }
    }

    #[test]
    fn box_muller_produces_normals() {
        let mut rng = Xoshiro256PlusPlus::new(42);
        let mut sum = 0.0;
        let n = 10_000;
        for _ in 0..n / 2 {
            let (z1, z2) = box_muller(&mut rng);
            sum += z1 + z2;
        }
        let mean = sum / n as f64;
        // Mean should be close to 0
        assert!(mean.abs() < 0.05, "mean={}", mean);
    }

    #[test]
    fn normal_cdf_at_zero() {
        let p = normal_cdf(0.0);
        assert!((p - 0.5).abs() < 1e-6);
    }

    #[test]
    fn normal_cdf_monotonic() {
        let p1 = normal_cdf(-1.0);
        let p2 = normal_cdf(0.0);
        let p3 = normal_cdf(1.0);
        assert!(p1 < p2);
        assert!(p2 < p3);
    }

    #[test]
    fn normal_cdf_tails() {
        assert!(normal_cdf(-3.0) < 0.01);
        assert!(normal_cdf(3.0) > 0.99);
    }

    #[test]
    fn cholesky_identity() {
        let m = vec![vec![1.0, 0.0], vec![0.0, 1.0]];
        let l = cholesky(&m);
        assert!((l[0][0] - 1.0).abs() < 1e-10);
        assert!((l[1][1] - 1.0).abs() < 1e-10);
        assert!((l[0][1]).abs() < 1e-10);
        assert!((l[1][0]).abs() < 1e-10);
    }

    #[test]
    fn cholesky_correlated() {
        let m = vec![vec![1.0, 0.5], vec![0.5, 1.0]];
        let l = cholesky(&m);
        // L * L^T should equal m
        let reconstructed_00 = l[0][0] * l[0][0];
        let reconstructed_01 = l[1][0] * l[0][0];
        let reconstructed_11 = l[1][0] * l[1][0] + l[1][1] * l[1][1];
        assert!((reconstructed_00 - 1.0).abs() < 1e-10);
        assert!((reconstructed_01 - 0.5).abs() < 1e-10);
        assert!((reconstructed_11 - 1.0).abs() < 1e-10);
    }

    #[test]
    fn cholesky_not_positive_definite_falls_back() {
        let m = vec![vec![1.0, 2.0], vec![2.0, 1.0]]; // Not PD
        let l = cholesky(&m);
        // Should fallback to identity
        assert!((l[0][0] - 1.0).abs() < 1e-10);
        assert!((l[1][1] - 1.0).abs() < 1e-10);
    }

    #[test]
    fn monte_carlo_empty_positions() {
        let result = monte_carlo(vec![], 1000, None, Some(42)).unwrap();
        assert_eq!(result.mean_pnl, 0.0);
        assert!(result.scenario_pnl.is_empty());
    }

    #[test]
    fn monte_carlo_deterministic() {
        let pos = vec![SimPosition::new(
            "mkt1".to_string(),
            "yes".to_string(),
            100.0,
            0.50,
            0.60,
        )];
        let r1 = monte_carlo(pos.clone(), 10000, None, Some(42)).unwrap();
        let r2 = monte_carlo(pos, 10000, None, Some(42)).unwrap();
        assert!((r1.mean_pnl - r2.mean_pnl).abs() < 1e-10);
        assert!((r1.var_95 - r2.var_95).abs() < 1e-10);
    }

    #[test]
    fn monte_carlo_yes_position_stats() {
        let pos = vec![SimPosition::new(
            "mkt1".to_string(),
            "yes".to_string(),
            100.0,
            0.50,
            0.60,
        )];
        let result = monte_carlo(pos, 50000, None, Some(42)).unwrap();

        // With current_price=0.60, U < 0.60 → outcome=1 (60% of the time)
        // Win PnL: 100 * (1 - 0.50) = 50
        // Lose PnL: 100 * (0 - 0.50) = -50
        // Expected: 0.60 * 50 + 0.40 * (-50) = 30 - 20 = 10
        assert!((result.mean_pnl - 10.0).abs() < 2.0, "mean={}", result.mean_pnl);
        assert!(result.win_probability > 0.55 && result.win_probability < 0.65);
        assert!(result.max_loss < 0.0);
        assert!(result.max_gain > 0.0);
        assert!(result.var_95 <= result.median_pnl);
    }

    #[test]
    fn monte_carlo_no_position() {
        let pos = vec![SimPosition::new(
            "mkt1".to_string(),
            "no".to_string(),
            100.0,
            0.40,
            0.60,
        )];
        let result = monte_carlo(pos, 50000, None, Some(42)).unwrap();

        // current_price=0.60, U < 0.60 → outcome=1 (60%)
        // outcome=1: NO PnL = 100*(0-0.40) = -40
        // outcome=0: NO PnL = 100*(1-0.40) = 60
        // Expected: 0.60*(-40) + 0.40*60 = -24 + 24 = 0
        assert!((result.mean_pnl).abs() < 2.0, "mean={}", result.mean_pnl);
    }

    #[test]
    fn monte_carlo_with_correlation() {
        let pos = vec![
            SimPosition::new("mkt1".to_string(), "yes".to_string(), 50.0, 0.50, 0.60),
            SimPosition::new("mkt2".to_string(), "yes".to_string(), 50.0, 0.50, 0.60),
        ];
        let corr = vec![vec![1.0, 0.8], vec![0.8, 1.0]];
        let result = monte_carlo(pos, 10000, Some(corr), Some(42)).unwrap();
        // With high correlation, variance should be higher than independent
        assert!(result.std_dev > 0.0);
    }

    #[test]
    fn monte_carlo_bad_correlation_dims() {
        let pos = vec![SimPosition::new(
            "mkt1".to_string(),
            "yes".to_string(),
            100.0,
            0.50,
            0.60,
        )];
        let corr = vec![vec![1.0, 0.5], vec![0.5, 1.0]]; // 2x2 but only 1 position
        let result = monte_carlo(pos, 100, Some(corr), Some(42));
        assert!(result.is_err());
    }

    #[test]
    fn monte_carlo_percentiles_sorted() {
        let pos = vec![SimPosition::new(
            "mkt1".to_string(),
            "yes".to_string(),
            100.0,
            0.50,
            0.60,
        )];
        let result = monte_carlo(pos, 10000, None, Some(42)).unwrap();
        for i in 1..result.percentiles.len() {
            assert!(result.percentiles[i].1 >= result.percentiles[i - 1].1);
        }
    }

    #[test]
    fn monte_carlo_cvar_less_than_var() {
        let pos = vec![SimPosition::new(
            "mkt1".to_string(),
            "yes".to_string(),
            100.0,
            0.50,
            0.60,
        )];
        let result = monte_carlo(pos, 10000, None, Some(42)).unwrap();
        // CVaR (expected shortfall) should be <= VaR (it's the average of the worst tail)
        assert!(result.cvar_95 <= result.var_95 + 1e-10);
    }
}
