foxhunt/ml/examples/stress_test_dqn.rs

//! DQN Stress Testing CLI
//!
//! Command-line interface for running comprehensive stress tests on DQN models.
//! Supports 8 predefined scenarios and custom scenario configuration.
//!
//! **Usage**:
//! ```bash
//! # Run all predefined scenarios
//! cargo run -p ml --example stress_test_dqn --release --features cuda
//!
//! # Run specific scenario
//! cargo run -p ml --example stress_test_dqn --release --features cuda -- \
//!   --scenario flash_crash
//!
//! # Custom scenario configuration
//! cargo run -p ml --example stress_test_dqn --release --features cuda -- \
//!   --price-shock -15.0 \
//!   --volatility 8.0 \
//!   --duration 600
//! ```

use anyhow::{Context, Result};
use clap::Parser;
use ml::dqn::stress_testing::{
    correlation_breakdown_scenario, flash_crash_scenario, gap_risk_scenario,
    liquidity_crisis_scenario, multi_asset_stress_scenario, trending_market_scenario,
    vix_spike_scenario, whipsaw_scenario, DQNStressTester, StressScenario,
};
use ml::dqn::dqn::WorkingDQNConfig;
use ml::trainers::{DQNHyperparameters, DQNTrainer, TargetUpdateMode};
use std::path::PathBuf;
use tracing::{info, Level};
use tracing_subscriber;

#[derive(Parser, Debug)]
#[command(name = "DQN Stress Test")]
#[command(about = "Run stress tests on DQN trading models", long_about = None)]
struct Args {
    /// Specific scenario to run (flash_crash, liquidity_crisis, vix_spike, trending, whipsaw, gap_risk, correlation_breakdown, multi_asset)
    #[arg(short, long)]
    scenario: Option<String>,

    /// Custom price shock percentage (negative = crash)
    #[arg(long)]
    price_shock: Option<f64>,

    /// Custom volatility multiplier
    #[arg(long)]
    volatility: Option<f64>,

    /// Custom spread multiplier
    #[arg(long)]
    spread: Option<f64>,

    /// Custom duration in trading steps
    #[arg(long)]
    duration: Option<usize>,

    /// Path to trained DQN model (optional)
    #[arg(short, long)]
    model_path: Option<PathBuf>,

    /// Number of training epochs for stress test
    #[arg(short, long, default_value = "10")]
    epochs: usize,

    /// Output report to file (JSON format)
    #[arg(short, long)]
    output: Option<PathBuf>,

    /// Verbose logging
    #[arg(short, long)]
    verbose: bool,
}

fn main() -> Result<()> {
    let args = Args::parse();

    // Initialize logging
    let log_level = if args.verbose {
        Level::DEBUG
    } else {
        Level::INFO
    };
    tracing_subscriber::fmt().with_max_level(log_level).init();

    info!("DQN Stress Testing CLI");
    info!("=".repeat(80));

    // Initialize DQN trainer
    let trainer = create_dqn_trainer(args.epochs, args.model_path.as_deref())?;

    // Create stress tester
    let mut stress_tester = DQNStressTester::new(trainer);

    // Determine which scenarios to run
    let results = if let Some(ref scenario_name) = args.scenario {
        // Run single specified scenario
        let scenario = get_scenario_by_name(scenario_name)?;
        info!("Running single scenario: {}", scenario.name);
        vec![stress_tester.run_scenario(&scenario)?]
    } else if args.price_shock.is_some()
        || args.volatility.is_some()
        || args.spread.is_some()
        || args.duration.is_some()
    {
        // Run custom scenario
        let scenario = StressScenario {
            name: "Custom Scenario".to_string(),
            price_shock_pct: args.price_shock.unwrap_or(-10.0),
            volatility_multiplier: args.volatility.unwrap_or(3.0),
            spread_multiplier: args.spread.unwrap_or(10.0),
            duration_steps: args.duration.unwrap_or(300),
            max_drawdown_threshold: 25.0,
            min_action_diversity: 20.0,
        };
        info!("Running custom scenario: {}", scenario.name);
        vec![stress_tester.run_scenario(&scenario)?]
    } else {
        // Run all predefined scenarios
        info!("Running all 8 predefined scenarios");
        stress_tester.run_all_scenarios()
    };

    // Generate comprehensive report
    let report = if results.len() > 1 {
        let full_report = stress_tester.generate_report();
        full_report.print_summary();
        full_report
    } else {
        // Single scenario - print detailed result
        if let Some(result) = results.first() {
            print_detailed_result(result);
        }
        // Create minimal report for consistency
        ml::dqn::stress_testing::StressTestReport {
            total_scenarios: results.len(),
            passed: results.iter().filter(|r| r.passed).count(),
            failed: results.iter().filter(|r| !r.passed).count(),
            avg_max_drawdown: results.first().map(|r| r.max_drawdown).unwrap_or(0.0),
            avg_action_diversity: results.first().map(|r| r.action_diversity).unwrap_or(0.0),
            worst_scenario: results.first().map(|r| r.scenario_name.clone()),
            results,
        }
    };

    // Save report to file if requested
    if let Some(output_path) = args.output {
        save_report_to_file(&report, &output_path)?;
        info!("Report saved to: {}", output_path.display());
    }

    // Exit with appropriate code
    if report.failed > 0 {
        std::process::exit(1);
    } else {
        std::process::exit(0);
    }
}

