//! Backtesting engine for trading strategies.
//!
//! Tests strategies against generated scenarios.

use crate::MarketSimError;
use marketsim_shared::{
    AggregateBacktestStats, BacktestRequest, BacktestResult, ScenarioBacktestResult,
    ScenarioResult, Strategy, StrategyType,
};

/// Backtest engine.
#[derive(Debug)]
pub struct BacktestEngine {
    /// Risk-free rate for Sharpe calculation.
    risk_free_rate: f64,
}

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

impl BacktestEngine {
    /// Create a new backtest engine.
    #[must_use]
    pub fn new() -> Self {
        Self {
            risk_free_rate: 0.04,
        }
    }

    /// Run backtest.
    pub fn run(&self, request: &BacktestRequest) -> Result<BacktestResult, MarketSimError> {
        let mut scenario_results = Vec::with_capacity(request.scenarios.len());

        for scenario in &request.scenarios {
            let result = self.backtest_scenario(&request.strategy, scenario, request)?;
            scenario_results.push(result);
        }

        // Calculate aggregate statistics
        let aggregate_stats = self.calculate_aggregate_stats(&scenario_results);

        Ok(BacktestResult {
            strategy: request.strategy.clone(),
            scenario_results,
            aggregate_stats,
        })
    }

    /// Backtest a single scenario.
    fn backtest_scenario(
        &self,
        strategy: &Strategy,
        scenario: &ScenarioResult,
        request: &BacktestRequest,
    ) -> Result<ScenarioBacktestResult, MarketSimError> {
        if scenario.price_paths.is_empty() {
            return Err(MarketSimError::BacktestFailed("No price paths".to_string()));
        }

        let num_paths = scenario.price_paths[0].paths.len();
        let _num_assets = scenario.price_paths.len();
        let _horizon = scenario.price_paths[0].paths[0].len() - 1;

        let mut final_values = Vec::with_capacity(num_paths);
        let mut total_returns = Vec::with_capacity(num_paths);
        let mut sharpe_ratios = Vec::with_capacity(num_paths);
        let mut max_drawdowns = Vec::with_capacity(num_paths);

        for path_idx in 0..num_paths {
            // Get prices for this path
            let prices: Vec<Vec<f64>> = scenario
                .price_paths
                .iter()
                .map(|ap| ap.paths[path_idx].clone())
                .collect();

            // Run strategy
            let (portfolio_values, returns) = self.run_strategy(
                strategy,
                &prices,
                request.initial_capital,
                request.transaction_cost_bps,
            );

            let final_value = *portfolio_values.last().unwrap();
            let total_return = (final_value / request.initial_capital) - 1.0;

            // Calculate Sharpe ratio
            let mean_return: f64 = returns.iter().sum::<f64>() / returns.len() as f64;
            let variance: f64 = returns
                .iter()
                .map(|r| (r - mean_return).powi(2))
                .sum::<f64>()
                / returns.len() as f64;
            let std_dev = variance.sqrt();
            let annualized_return = mean_return * 252.0;
            let annualized_std = std_dev * (252.0_f64).sqrt();
            let sharpe = if annualized_std > 0.0 {
                (annualized_return - self.risk_free_rate) / annualized_std
            } else {
                0.0
            };

            // Calculate max drawdown
            let max_dd = self.calculate_max_drawdown(&portfolio_values);

            final_values.push(final_value);
            total_returns.push(total_return);
            sharpe_ratios.push(sharpe);
            max_drawdowns.push(max_dd);
        }

        let mean_final_value = final_values.iter().sum::<f64>() / num_paths as f64;
        let mean_total_return = total_returns.iter().sum::<f64>() / num_paths as f64;
        let mean_sharpe = sharpe_ratios.iter().sum::<f64>() / num_paths as f64;
        let mean_max_dd = max_drawdowns.iter().sum::<f64>() / num_paths as f64;
        let prob_loss =
            total_returns.iter().filter(|&&r| r < 0.0).count() as f64 / num_paths as f64;

        Ok(ScenarioBacktestResult {
            scenario_name: scenario.scenario.name.clone(),
            mean_final_value,
            mean_total_return,
            mean_sharpe_ratio: mean_sharpe,
            mean_max_drawdown: mean_max_dd,
            probability_of_loss: prob_loss,
        })
    }

