//! Stochastic simulation runner — Euler-Maruyama with Rayon parallelism.
//!
//! Generates synthetic price paths from an SDE defined via the expression DSL.

use rand::{Rng, SeedableRng};
use rand_chacha::ChaCha8Rng;
use rand_distr::{Distribution, Poisson};
use rayon::prelude::*;
use serde::{Deserialize, Serialize};

use bt_kernel::{BtError, Result};

use super::compiler::SdeProgram;
use super::sde_ops::{eval_tape, SdeBuiltins};
use crate::monte_carlo::DistributionStats;

// ---------------------------------------------------------------------------
// Config
// ---------------------------------------------------------------------------

/// Configuration for a stochastic simulation run.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct StochasticSimConfig {
    /// Number of simulation paths.
    pub n_paths: usize,
    /// Number of time steps per path.
    pub n_steps: usize,
    /// Time step size in years (e.g. 1.0/252.0 for daily).
    pub dt: f64,
    /// Initial price.
    pub s0: f64,
    /// Quantile levels for reporting.
    #[serde(default = "default_confidence_levels")]
    pub confidence_levels: Vec<f64>,
    /// Fixed RNG seed for reproducibility. `None` → random.
    pub rng_seed: Option<u64>,
    /// Store full price paths (memory-intensive).
    #[serde(default)]
    pub store_paths: bool,
    /// GPU precision: "f32" or "f64" (default). f32 is ~10-20x faster on consumer GPUs.
    #[serde(default = "default_precision")]
    pub precision: String,
}

fn default_precision() -> String {
    "f64".into()
}

fn default_confidence_levels() -> Vec<f64> {
    vec![0.05, 0.10, 0.25, 0.50, 0.75, 0.90, 0.95]
}

// ---------------------------------------------------------------------------
// Result
// ---------------------------------------------------------------------------

/// Output of a stochastic simulation.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct StochasticResult {
    pub n_paths: usize,
    pub n_steps: usize,
    pub model_name: String,
    /// Distribution of final prices.
    pub final_price: DistributionStats,
    /// Distribution of total return (S_T / S_0 - 1).
    pub final_return: DistributionStats,
    /// Distribution of max drawdown (0.0–1.0).
    pub max_drawdown: DistributionStats,
    /// Distribution of annualized return.
    pub annualized_return: DistributionStats,
    /// Distribution of annualized volatility.
    pub annualized_vol: DistributionStats,
    /// Full price paths `[path_idx][step]`. `None` when `store_paths = false`.
    pub paths: Option<Vec<Vec<f64>>>,
}

// ---------------------------------------------------------------------------
// Runner
// ---------------------------------------------------------------------------

/// Run a stochastic simulation using a compiled SDE program.
pub fn run_stochastic(
    program: &SdeProgram,
    config: &StochasticSimConfig,
) -> Result<StochasticResult> {
    // Validation
    if config.n_paths == 0 {
        return Err(BtError::Strategy("stochastic: n_paths must be > 0".into()));
    }
    if config.n_steps == 0 {
        return Err(BtError::Strategy("stochastic: n_steps must be > 0".into()));
    }
    if config.dt <= 0.0 {
        return Err(BtError::Strategy("stochastic: dt must be > 0".into()));
    }
    if config.s0 <= 0.0 {
        return Err(BtError::Strategy("stochastic: s0 must be > 0".into()));
    }

    // License gate
    let max_sims = bt_license::max_monte_carlo_sims();
    if config.n_paths > max_sims {
        return Err(BtError::Strategy(format!(
            "stochastic: {} paths requested but Community license allows max {}. \
             Upgrade to Pro for unlimited simulations.",
            config.n_paths, max_sims
        )));
    }

    let base_seed = resolve_base_seed(config.rng_seed);
    let store = config.store_paths;
    let dt = config.dt;
    let s0 = config.s0;
    let n_steps = config.n_steps;
    let params = &program.param_values;
    let has_jumps = program.jump_intensity.is_some();

    // Run paths in parallel
    let results: Vec<PathResult> = (0..config.n_paths)
        .into_par_iter()
        .map(|i| {
            let mut rng = path_rng(base_seed, i);
            simulate_path(program, params, s0, n_steps, dt, has_jumps, store, &mut rng)
        })
        .collect();

    // Aggregate statistics
    let mut final_prices = Vec::with_capacity(config.n_paths);
    let mut final_returns = Vec::with_capacity(config.n_paths);
    let mut max_drawdowns = Vec::with_capacity(config.n_paths);
    let mut ann_returns = Vec::with_capacity(config.n_paths);
    let mut ann_vols = Vec::with_capacity(config.n_paths);
    let mut paths: Option<Vec<Vec<f64>>> = None;

    let total_time = n_steps as f64 * dt;

    for r in results {
        final_prices.push(r.final_price);

        let ret = r.final_price / s0 - 1.0;
        final_returns.push(ret);
        max_drawdowns.push(r.max_drawdown);

        // Annualized return: (S_T / S_0)^(1/T) - 1
        let ratio = r.final_price / s0;
        let ann_ret = if total_time > 0.0 && ratio > 0.0 {
            ratio.powf(1.0 / total_time) - 1.0
        } else {
            0.0
        };
        ann_returns.push(ann_ret);

        // Annualized vol from log returns
        let ann_vol = if r.sum_log_ret_sq > 0.0 && n_steps > 1 {
            let var = r.sum_log_ret_sq / n_steps as f64;
            (var / dt).sqrt()
        } else {
            0.0
        };
        ann_vols.push(ann_vol);

        if let Some(p) = r.path {
            paths.get_or_insert_with(Vec::new).push(p);
        }
    }

    let cl = &config.confidence_levels;

    // Compute all 5 stats in parallel
    let ((final_price, final_return), (max_drawdown, (annualized_return, annualized_vol))) =
        rayon::join(
            || rayon::join(|| build_stats(&final_prices, cl), || build_stats(&final_returns, cl)),
            || rayon::join(
                || build_stats(&max_drawdowns, cl),
                || rayon::join(|| build_stats(&ann_returns, cl), || build_stats(&ann_vols, cl)),
            ),
        );

    Ok(StochasticResult {
        n_paths: config.n_paths,
        n_steps: config.n_steps,
        model_name: program.state_names.first().map(|_| &program.param_names)
            .map(|_| String::new())
            .unwrap_or_default(),
        final_price,
        final_return,
        max_drawdown,
        annualized_return,
        annualized_vol,
        paths,
    })
}

