//! Scenario generation using world models.
//!
//! Generates multiple price paths for market scenarios.

use crate::MarketSimError;
use crate::world_model::MarketWorldModel;
use marketsim_shared::{
    AssetPathStats, AssetPricePaths, PathStatistics, ScenarioRequest, ScenarioResult,
    SimulationMetadata,
};
use std::time::Instant;

/// Scenario generator.
#[derive(Debug)]
pub struct ScenarioGenerator {
    /// Default number of paths if not specified.
    #[allow(dead_code)]
    default_paths: usize,
}

impl Default for ScenarioGenerator {
    fn default() -> Self {
        Self::new()
    }
}

impl ScenarioGenerator {
    /// Create a new generator.
    #[must_use]
    pub fn new() -> Self {
        Self { default_paths: 100 }
    }

    /// Generate scenarios.
    pub fn generate(
        &self,
        request: &ScenarioRequest,
        world_model: &MarketWorldModel,
    ) -> Result<ScenarioResult, MarketSimError> {
        let start = Instant::now();

        // Get scenario-adjusted parameters
        let params = world_model
            .get_scenario_params(request.scenario.scenario_type, request.scenario.severity);

        // Initialize RNG with seed
        let seed = request.seed.unwrap_or(42);
        let mut rng = SimpleRng::new(seed);

        // Time step (daily)
        let dt = 1.0 / 252.0;

        // Generate paths for each asset
        let mut price_paths = Vec::with_capacity(request.assets.len());

        for asset in &request.assets {
            // Get initial price from state
            let initial_price = request
                .initial_state
                .prices
                .iter()
                .find(|p| p.symbol == asset.symbol)
                .map_or(asset.current_price, |p| p.price);

            // Generate paths
            let mut paths = Vec::with_capacity(request.num_paths);
            let mut return_paths = Vec::with_capacity(request.num_paths);

            for _ in 0..request.num_paths {
                let mut path = Vec::with_capacity(request.horizon_days + 1);
                let mut returns = Vec::with_capacity(request.horizon_days);

                path.push(initial_price);

                for _ in 0..request.horizon_days {
                    let z = rng.normal();
                    let z_jump = rng.normal();
                    let u_jump = rng.uniform();

                    let new_price =
                        world_model.step(*path.last().unwrap(), &params, dt, z, z_jump, u_jump);

                    let ret = (new_price / path.last().unwrap()).ln();
                    returns.push(ret);
                    path.push(new_price);
                }

                paths.push(path);
                return_paths.push(returns);
            }

            price_paths.push(AssetPricePaths {
                symbol: asset.symbol.clone(),
                paths,
                return_paths,
            });
        }

        // Calculate statistics
        let statistics = self.calculate_statistics(&price_paths, request.horizon_days);

        let computation_time = start.elapsed().as_millis() as u64;

        Ok(ScenarioResult {
            scenario: request.scenario.clone(),
            price_paths,
            statistics,
            metadata: SimulationMetadata {
                computation_time_ms: computation_time,
                model_version: "1.0.0".to_string(),
                seed_used: seed,
            },
        })
    }

    /// Calculate statistics from price paths.
    fn calculate_statistics(
        &self,
        price_paths: &[AssetPricePaths],
        horizon: usize,
    ) -> PathStatistics {
        let mut asset_stats = Vec::with_capacity(price_paths.len());

        for asset_paths in price_paths {
            // Collect final prices
            let mut final_prices: Vec<f64> = asset_paths.paths.iter().map(|p| p[horizon]).collect();

            final_prices.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

            let n = final_prices.len();
            let initial_price = asset_paths.paths[0][0];

            // Calculate returns
            let returns: Vec<f64> = final_prices
                .iter()
                .map(|p| (p / initial_price).ln())
                .collect();

            let mean_return: f64 = returns.iter().sum::<f64>() / n as f64;

            // Volatility from return paths
            let all_returns: Vec<f64> = asset_paths
                .return_paths
                .iter()
                .flat_map(|r| r.iter())
                .copied()
                .collect();

            let mean_daily: f64 = all_returns.iter().sum::<f64>() / all_returns.len() as f64;
            let variance: f64 = all_returns
                .iter()
                .map(|r| (r - mean_daily).powi(2))
                .sum::<f64>()
                / all_returns.len() as f64;
            let volatility = variance.sqrt() * (252.0_f64).sqrt();

            // Max drawdown
            let mean_max_dd: f64 = asset_paths
                .paths
                .iter()
                .map(|path| self.max_drawdown(path))
                .sum::<f64>()
                / n as f64;

            asset_stats.push(AssetPathStats {
                symbol: asset_paths.symbol.clone(),
                mean_final_price: final_prices.iter().sum::<f64>() / n as f64,
                median_final_price: final_prices[n / 2],
                pct_5_final_price: final_prices[(0.05 * n as f64) as usize],
                pct_95_final_price: final_prices[(0.95 * n as f64).min((n - 1) as f64) as usize],
                mean_return,
                volatility,
                mean_max_drawdown: mean_max_dd,
            });
        }

        // Correlation at final time (simplified - just returning identity for now)
        let n_assets = price_paths.len();
        let final_correlation = (0..n_assets)
            .map(|i| {
                (0..n_assets)
                    .map(|j| if i == j { 1.0 } else { 0.5 })
                    .collect()
            })
            .collect();

        PathStatistics {
            asset_stats,
            final_correlation,
        }
    }

