foxhunt/ml/examples/train_dqn_es_fut.rs

//! **DQN Training on ES.FUT Real Market Data**
//!
//! Production training script for DQN model on ES.FUT futures data.
//!
//! ## Usage
//!
//! ```bash
//! # Fast training (10 epochs, ~5 seconds)
//! cargo run -p ml --example train_dqn_es_fut --release
//!
//! # Production training (50 epochs, ~20 seconds)
//! cargo run -p ml --example train_dqn_es_fut --release -- --epochs 50
//!
//! # Full training (200 epochs, ~80 seconds)
//! cargo run -p ml --example train_dqn_es_fut --release -- --epochs 200
//! ```
//!
//! ## Expected Results
//!
//! - **10 epochs**: Loss ~0.15, Q-value ~3.0
//! - **50 epochs**: Loss ~0.04, Q-value ~0.9 (production checkpoint)
//! - **200 epochs**: Loss ~0.01, Q-value ~0.5 (maximum convergence)
//!
//! ## Output
//!
//! Checkpoint saved to: `ml/checkpoints/dqn_es_fut_v1.safetensors`

use anyhow::{Context, Result};
use clap::Parser;
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
use std::path::PathBuf;
use std::time::Instant;
use tracing::{info, Level};
use tracing_subscriber::FmtSubscriber;

#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// Number of training epochs
    #[arg(short, long, default_value_t = 10)]
    epochs: usize,

    /// Batch size (max 230 for RTX 3050 Ti 4GB VRAM)
    #[arg(short, long, default_value_t = 128)]
    batch_size: usize,

    /// Learning rate
    #[arg(short, long, default_value_t = 0.0001)]
    learning_rate: f64,

    /// Data directory
    #[arg(
        short,
        long,
        default_value = "../test_data/real/databento/ml_training_small"
    )]
    data_dir: String,

    /// Output checkpoint path
    #[arg(short, long, default_value = "checkpoints/dqn_es_fut_v1.safetensors")]
    output: String,

    /// Enable early stopping
    #[arg(long, default_value_t = true)]
    early_stopping: bool,

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

