//! Monte Carlo simulation for risk scenarios
//!
//! Simulates future portfolio paths using geometric Brownian motion

use crate::error::{Result, RiskAnalyzerError};
use rand::SeedableRng;
use rand_distr::{Distribution, Normal};
use risk_analyzer_shared::MonteCarloResult;

/// Run Monte Carlo simulation for portfolio returns
///
/// # Arguments
/// * `initial_value` - Initial portfolio value
/// * `mean_return` - Expected return per period
/// * `volatility` - Volatility per period
/// * `num_steps` - Number of time steps
/// * `num_simulations` - Number of simulation paths
/// * `num_paths_to_return` - Number of paths to return for visualization (subset)
pub fn run_monte_carlo_simulation(
    initial_value: f64,
    mean_return: f64,
    volatility: f64,
    num_steps: usize,
    num_simulations: usize,
    num_paths_to_return: usize,
) -> Result<MonteCarloResult> {
    if num_simulations == 0 {
        return Err(RiskAnalyzerError::InvalidConfig(
            "Number of simulations must be positive".to_string(),
        ));
    }

    if num_steps == 0 {
        return Err(RiskAnalyzerError::InvalidConfig(
            "Number of steps must be positive".to_string(),
        ));
    }

    if initial_value <= 0.0 {
        return Err(RiskAnalyzerError::InvalidConfig(
            "Initial value must be positive".to_string(),
        ));
    }

    if volatility < 0.0 {
        return Err(RiskAnalyzerError::InvalidConfig(
            "Volatility must be non-negative".to_string(),
        ));
    }

    let mut rng = rand::rngs::StdRng::seed_from_u64(42);
    let normal = Normal::new(0.0, 1.0).map_err(|e| {
        RiskAnalyzerError::CalculationError(format!("Failed to create normal distribution: {e}"))
    })?;

    let dt = 1.0;
    let drift = mean_return - 0.5 * volatility * volatility;

    let mut all_final_values = Vec::with_capacity(num_simulations);
    let mut paths_to_store = Vec::with_capacity(num_paths_to_return.min(num_simulations));

    for sim_idx in 0..num_simulations {
        let mut path = Vec::with_capacity(num_steps + 1);
        let mut value = initial_value;
        path.push(value);

        for _ in 0..num_steps {
            let z = normal.sample(&mut rng);
            let return_val = drift * dt + volatility * dt.sqrt() * z;
            value *= (1.0 + return_val).max(0.01);
            path.push(value);
        }

        all_final_values.push(value);

        if sim_idx < num_paths_to_return {
            paths_to_store.push(path);
        }
    }

    let mut sorted_finals = all_final_values.clone();
    sorted_finals.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

    let percentiles = vec![
        (
            0.01,
            sorted_finals[(0.01 * num_simulations as f64) as usize],
        ),
        (
            0.05,
            sorted_finals[(0.05 * num_simulations as f64) as usize],
        ),
        (
            0.10,
            sorted_finals[(0.10 * num_simulations as f64) as usize],
        ),
        (
            0.25,
            sorted_finals[(0.25 * num_simulations as f64) as usize],
        ),
        (
            0.50,
            sorted_finals[(0.50 * num_simulations as f64) as usize],
        ),
        (
            0.75,
            sorted_finals[(0.75 * num_simulations as f64) as usize],
        ),
        (
            0.90,
            sorted_finals[(0.90 * num_simulations as f64) as usize],
        ),
        (
            0.95,
            sorted_finals[(0.95 * num_simulations as f64) as usize],
        ),
        (
            0.99,
            sorted_finals[((0.99 * num_simulations as f64) as usize).min(num_simulations - 1)],
        ),
    ];

    let mean_final = all_final_values.iter().sum::<f64>() / num_simulations as f64;
    let median_final = sorted_finals[num_simulations / 2];

    let variance = all_final_values
        .iter()
        .map(|&v| (v - mean_final).powi(2))
        .sum::<f64>()
        / num_simulations as f64;
    let std_final = variance.sqrt();

    Ok(MonteCarloResult {
        paths: paths_to_store,
        final_values: all_final_values,
        percentiles,
        mean_final_value: mean_final,
        median_final_value: median_final,
        std_final_value: std_final,
    })
}