/// Create DQN trainer for stress testing
fn create_dqn_trainer(epochs: usize, model_path: Option<&std::path::Path>) -> Result<DQNTrainer> {
    let hyperparams = DQNHyperparameters {
        learning_rate: 1e-4,
        batch_size: 128,
        gamma: 0.99,
        epsilon_start: 0.3, // Higher exploration for stress tests
        epsilon_end: 0.05,
        epsilon_decay: 0.995,
        buffer_size: 50000,
        min_replay_size: 1000,
        epochs,
        checkpoint_frequency: 10,
        early_stopping_enabled: false, // No early stopping during stress tests
        q_value_floor: 0.5,
        min_loss_improvement_pct: 2.0,
        plateau_window: 30,
        min_epochs_before_stopping: epochs, // Disable early stopping
        hold_penalty: 0.01,
        use_huber_loss: true,
        huber_delta: 1.0,
        use_double_dqn: true,
        gradient_clip_norm: Some(10.0),
        hold_penalty_weight: 1.0,
        movement_threshold: 0.02,
        enable_preprocessing: true,
        preprocessing_window: 50,
        preprocessing_clip_sigma: 5.0,
        tau: 0.001,
        target_update_mode: TargetUpdateMode::Soft,
        target_update_frequency: 10000,
    };

    let dqn_config = WorkingDQNConfig {
        input_dim: 128, // 125 market features + 3 portfolio features
        hidden_dims: vec![256, 128, 64],
        output_dim: 45, // 45-action factored space
        learning_rate: hyperparams.learning_rate,
        gamma: hyperparams.gamma,
        use_double_dqn: hyperparams.use_double_dqn,
        tau: hyperparams.tau,
    };

    let trainer = DQNTrainer::new(hyperparams, dqn_config)
        .context("Failed to create DQN trainer")?;

    // Load model if path provided
    if let Some(path) = model_path {
        info!("Loading model from: {}", path.display());
        // Model loading would go here
    }

    Ok(trainer)
}

/// Get scenario by name
fn get_scenario_by_name(name: &str) -> Result<StressScenario> {
    match name.to_lowercase().as_str() {
        "flash_crash" | "flash" => Ok(flash_crash_scenario()),
        "liquidity_crisis" | "liquidity" => Ok(liquidity_crisis_scenario()),
        "vix_spike" | "vix" => Ok(vix_spike_scenario()),
        "trending" | "trending_market" => Ok(trending_market_scenario()),
        "whipsaw" => Ok(whipsaw_scenario()),
        "gap_risk" | "gap" => Ok(gap_risk_scenario()),
        "correlation_breakdown" | "correlation" => Ok(correlation_breakdown_scenario()),
        "multi_asset" | "multi_asset_stress" => Ok(multi_asset_stress_scenario()),
        _ => anyhow::bail!("Unknown scenario: {}. Valid options: flash_crash, liquidity_crisis, vix_spike, trending, whipsaw, gap_risk, correlation_breakdown, multi_asset", name),
    }
}

/// Print detailed result for single scenario
fn print_detailed_result(result: &ml::dqn::stress_testing::StressResult) {
    println!("\n{}", "=".repeat(80));
    println!("STRESS TEST RESULT: {}", result.scenario_name);
    println!("{}", "=".repeat(80));
    println!("Status:              {}", if result.passed { "✅ PASSED" } else { "❌ FAILED" });
    println!("Max Drawdown:        {:.2}%", result.max_drawdown);
    println!("Final Portfolio:     {:.2}%", result.final_portfolio_pct);
    println!("Action Diversity:    {:.2}%", result.action_diversity);
    println!("Total Trades:        {}", result.total_trades);
    println!("Q-Value Std Dev:     {:.4}", result.q_value_std);
    println!("Bankruptcy:          {}", if result.bankruptcy { "YES ⚠️" } else { "NO" });
    println!("Circuit Breaker:     {}", if result.circuit_breaker_triggered { "TRIGGERED" } else { "OK" });
    println!("Execution Time:      {}ms", result.execution_time_ms);

    if let Some(recovery) = result.recovery_steps {
        println!("Recovery Time:       {} steps", recovery);
    } else {
        println!("Recovery Time:       Not achieved");
    }

    if !result.passed {
        println!("\nFailure Reasons:");
        for (i, reason) in result.failure_reasons.iter().enumerate() {
            println!("  {}. {}", i + 1, reason);
        }
    }

    println!("{}", "=".repeat(80));
}

/// Save report to JSON file
fn save_report_to_file(
    report: &ml::dqn::stress_testing::StressTestReport,
    path: &std::path::Path,
) -> Result<()> {
    let json = serde_json::to_string_pretty(report)
        .context("Failed to serialize report to JSON")?;
    std::fs::write(path, json).context("Failed to write report to file")?;
    Ok(())
}