    /// Run strategy on price paths.
    fn run_strategy(
        &self,
        strategy: &Strategy,
        prices: &[Vec<f64>],
        initial_capital: f64,
        transaction_cost_bps: f64,
    ) -> (Vec<f64>, Vec<f64>) {
        let num_assets = prices.len();
        let horizon = prices[0].len();

        // Initialize equal weights
        let mut weights: Vec<f64> = vec![1.0 / num_assets as f64; num_assets];
        let mut portfolio_value: f64 = initial_capital;
        let mut portfolio_values = Vec::with_capacity(horizon);
        let mut returns = Vec::with_capacity(horizon - 1);

        portfolio_values.push(portfolio_value);

        for t in 1..horizon {
            // Calculate asset returns
            let asset_returns: Vec<f64> = (0..num_assets)
                .map(|i| (prices[i][t] / prices[i][t - 1]) - 1.0)
                .collect();

            // Portfolio return
            let portfolio_return: f64 = weights
                .iter()
                .zip(asset_returns.iter())
                .map(|(w, r)| w * r)
                .sum();

            // Update portfolio value
            portfolio_value *= 1.0 + portfolio_return;

            // Rebalance if needed
            if t % strategy.parameters.rebalance_frequency == 0 {
                let new_weights = self.calculate_weights(strategy, prices, t);

                // Transaction costs
                let turnover: f64 = weights
                    .iter()
                    .zip(new_weights.iter())
                    .map(|(old, new)| (new - old).abs())
                    .sum::<f64>()
                    / 2.0;

                let cost = turnover * transaction_cost_bps / 10000.0;
                portfolio_value *= 1.0 - cost;

                weights = new_weights;
            }

            returns.push(portfolio_return);
            portfolio_values.push(portfolio_value);
        }

        (portfolio_values, returns)
    }

    /// Calculate strategy weights.
    fn calculate_weights(&self, strategy: &Strategy, prices: &[Vec<f64>], t: usize) -> Vec<f64> {
        let num_assets = prices.len();
        let lookback = strategy.parameters.lookback.min(t);

        match strategy.strategy_type {
            StrategyType::BuyAndHold => {
                // Equal weight
                vec![1.0 / num_assets as f64; num_assets]
            }
            StrategyType::Momentum => {
                // Weight by past returns
                let mut momentum: Vec<f64> = (0..num_assets)
                    .map(|i| {
                        if t > lookback && prices[i][t - lookback] > 0.0 {
                            (prices[i][t] / prices[i][t - lookback]) - 1.0
                        } else {
                            0.0
                        }
                    })
                    .collect();

                // Normalize to positive weights
                let min_mom = momentum.iter().copied().fold(f64::INFINITY, f64::min);
                for m in &mut momentum {
                    *m -= min_mom - 0.001;
                }

                let sum: f64 = momentum.iter().sum();
                if sum > 0.0 {
                    momentum.iter().map(|m| m / sum).collect()
                } else {
                    vec![1.0 / num_assets as f64; num_assets]
                }
            }
            StrategyType::MeanReversion => {
                // Weight inversely by past returns
                let mut z_scores: Vec<f64> = (0..num_assets)
                    .map(|i| {
                        if t > lookback {
                            let returns: Vec<f64> = (t - lookback..t)
                                .map(|j| (prices[i][j + 1] / prices[i][j]) - 1.0)
                                .collect();
                            let mean: f64 = returns.iter().sum::<f64>() / returns.len() as f64;
                            -mean // Negative = mean reversion
                        } else {
                            0.0
                        }
                    })
                    .collect();

                // Normalize
                let min_z = z_scores.iter().copied().fold(f64::INFINITY, f64::min);
                for z in &mut z_scores {
                    *z -= min_z - 0.001;
                }

                let sum: f64 = z_scores.iter().sum();
                if sum > 0.0 {
                    z_scores.iter().map(|z| z / sum).collect()
                } else {
                    vec![1.0 / num_assets as f64; num_assets]
                }
            }
            StrategyType::RiskParity => {
                // Weight inversely by volatility
                let inv_vol: Vec<f64> = (0..num_assets)
                    .map(|i| {
                        if t > lookback {
                            let returns: Vec<f64> = (t - lookback..t)
                                .map(|j| (prices[i][j + 1] / prices[i][j]) - 1.0)
                                .collect();
                            let mean: f64 = returns.iter().sum::<f64>() / returns.len() as f64;
                            let variance: f64 =
                                returns.iter().map(|r| (r - mean).powi(2)).sum::<f64>()
                                    / returns.len() as f64;
                            let vol = variance.sqrt();
                            if vol > 0.001 { 1.0 / vol } else { 1.0 }
                        } else {
                            1.0
                        }
                    })
                    .collect();

                let sum: f64 = inv_vol.iter().sum();
                inv_vol.iter().map(|v| v / sum).collect()
            }
            StrategyType::TrendFollowing => {
                // Binary: 100% if above MA, 0% if below
                let short_lookback = lookback / 2;
                let weights: Vec<f64> = (0..num_assets)
                    .map(|i| {
                        if t > lookback {
                            let long_ma: f64 = (t - lookback..t).map(|j| prices[i][j]).sum::<f64>()
                                / lookback as f64;
                            let short_ma: f64 = if t > short_lookback {
                                (t - short_lookback..t).map(|j| prices[i][j]).sum::<f64>()
                                    / short_lookback as f64
                            } else {
                                prices[i][t]
                            };

                            if short_ma > long_ma {
                                1.0
                            } else {
                                0.1 // Minimum weight
                            }
                        } else {
                            1.0
                        }
                    })
                    .collect();

                let sum: f64 = weights.iter().sum();
                weights.iter().map(|w| w / sum).collect()
            }
            StrategyType::Custom => {
                // Default to equal weight
                vec![1.0 / num_assets as f64; num_assets]
            }
        }
    }