/// Generate portfolio returns from Monte Carlo simulation
///
/// Returns a vector of simulated portfolio returns
pub fn simulate_portfolio_returns(
    mean_return: f64,
    volatility: f64,
    num_simulations: usize,
    horizon_days: usize,
) -> Result<Vec<f64>> {
    if num_simulations == 0 {
        return Err(RiskAnalyzerError::InvalidConfig(
            "Number of simulations must be positive".to_string(),
        ));
    }

    if horizon_days == 0 {
        return Err(RiskAnalyzerError::InvalidConfig(
            "Horizon must be positive".to_string(),
        ));
    }

    let mut rng = rand::rngs::StdRng::seed_from_u64(42);
    let normal = Normal::new(mean_return, volatility).map_err(|e| {
        RiskAnalyzerError::CalculationError(format!("Failed to create normal distribution: {e}"))
    })?;

    let mut returns = Vec::with_capacity(num_simulations);

    for _ in 0..num_simulations {
        let mut cumulative_return = 0.0;
        for _ in 0..horizon_days {
            cumulative_return += normal.sample(&mut rng);
        }
        returns.push(cumulative_return);
    }

    Ok(returns)
}

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

    #[test]
    fn test_run_monte_carlo_zero_simulations() {
        let result =
            run_monte_carlo_simulation(100_000.0, 0.08 / 252.0, 0.2 / 252.0_f64.sqrt(), 10, 0, 5);
        assert!(result.is_err());
    }

    #[test]
    fn test_run_monte_carlo_zero_steps() {
        let result =
            run_monte_carlo_simulation(100_000.0, 0.08 / 252.0, 0.2 / 252.0_f64.sqrt(), 0, 1000, 5);
        assert!(result.is_err());
    }

    #[test]
    fn test_run_monte_carlo_negative_initial() {
        let result = run_monte_carlo_simulation(
            -100_000.0,
            0.08 / 252.0,
            0.2 / 252.0_f64.sqrt(),
            10,
            1000,
            5,
        );
        assert!(result.is_err());
    }

    #[test]
    fn test_run_monte_carlo_negative_volatility() {
        let result = run_monte_carlo_simulation(100_000.0, 0.08 / 252.0, -0.2, 10, 1000, 5);
        assert!(result.is_err());
    }

    #[test]
    fn test_run_monte_carlo_basic() {
        let result = run_monte_carlo_simulation(
            100_000.0,
            0.08 / 252.0,
            0.2 / 252.0_f64.sqrt(),
            252,
            1000,
            10,
        )
        .unwrap();

        assert_eq!(result.paths.len(), 10);
        assert_eq!(result.final_values.len(), 1000);
        assert_eq!(result.percentiles.len(), 9);

        for path in &result.paths {
            assert_eq!(path.len(), 253);
            assert!((path[0] - 100_000.0).abs() < 1.0);
        }

        assert!(result.mean_final_value > 0.0);
        assert!(result.median_final_value > 0.0);
        assert!(result.std_final_value > 0.0);
    }

    #[test]
    fn test_run_monte_carlo_percentiles_ordered() {
        let result = run_monte_carlo_simulation(
            100_000.0,
            0.08 / 252.0,
            0.2 / 252.0_f64.sqrt(),
            100,
            10_000,
            5,
        )
        .unwrap();

        for i in 1..result.percentiles.len() {
            assert!(
                result.percentiles[i].1 >= result.percentiles[i - 1].1,
                "Percentiles should be ordered"
            );
        }
    }

    #[test]
    fn test_run_monte_carlo_positive_drift() {
        let result = run_monte_carlo_simulation(
            100_000.0,
            0.10 / 252.0,
            0.15 / 252.0_f64.sqrt(),
            252,
            5000,
            5,
        )
        .unwrap();

        assert!(
            result.mean_final_value > 100_000.0,
            "With positive drift, mean should exceed initial value"
        );
    }

    #[test]
    fn test_simulate_portfolio_returns_zero_simulations() {
        let result = simulate_portfolio_returns(0.001, 0.02, 0, 10);
        assert!(result.is_err());
    }

    #[test]
    fn test_simulate_portfolio_returns_zero_horizon() {
        let result = simulate_portfolio_returns(0.001, 0.02, 1000, 0);
        assert!(result.is_err());
    }

    #[test]
    fn test_simulate_portfolio_returns_basic() {
        let returns = simulate_portfolio_returns(0.001, 0.02, 10_000, 1).unwrap();
        assert_eq!(returns.len(), 10_000);
        let mean = returns.iter().sum::<f64>() / returns.len() as f64;
        assert!((mean - 0.001).abs() < 0.005);
    }

    #[test]
    fn test_simulate_portfolio_returns_horizon_scaling() {
        let returns_1 = simulate_portfolio_returns(0.001, 0.02, 10_000, 1).unwrap();
        let returns_10 = simulate_portfolio_returns(0.001, 0.02, 10_000, 10).unwrap();

        let std_1 = {
            let mean = returns_1.iter().sum::<f64>() / returns_1.len() as f64;
            (returns_1.iter().map(|&r| (r - mean).powi(2)).sum::<f64>() / returns_1.len() as f64)
                .sqrt()
        };

        let std_10 = {
            let mean = returns_10.iter().sum::<f64>() / returns_10.len() as f64;
            (returns_10.iter().map(|&r| (r - mean).powi(2)).sum::<f64>() / returns_10.len() as f64)
                .sqrt()
        };

        assert!(
            std_10 > std_1,
            "Longer horizon should have higher variability"
        );
    }

    #[test]
    fn test_monte_carlo_paths_start_at_initial_value() {
        let result = run_monte_carlo_simulation(
            50_000.0,
            0.05 / 252.0,
            0.18 / 252.0_f64.sqrt(),
            100,
            500,
            10,
        )
        .unwrap();

        for path in &result.paths {
            assert!((path[0] - 50_000.0).abs() < 0.01);
        }
    }

    #[test]
    fn test_monte_carlo_final_values_distribution() {
        let result = run_monte_carlo_simulation(
            100_000.0,
            0.08 / 252.0,
            0.2 / 252.0_f64.sqrt(),
            252,
            10_000,
            5,
        )
        .unwrap();

        let min_value = result
            .final_values
            .iter()
            .copied()
            .fold(f64::INFINITY, f64::min);
        let max_value = result
            .final_values
            .iter()
            .copied()
            .fold(f64::NEG_INFINITY, f64::max);

        assert!(min_value > 0.0);
        assert!(max_value > min_value);
        assert!(result.mean_final_value > min_value);
        assert!(result.mean_final_value < max_value);
    }
}