    /// Calculate maximum drawdown for a price path.
    fn max_drawdown(&self, path: &[f64]) -> f64 {
        let mut peak: f64 = path[0];
        let mut max_dd: f64 = 0.0;

        for &price in path.iter().skip(1) {
            peak = peak.max(price);
            let dd = (peak - price) / peak;
            max_dd = max_dd.max(dd);
        }

        max_dd
    }
}

/// Simple pseudo-random number generator.
struct SimpleRng {
    state: u64,
}

impl SimpleRng {
    fn new(seed: u64) -> Self {
        Self { state: seed }
    }

    fn next(&mut self) -> u64 {
        self.state = self
            .state
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        self.state
    }

    fn uniform(&mut self) -> f64 {
        (self.next() >> 11) as f64 / (1u64 << 53) as f64
    }

    fn normal(&mut self) -> f64 {
        // Box-Muller transform
        let u1 = self.uniform() + 1e-10;
        let u2 = self.uniform();
        (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use marketsim_shared::{
        Asset, AssetPrice, AssetType, EconomicFactors, MarketIndicators, MarketState,
        ScenarioDescription, ScenarioType,
    };

    fn create_test_request() -> ScenarioRequest {
        ScenarioRequest {
            assets: vec![Asset {
                symbol: "SPY".to_string(),
                name: "S&P 500".to_string(),
                asset_type: AssetType::Index,
                current_price: 500.0,
            }],
            historical_data: vec![],
            initial_state: MarketState {
                timestamp: "2024-01-01T00:00:00Z".to_string(),
                prices: vec![AssetPrice {
                    symbol: "SPY".to_string(),
                    price: 500.0,
                    change_pct: 0.0,
                }],
                indicators: MarketIndicators::default(),
                economic_factors: EconomicFactors::default(),
            },
            scenario: ScenarioDescription {
                name: "Test".to_string(),
                scenario_type: ScenarioType::Crisis,
                description: "Test".to_string(),
                severity: 0.5,
                duration_days: 20,
            },
            num_paths: 100,
            horizon_days: 20,
            seed: Some(42),
        }
    }

    #[test]
    fn test_generator_creation() {
        let generator = ScenarioGenerator::new();
        assert!(std::mem::size_of_val(&generator) > 0);
    }

    #[test]
    fn test_generate() {
        let generator = ScenarioGenerator::new();
        let model = MarketWorldModel::new();
        let request = create_test_request();

        let result = generator.generate(&request, &model);
        assert!(result.is_ok());

        let scenario = result.unwrap();
        assert_eq!(scenario.price_paths.len(), 1);
        assert_eq!(scenario.price_paths[0].paths.len(), 100);
        assert_eq!(scenario.price_paths[0].paths[0].len(), 21); // 20 days + initial
    }

    #[test]
    fn test_statistics() {
        let generator = ScenarioGenerator::new();
        let model = MarketWorldModel::new();
        let request = create_test_request();

        let scenario = generator.generate(&request, &model).unwrap();

        assert_eq!(scenario.statistics.asset_stats.len(), 1);
        assert!(scenario.statistics.asset_stats[0].volatility > 0.0);
    }

    #[test]
    fn test_max_drawdown() {
        let generator = ScenarioGenerator::new();
        let path = vec![100.0, 110.0, 105.0, 95.0, 100.0];
        let dd = generator.max_drawdown(&path);

        // Max drawdown from 110 to 95 = (110-95)/110 ≈ 0.136
        assert!((dd - 0.136).abs() < 0.01);
    }

    #[test]
    fn test_reproducibility() {
        let generator = ScenarioGenerator::new();
        let model = MarketWorldModel::new();
        let request = create_test_request();

        let result1 = generator.generate(&request, &model).unwrap();
        let result2 = generator.generate(&request, &model).unwrap();

        // With same seed, should get same results
        assert_eq!(
            result1.price_paths[0].paths[0][5],
            result2.price_paths[0].paths[0][5]
        );
    }
}