    /// Calculate max drawdown.
    fn calculate_max_drawdown(&self, values: &[f64]) -> f64 {
        let mut peak: f64 = values[0];
        let mut max_dd: f64 = 0.0;

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

        max_dd
    }

    /// Calculate aggregate statistics.
    fn calculate_aggregate_stats(
        &self,
        results: &[ScenarioBacktestResult],
    ) -> AggregateBacktestStats {
        let n = results.len() as f64;

        let overall_mean_return = results.iter().map(|r| r.mean_total_return).sum::<f64>() / n;
        let overall_sharpe = results.iter().map(|r| r.mean_sharpe_ratio).sum::<f64>() / n;
        let worst_case = results
            .iter()
            .map(|r| r.mean_total_return)
            .fold(f64::INFINITY, f64::min);
        let best_case = results
            .iter()
            .map(|r| r.mean_total_return)
            .fold(f64::NEG_INFINITY, f64::max);
        let win_rate = results.iter().filter(|r| r.mean_total_return > 0.0).count() as f64 / n;

        AggregateBacktestStats {
            overall_mean_return,
            overall_sharpe_ratio: overall_sharpe,
            worst_case_return: worst_case,
            best_case_return: best_case,
            win_rate,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use marketsim_shared::{
        AssetPricePaths, PathStatistics, ScenarioDescription, ScenarioType, SimulationMetadata,
        Strategy, StrategyParameters,
    };

    fn create_test_scenario() -> ScenarioResult {
        // Create simple price paths: 10 paths, 20 steps
        let paths: Vec<Vec<f64>> = (0..10)
            .map(|i| {
                let mut path = vec![100.0];
                for t in 1..21 {
                    // Simple random walk
                    let change = 1.0 + 0.01 * (i as f64 - 5.0) * 0.1;
                    path.push(path.last().unwrap() * change);
                }
                path
            })
            .collect();

        let return_paths: Vec<Vec<f64>> = paths
            .iter()
            .map(|p| p.windows(2).map(|w| (w[1] / w[0]).ln()).collect())
            .collect();

        ScenarioResult {
            scenario: ScenarioDescription {
                name: "Test".to_string(),
                scenario_type: ScenarioType::Crisis,
                description: "Test".to_string(),
                severity: 0.5,
                duration_days: 20,
            },
            price_paths: vec![AssetPricePaths {
                symbol: "SPY".to_string(),
                paths,
                return_paths,
            }],
            statistics: PathStatistics {
                asset_stats: vec![],
                final_correlation: vec![],
            },
            metadata: SimulationMetadata {
                computation_time_ms: 1,
                model_version: "1.0".to_string(),
                seed_used: 42,
            },
        }
    }

    #[test]
    fn test_backtest_engine_creation() {
        let engine = BacktestEngine::new();
        assert!(engine.risk_free_rate > 0.0);
    }

    #[test]
    fn test_run_backtest() {
        let engine = BacktestEngine::new();
        let scenario = create_test_scenario();

        let request = BacktestRequest {
            strategy: Strategy {
                name: "Momentum".to_string(),
                strategy_type: StrategyType::Momentum,
                parameters: StrategyParameters::default(),
            },
            scenarios: vec![scenario],
            initial_capital: 100000.0,
            transaction_cost_bps: 10.0,
        };

        let result = engine.run(&request);
        assert!(result.is_ok());

        let backtest = result.unwrap();
        assert_eq!(backtest.scenario_results.len(), 1);
    }

    #[test]
    fn test_strategy_types() {
        let engine = BacktestEngine::new();
        let scenario = create_test_scenario();

        for strategy_type in [
            StrategyType::BuyAndHold,
            StrategyType::Momentum,
            StrategyType::MeanReversion,
            StrategyType::RiskParity,
        ] {
            let request = BacktestRequest {
                strategy: Strategy {
                    name: "Test".to_string(),
                    strategy_type,
                    parameters: StrategyParameters::default(),
                },
                scenarios: vec![scenario.clone()],
                initial_capital: 100000.0,
                transaction_cost_bps: 10.0,
            };

            let result = engine.run(&request);
            assert!(result.is_ok());
        }
    }

    #[test]
    fn test_max_drawdown() {
        let engine = BacktestEngine::new();
        let values = vec![100.0, 110.0, 105.0, 95.0, 100.0];
        let dd = engine.calculate_max_drawdown(&values);

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