Documentation
use crate::strategy_engine::StrategyEngine;
use plotly::{Scatter, Plot};
use plotly::layout::Layout;

pub struct Simulator {
    engine: StrategyEngine,
}

impl Simulator {
    pub fn new(engine: StrategyEngine) -> Self {
        Self { engine }
    }

    pub fn simulate(&self, prices: Vec<f64>) {
        let results = self.engine.backtest(prices.clone());
        let indicators = self.engine.calculate_indicators(prices.clone());

        // 创建绘图
        let mut plot = Plot::new();

        // 设置布局
        let layout = Layout::new()
            .title("Strategy Performance")
            .x_axis(plotly::layout::Axis::new().title("Time"))
            .y_axis(plotly::layout::Axis::new().title("Return"));
        plot.set_layout(layout);

        // 添加每个策略的回测结果
        for (i, result) in results.iter().enumerate() {
            let trace = Scatter::new((0..prices.len()).collect(), vec![*result; prices.len()])
                .name(&format!("Strategy {}", i + 1));
            plot.add_trace(trace);
        }

        // 将图表保存为 HTML 文件
        plot.write_html("output/strategy_performance.html");

        // 打印指标
        for (i, indicator) in indicators.iter().enumerate() {
            println!("Strategy {}: Indicator = {}", i + 1, indicator);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::strategy_engine::StrategyEngine;
    use crate::trend_following::TrendFollowing;
    use crate::moving_average_cross::MovingAverageCross;
    use crate::bollinger_bands::BollingerBands;
    use crate::mean_reversion::MeanReversion;
    use crate::momentum::Momentum;
    use crate::arbitrage::Arbitrage;

    #[test]
    fn test_simulator_with_multiple_strategies() {
        let mut engine = StrategyEngine::new();

        let mut trend_following = TrendFollowing::new(3, 6);
        let mut moving_average_cross = MovingAverageCross::new(3, 6);
        let mut bollinger_bands = BollingerBands::new(10);
        let mut mean_reversion = MeanReversion::new(5);
        let mut momentum = Momentum::new(3);
        let arbitrage = Arbitrage::new();

        engine.add_strategy(Box::new(trend_following.clone()));
        engine.add_strategy(Box::new(moving_average_cross.clone()));
        engine.add_strategy(Box::new(bollinger_bands.clone()));
        engine.add_strategy(Box::new(mean_reversion.clone()));
        engine.add_strategy(Box::new(momentum.clone()));
        engine.add_strategy(Box::new(arbitrage.clone()));

        let prices = vec![10.0, 12.0, 15.0, 14.0, 13.0, 12.0];

        for price in &prices {
            trend_following.update_price(*price);
            moving_average_cross.update_price(*price);
            bollinger_bands.update_price(*price);
            mean_reversion.update_price(*price);
            momentum.update_price(*price);
        }

        let simulator = Simulator::new(engine);
        simulator.simulate(prices.clone());
    }

    #[test]
    fn test_simulator_with_empty_strategies() {
        let engine = StrategyEngine::new();
        let simulator = Simulator::new(engine);

        let prices = vec![10.0, 12.0, 15.0, 14.0, 13.0, 12.0];

        simulator.simulate(prices);
    }

    #[test]
    fn test_simulator_with_single_strategy() {
        let mut engine = StrategyEngine::new();
        let mut trend_following = TrendFollowing::new(3, 6);

        engine.add_strategy(Box::new(trend_following.clone()));

        let prices = vec![10.0, 12.0, 15.0, 14.0, 13.0, 12.0];

        for price in &prices {
            trend_following.update_price(*price);
        }

        let simulator = Simulator::new(engine);
        simulator.simulate(prices.clone());
    }
}