// ---------------------------------------------------------------------------
// Single path simulation (Euler-Maruyama)
// ---------------------------------------------------------------------------

struct PathResult {
    final_price: f64,
    max_drawdown: f64,
    sum_log_ret_sq: f64,
    path: Option<Vec<f64>>,
}

fn simulate_path(
    program: &SdeProgram,
    params: &[f64],
    s0: f64,
    n_steps: usize,
    dt: f64,
    has_jumps: bool,
    store: bool,
    rng: &mut ChaCha8Rng,
) -> PathResult {
    // Initialize state: index 0 = S
    let mut state = program.state_inits.clone();
    state[0] = s0;

    // Pre-allocate register scratch space
    let mut registers = vec![0.0_f64; program.max_registers];

    let mut path = if store {
        let mut p = Vec::with_capacity(n_steps + 1);
        p.push(s0);
        Some(p)
    } else {
        None
    };

    let sqrt_dt = dt.sqrt();
    let mut peak = s0;
    let mut max_dd = 0.0_f64;
    let mut sum_log_ret_sq = 0.0_f64;
    let mut last_ret = 0.0_f64;

    for step in 0..n_steps {
        let builtins = SdeBuiltins {
            dt,
            t: step as f64 * dt,
            step_idx: step,
            ret: last_ret,
        };

        // Evaluate drift: μ
        let drift_val = eval_tape(
            &program.drift,
            &mut registers,
            &state,
            params,
            &builtins,
            rng,
        );

        // Evaluate diffusion: σ
        let diff_val = eval_tape(
            &program.diffusion,
            &mut registers,
            &state,
            params,
            &builtins,
            rng,
        );

        // Brownian increment
        let z: f64 = {
            use rand_distr::{Distribution, Normal};
            Normal::new(0.0, 1.0).unwrap().sample(rng)
        };

        // Euler-Maruyama: log_return = (μ - σ²/2)·dt + σ·√dt·Z
        let mut log_ret = (drift_val - 0.5 * diff_val * diff_val) * dt + diff_val * sqrt_dt * z;

        // Jump component (compound Poisson)
        if has_jumps {
            if let Some(ref ji_tape) = program.jump_intensity {
                let lambda = eval_tape(ji_tape, &mut registers, &state, params, &builtins, rng);
                if lambda > 0.0 {
                    let expected = lambda * dt;
                    // For small λ·dt, use direct Bernoulli approximation
                    let n_jumps = if expected < 0.1 {
                        let u: f64 = rng.random();
                        if u < expected { 1.0 } else { 0.0 }
                    } else {
                        let pois = Poisson::new(expected).unwrap_or_else(|_| Poisson::new(0.01).unwrap());
                        pois.sample(rng)
                    };

                    if n_jumps > 0.0 {
                        if let Some(ref js_tape) = program.jump_size {
                            // Evaluate jump size for each jump
                            for _ in 0..n_jumps as usize {
                                let j = eval_tape(
                                    js_tape,
                                    &mut registers,
                                    &state,
                                    params,
                                    &builtins,
                                    rng,
                                );
                                log_ret += j;
                            }
                        }
                    }
                }
            }
        }

        // Update price
        let new_price = state[0] * log_ret.exp();
        state[0] = new_price;
        last_ret = log_ret;
        sum_log_ret_sq += log_ret * log_ret;

        // State variable updates (executed after price update)
        for &(idx, ref tape) in &program.state_updates {
            let new_val = eval_tape(tape, &mut registers, &state, params, &builtins, rng);
            state[idx] = new_val;
        }

        // Track drawdown
        if new_price > peak {
            peak = new_price;
        }
        if peak > 0.0 {
            let dd = (peak - new_price) / peak;
            max_dd = max_dd.max(dd);
        }

        if let Some(ref mut p) = path {
            p.push(new_price);
        }
    }

    PathResult {
        final_price: state[0],
        max_drawdown: max_dd,
        sum_log_ret_sq,
        path,
    }
}