#[tokio::main]
async fn main() -> Result<()> {
    let args = Args::parse();

    // Setup logging
    let log_level = if args.verbose {
        Level::DEBUG
    } else {
        Level::INFO
    };
    let subscriber = FmtSubscriber::builder()
        .with_max_level(log_level)
        .with_target(false)
        .with_thread_ids(false)
        .with_file(false)
        .with_line_number(false)
        .finish();
    tracing::subscriber::set_global_default(subscriber)?;

    println!("\n{}", "=".repeat(80));
    println!("🚀 DQN Training on ES.FUT Real Market Data");
    println!("{}", "=".repeat(80));
    println!();
    println!("⚙️  Configuration:");
    println!("   Epochs:         {}", args.epochs);
    println!("   Batch Size:     {}", args.batch_size);
    println!("   Learning Rate:  {}", args.learning_rate);
    println!("   Data Dir:       {}", args.data_dir);
    println!("   Output:         {}", args.output);
    println!("   Early Stopping: {}", args.early_stopping);
    println!();

    let start_time = Instant::now();

    // ========================================================================
    // Step 1: Verify data directory exists
    // ========================================================================
    info!("Verifying data directory...");

    let data_path = PathBuf::from(&args.data_dir);
    if !data_path.exists() {
        eprintln!("❌ Error: Data directory not found: {}", args.data_dir);
        eprintln!("   Run data acquisition first or check path.");
        std::process::exit(1);
    }

    // Count DBN files
    let dbn_files: Vec<_> = std::fs::read_dir(&data_path)?
        .filter_map(|entry| entry.ok())
        .filter(|entry| entry.path().extension().and_then(|s| s.to_str()) == Some("dbn"))
        .collect();

    if dbn_files.is_empty() {
        eprintln!("❌ Error: No DBN files found in: {}", args.data_dir);
        std::process::exit(1);
    }

    info!("Found {} DBN files", dbn_files.len());
    println!(
        "✅ Data directory validated ({} DBN files)\n",
        dbn_files.len()
    );

    // ========================================================================
    // Step 2: Configure DQN hyperparameters
    // ========================================================================
    info!("Configuring DQN hyperparameters...");

    let mut hyperparams = DQNHyperparameters::conservative();
    hyperparams.epochs = args.epochs;
    hyperparams.batch_size = args.batch_size;
    hyperparams.learning_rate = args.learning_rate;
    hyperparams.gamma = 0.99;
    hyperparams.epsilon_start = 1.0;
    hyperparams.epsilon_end = 0.01;
    hyperparams.epsilon_decay = 0.995;
    hyperparams.buffer_size = 100_000;
    hyperparams.checkpoint_frequency = args.epochs / 5; // Save 5 checkpoints
    hyperparams.early_stopping_enabled = args.early_stopping;
    hyperparams.q_value_floor = 0.5;
    hyperparams.min_loss_improvement_pct = 2.0;
    hyperparams.plateau_window = 30;
    hyperparams.min_epochs_before_stopping = args.epochs / 2;

    // Validate batch size
    if hyperparams.batch_size > 230 {
        eprintln!(
            "❌ Error: Batch size {} exceeds GPU limit (230)",
            hyperparams.batch_size
        );
        eprintln!("   Reduce batch size to fit in 4GB VRAM.");
        std::process::exit(1);
    }

    println!("✅ Hyperparameters configured\n");

    // ========================================================================
    // Step 3: Create DQN trainer
    // ========================================================================
    info!("Initializing DQN trainer...");

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

    println!("✅ DQN trainer initialized\n");

    // ========================================================================
    // Step 4: Setup checkpoint directory
    // ========================================================================
    info!("Setting up checkpoint directory...");

    let output_path = PathBuf::from(&args.output);
    let checkpoint_dir = output_path.parent().context("Invalid output path")?;

    std::fs::create_dir_all(checkpoint_dir)?;

    println!(
        "✅ Checkpoint directory ready: {}\n",
        checkpoint_dir.display()
    );

    // ========================================================================
    // Step 5: Run training
    // ========================================================================
    println!("{}", "=".repeat(80));
    println!("🏋️  Starting DQN Training");
    println!("{}", "=".repeat(80));
    println!();

    let training_start = Instant::now();
    let mut checkpoint_count = 0;

    let metrics = trainer
        .train(&args.data_dir, |epoch, checkpoint_data| {
            checkpoint_count += 1;

            let checkpoint_path = if epoch == args.epochs {
                // Final checkpoint
                output_path.clone()
            } else {
                // Intermediate checkpoint
                checkpoint_dir.join(format!("dqn_es_fut_epoch_{}.safetensors", epoch))
            };

            std::fs::write(&checkpoint_path, checkpoint_data)
                .context("Failed to write checkpoint")?;

            let size_kb = std::fs::metadata(&checkpoint_path)?.len() / 1024;
            info!(
                "Checkpoint saved: epoch {} ({} KB) -> {}",
                epoch,
                size_kb,
                checkpoint_path.display()
            );

            Ok(checkpoint_path.to_string_lossy().to_string())
        })
        .await
        .context("Training failed")?;

    let training_time = training_start.elapsed();

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

    // ========================================================================
    // Step 6: Report results
    // ========================================================================
    println!("📊 Training Metrics:");
    println!();
    println!("   Epochs Completed:   {}", metrics.epochs_trained);
    println!("   Final Loss:         {:.6}", metrics.loss);
    println!("   Convergence:        {}", metrics.convergence_achieved);
    println!();

    if let Some(avg_q_value) = metrics.additional_metrics.get("avg_q_value") {
        println!("   Avg Q-value:        {:.4}", avg_q_value);
    }

    if let Some(avg_grad_norm) = metrics.additional_metrics.get("avg_gradient_norm") {
        println!("   Avg Gradient Norm:  {:.6}", avg_grad_norm);
    }

    if let Some(final_epsilon) = metrics.additional_metrics.get("final_epsilon") {
        println!("   Final Epsilon:      {:.4}", final_epsilon);
    }

    println!();
    println!("⏱️  Performance:");
    println!();
    println!(
        "   Training Time:      {:.2}s ({:.1} min)",
        training_time.as_secs_f64(),
        training_time.as_secs_f64() / 60.0
    );
    println!(
        "   Avg Epoch Time:     {:.3}s",
        training_time.as_secs_f64() / metrics.epochs_trained as f64
    );
    println!("   Checkpoints Saved:  {}", checkpoint_count);
    println!();

    // ========================================================================
    // Step 7: Verify final checkpoint
    // ========================================================================
    if output_path.exists() {
        let checkpoint_size = std::fs::metadata(&output_path)?.len();
        println!("💾 Final Checkpoint:");
        println!();
        println!("   Path: {}", output_path.display());
        println!(
            "   Size: {} KB ({} bytes)",
            checkpoint_size / 1024,
            checkpoint_size
        );
        println!();
    }

    // ========================================================================
    // Summary
    // ========================================================================
    let total_time = start_time.elapsed();

    println!("{}", "=".repeat(80));
    println!("🎉 DQN Training Successful");
    println!("{}", "=".repeat(80));
    println!();
    println!("✅ Model trained and saved to: {}", args.output);
    println!(
        "⏱️  Total time: {:.2}s ({:.1} min)",
        total_time.as_secs_f64(),
        total_time.as_secs_f64() / 60.0
    );
    println!();

    // Next steps
    println!("📌 Next Steps:");
    println!();
    println!("   1. Run inference test:");
    println!("      cargo test -p ml dqn_training_pipeline_test");
    println!();
    println!("   2. Integrate with paper trading:");
    println!("      See services/trading_service/src/paper_trading_executor.rs");
    println!();
    println!("   3. Monitor performance:");
    println!("      Check Grafana dashboard for ML metrics");
    println!();

    Ok(())
}