rustytorch/demos/rtx-marketsim-demo/src/lib.rs

//! MarketSim - Financial World Model Demo.
//!
//! This demo showcases market scenario generation using world models,
//! enabling stress testing and scenario analysis for portfolios.

pub mod backtest;
pub mod sample_data;
pub mod scenario_generator;
pub mod world_model;

use marketsim_shared::{BacktestRequest, BacktestResult, ScenarioRequest, ScenarioResult};
use thiserror::Error;

use backtest::BacktestEngine;
use scenario_generator::ScenarioGenerator;
use world_model::MarketWorldModel;

/// Errors that can occur during market simulation.
#[derive(Debug, Error)]
pub enum MarketSimError {
    /// No assets provided.
    #[error("No assets provided for simulation")]
    NoAssets,

    /// No historical data provided.
    #[error("No historical data provided")]
    NoHistoricalData,

    /// Invalid scenario parameters.
    #[error("Invalid scenario parameters: {0}")]
    InvalidScenario(String),

    /// Simulation failed.
    #[error("Simulation failed: {0}")]
    SimulationFailed(String),

    /// Backtest failed.
    #[error("Backtest failed: {0}")]
    BacktestFailed(String),
}

/// Main MarketSim system.
#[derive(Debug)]
pub struct MarketSim {
    /// World model for market dynamics.
    world_model: MarketWorldModel,
    /// Scenario generator.
    scenario_gen: ScenarioGenerator,
    /// Backtest engine.
    backtest_engine: BacktestEngine,
}

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

impl MarketSim {
    /// Create a new MarketSim instance.
    #[must_use]
    pub fn new() -> Self {
        Self {
            world_model: MarketWorldModel::new(),
            scenario_gen: ScenarioGenerator::new(),
            backtest_engine: BacktestEngine::new(),
        }
    }

    /// Generate market scenarios.
    pub fn generate_scenarios(
        &self,
        request: &ScenarioRequest,
    ) -> Result<ScenarioResult, MarketSimError> {
        // Validate request
        if request.assets.is_empty() {
            return Err(MarketSimError::NoAssets);
        }

        if request.num_paths == 0 {
            return Err(MarketSimError::InvalidScenario(
                "num_paths must be > 0".to_string(),
            ));
        }

        if request.horizon_days == 0 {
            return Err(MarketSimError::InvalidScenario(
                "horizon_days must be > 0".to_string(),
            ));
        }

        // Generate scenarios using world model
        self.scenario_gen.generate(request, &self.world_model)
    }

    /// Run backtest on generated scenarios.
    pub fn run_backtest(
        &self,
        request: &BacktestRequest,
    ) -> Result<BacktestResult, MarketSimError> {
        if request.scenarios.is_empty() {
            return Err(MarketSimError::BacktestFailed(
                "No scenarios provided".to_string(),
            ));
        }

        self.backtest_engine.run(request)
    }

    /// Get the world model.
    #[must_use]
    pub fn world_model(&self) -> &MarketWorldModel {
        &self.world_model
    }
}

/// Run the full demo.
pub fn run_demo() -> Result<(ScenarioResult, BacktestResult), MarketSimError> {
    let sim = MarketSim::new();

    // Generate crisis scenario
    let request = sample_data::create_crisis_simulation_request();
    let scenario_result = sim.generate_scenarios(&request)?;

    // Run backtest with momentum strategy
    let backtest_request = sample_data::create_backtest_request(vec![scenario_result.clone()]);
    let backtest_result = sim.run_backtest(&backtest_request)?;

    Ok((scenario_result, backtest_result))
}

#[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,
                },
                Asset {
                    symbol: "TLT".to_string(),
                    name: "Treasury".to_string(),
                    asset_type: AssetType::Bond,
                    current_price: 90.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,
                    },
                    AssetPrice {
                        symbol: "TLT".to_string(),
                        price: 90.0,
                        change_pct: 0.0,
                    },
                ],
                indicators: MarketIndicators::default(),
                economic_factors: EconomicFactors::default(),
            },
            scenario: ScenarioDescription {
                name: "Test Crisis".to_string(),
                scenario_type: ScenarioType::Crisis,
                description: "Test".to_string(),
                severity: 0.5,
                duration_days: 20,
            },
            num_paths: 10,
            horizon_days: 20,
            seed: Some(42),
        }
    }

    #[test]
    fn test_marketsim_creation() {
        let sim = MarketSim::new();
        assert!(std::mem::size_of_val(&sim) > 0);
    }

    #[test]
    fn test_generate_scenarios() {
        let sim = MarketSim::new();
        let request = create_test_request();

        let result = sim.generate_scenarios(&request);
        assert!(result.is_ok());

        let scenario = result.unwrap();
        assert_eq!(scenario.price_paths.len(), 2);
        assert_eq!(scenario.price_paths[0].paths.len(), 10);
    }

    #[test]
    fn test_no_assets_error() {
        let sim = MarketSim::new();
        let mut request = create_test_request();
        request.assets.clear();

        let result = sim.generate_scenarios(&request);
        assert!(matches!(result, Err(MarketSimError::NoAssets)));
    }

    #[test]
    fn test_invalid_num_paths() {
        let sim = MarketSim::new();
        let mut request = create_test_request();
        request.num_paths = 0;

        let result = sim.generate_scenarios(&request);
        assert!(matches!(result, Err(MarketSimError::InvalidScenario(_))));
    }

    #[test]
    fn test_run_demo() {
        let result = run_demo();
        assert!(result.is_ok());

        let (scenario, backtest) = result.unwrap();
        assert!(!scenario.price_paths.is_empty());
        assert!(!backtest.scenario_results.is_empty());
    }
}