// ---------------------------------------------------------------------------
// Helpers (same patterns as monte_carlo.rs)
// ---------------------------------------------------------------------------

fn path_rng(base_seed: u64, path_idx: usize) -> ChaCha8Rng {
    let seed = base_seed
        .wrapping_add(path_idx as u64)
        .wrapping_mul(0x9e3779b97f4a7c15);
    ChaCha8Rng::seed_from_u64(seed)
}

fn resolve_base_seed(seed: Option<u64>) -> u64 {
    seed.unwrap_or_else(rand::random::<u64>)
}

fn build_stats(values: &[f64], confidence_levels: &[f64]) -> DistributionStats {
    if values.is_empty() {
        return DistributionStats {
            percentiles: confidence_levels.iter().map(|&q| (q, 0.0)).collect(),
            mean: 0.0,
            std: 0.0,
            min: 0.0,
            max: 0.0,
        };
    }

    let n = values.len();
    let cmp = |a: &f64, b: &f64| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal);

    // Single parallel pass: sum + min + max
    let (sum, min, max) = values
        .par_iter()
        .fold(
            || (0.0f64, f64::INFINITY, f64::NEG_INFINITY),
            |(s, mn, mx), &v| (s + v, mn.min(v), mx.max(v)),
        )
        .reduce(
            || (0.0, f64::INFINITY, f64::NEG_INFINITY),
            |(s1, mn1, mx1), (s2, mn2, mx2)| (s1 + s2, mn1.min(mn2), mx1.max(mx2)),
        );
    let mean = sum / n as f64;
    let variance: f64 =
        values.par_iter().map(|v| (v - mean).powi(2)).sum::<f64>() / n as f64;
    let std = variance.sqrt();

    // Percentiles via select_nth_unstable: O(n) per percentile instead of O(n log n) sort.
    // Process in ascending index order so each selection shrinks the working partition.
    let mut data = values.to_vec();

    let mut targets: Vec<(usize, f64, usize)> = confidence_levels
        .iter()
        .enumerate()
        .map(|(orig, &q)| {
            let idx = ((q * (n - 1) as f64).round() as usize).min(n - 1);
            (orig, q, idx)
        })
        .collect();
    targets.sort_by_key(|&(_, _, idx)| idx);

    let mut percentiles = vec![(0.0, 0.0); confidence_levels.len()];
    let mut lo = 0usize;
    for &(orig, q, idx) in &targets {
        if idx >= lo {
            data[lo..].select_nth_unstable_by(idx - lo, cmp);
        }
        percentiles[orig] = (q, data[idx]);
        lo = idx + 1;
    }

    DistributionStats {
        percentiles,
        mean,
        std,
        min,
        max,
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use crate::stochastic::compiler::{compile_sde, SdeModelDef};
    use std::collections::HashMap;

    fn gbm_program() -> SdeProgram {
        let def = SdeModelDef {
            name: "gbm".into(),
            drift: "mu".into(),
            diffusion: "sigma".into(),
            jump_intensity: None,
            jump_size: None,
            state_vars: HashMap::new(),
            state_update: HashMap::new(),
            params: [("mu".into(), 0.05), ("sigma".into(), 0.2)]
                .into_iter()
                .collect(),
        };
        compile_sde(&def).unwrap()
    }

    #[test]
    fn test_gbm_runs() {
        let program = gbm_program();
        let config = StochasticSimConfig {
            n_paths: 100,
            n_steps: 252,
            dt: 1.0 / 252.0,
            s0: 100.0,
            confidence_levels: vec![0.05, 0.50, 0.95],
            rng_seed: Some(42),
            store_paths: false,
            precision: "f64".into(),
        };
        let result = run_stochastic(&program, &config).unwrap();
        assert_eq!(result.n_paths, 100);
        assert_eq!(result.n_steps, 252);
        assert!(result.final_price.mean > 0.0);
        assert!(result.paths.is_none());
    }

    #[test]
    fn test_gbm_deterministic() {
        let program = gbm_program();
        let config = StochasticSimConfig {
            n_paths: 50,
            n_steps: 100,
            dt: 1.0 / 252.0,
            s0: 100.0,
            confidence_levels: vec![0.50],
            rng_seed: Some(123),
            store_paths: true,
            precision: "f64".into(),
        };
        let r1 = run_stochastic(&program, &config).unwrap();
        let r2 = run_stochastic(&program, &config).unwrap();

        // Same seed → identical paths
        let p1 = r1.paths.unwrap();
        let p2 = r2.paths.unwrap();
        assert_eq!(p1.len(), p2.len());
        for (a, b) in p1.iter().zip(p2.iter()) {
            for (x, y) in a.iter().zip(b.iter()) {
                assert_eq!(x, y);
            }
        }
    }

    #[test]
    fn test_gbm_store_paths_shape() {
        let program = gbm_program();
        let config = StochasticSimConfig {
            n_paths: 10,
            n_steps: 50,
            dt: 1.0 / 252.0,
            s0: 100.0,
            confidence_levels: vec![0.50],
            rng_seed: Some(42),
            store_paths: true,
            precision: "f64".into(),
        };
        let result = run_stochastic(&program, &config).unwrap();
        let paths = result.paths.unwrap();
        assert_eq!(paths.len(), 10);
        assert_eq!(paths[0].len(), 51); // n_steps + 1 (includes S0)
        assert_eq!(paths[0][0], 100.0); // first value is S0
    }

    #[test]
    fn test_gbm_mean_convergence() {
        // With enough paths, the mean final price should be close to
        // S0 * exp(mu * T) = 100 * exp(0.05) ≈ 105.127
        let program = gbm_program();
        let config = StochasticSimConfig {
            n_paths: 1_000,
            n_steps: 252,
            dt: 1.0 / 252.0,
            s0: 100.0,
            confidence_levels: vec![0.50],
            rng_seed: Some(42),
            store_paths: false,
            precision: "f64".into(),
        };
        let result = run_stochastic(&program, &config).unwrap();
        let expected = 100.0 * (0.05_f64).exp(); // ~105.127
        let error = (result.final_price.mean - expected).abs() / expected;
        assert!(
            error < 0.10,
            "GBM mean final price {:.2} too far from expected {:.2} (error {:.1}%)",
            result.final_price.mean,
            expected,
            error * 100.0
        );
    }

    #[test]
    fn test_garch_jd_runs() {
        let def = SdeModelDef {
            name: "garch_jd".into(),
            drift: "mu".into(),
            diffusion: "sqrt(h)".into(),
            jump_intensity: Some("lambda".into()),
            jump_size: Some("normal(mu_j, sigma_j)".into()),
            state_vars: [("h".into(), 1e-4)].into_iter().collect(),
            state_update: [("h".into(), "omega + alpha * (ret - mu) ** 2 + beta * h".into())]
                .into_iter()
                .collect(),
            params: [
                ("mu".into(), 0.08),
                ("omega".into(), 1e-6),
                ("alpha".into(), 0.1),
                ("beta".into(), 0.85),
                ("lambda".into(), 5.0),
                ("mu_j".into(), -0.02),
                ("sigma_j".into(), 0.04),
            ]
            .into_iter()
            .collect(),
        };
        let program = compile_sde(&def).unwrap();
        let config = StochasticSimConfig {
            n_paths: 100,
            n_steps: 252,
            dt: 1.0 / 252.0,
            s0: 100.0,
            confidence_levels: vec![0.05, 0.50, 0.95],
            rng_seed: Some(42),
            store_paths: false,
            precision: "f64".into(),
        };
        let result = run_stochastic(&program, &config).unwrap();
        assert_eq!(result.n_paths, 100);
        // Prices should be positive
        assert!(result.final_price.min > 0.0);
    }

    #[test]
    fn test_validation_errors() {
        let program = gbm_program();

        let config = StochasticSimConfig {
            n_paths: 0,
            n_steps: 100,
            dt: 1.0 / 252.0,
            s0: 100.0,
            confidence_levels: vec![],
            rng_seed: None,
            store_paths: false,
            precision: "f64".into(),
        };
        assert!(run_stochastic(&program, &config).is_err());

        let config = StochasticSimConfig {
            n_paths: 10,
            n_steps: 100,
            dt: -1.0,
            s0: 100.0,
            confidence_levels: vec![],
            rng_seed: None,
            store_paths: false,
            precision: "f64".into(),
        };
        assert!(run_stochastic(&program, &config).is_err());
    }
}
