foxhunt/crates/ml/examples/hyperopt_baseline_rl.rs

//! Hyperopt RL Runner for DQN/PPO on Real Databento Market Data
//!
//! Runs hyperparameter optimization using Particle Swarm Optimization (PSO) for
//! DQN, PPO, or both models on downloaded Databento futures data. This binary is
//! part of the walk-forward real-data training pipeline.
//!
//! ## Usage
//!
//! ```bash
//! # Run DQN hyperopt only (10 trials, 10 epochs each)
//! SQLX_OFFLINE=true cargo run -p ml --example hyperopt_baseline --release -- \
//!   --model dqn --trials 10 --epochs 10 \
//!   --data-dir test_data/futures-baseline
//!
//! # Run PPO hyperopt only
//! SQLX_OFFLINE=true cargo run -p ml --example hyperopt_baseline --release -- \
//!   --model ppo --trials 20 --epochs 15 \
//!   --data-dir test_data/futures-baseline
//!
//! # Run both models (default)
//! SQLX_OFFLINE=true cargo run -p ml --example hyperopt_baseline --release -- \
//!   --data-dir test_data/futures-baseline
//! ```
//!
//! ## Output
//!
//! Results are written as JSON to `--output` (default: `ml/trained_models/hyperopt_results.json`).
//!
//! ```json
//! {
//!   "dqn": { "best_objective": 0.123, "best_params": {...}, "trials": 20, "elapsed_secs": 45.3 },
//!   "ppo": { "best_objective": 0.456, "best_params": {...}, "trials": 20, "elapsed_secs": 67.8 }
//! }
//! ```

#![allow(unused_crate_dependencies, unsafe_code)]

use anyhow::{Context, Result};
use clap::Parser;
use serde_json::Value;
use std::path::PathBuf;
use std::time::Instant;
use tracing::{error, info, warn};

use ml::hyperopt::adapters::dqn::DQNTrainer;
use ml::hyperopt::adapters::ppo::PPOTrainer;
use ml::hyperopt::paths::{generate_run_id, TrainingPaths};
use ml::hyperopt::ArgminOptimizer;

use common::metrics::{server as metrics_server, training_metrics};

/// Hyperparameter optimization runner for DQN/PPO on Databento market data
#[derive(Parser, Debug)]
#[command(name = "hyperopt-baseline-rl")]
#[command(about = "Run hyperparameter optimization for DQN/PPO on real Databento futures data")]
struct Args {
    /// Model to optimize: "dqn", "ppo", or "both"
    #[arg(long, default_value = "both")]
    model: String,

    /// Number of PSO trials per model
    #[arg(long, default_value = "20")]
    trials: usize,

    /// Number of initial LHS (Latin Hypercube Sampling) samples
    #[arg(long, default_value = "5")]
    n_initial: usize,

    /// Training epochs per trial (DQN) / episodes per trial (PPO)
    #[arg(long, default_value = "10")]
    epochs: usize,

    /// Path to downloaded DBN data directory
    #[arg(long, default_value = "test_data/futures-baseline")]
    data_dir: PathBuf,

    /// Output path for JSON results
    #[arg(long, default_value = "ml/trained_models/hyperopt_results.json")]
    output: PathBuf,

    /// Random seed for reproducibility
    #[arg(long, default_value = "42")]
    seed: u64,

    /// Base directory for training run outputs (checkpoints, logs, metrics)
    #[arg(long, default_value = "/tmp/ml_training")]
    base_dir: String,

    /// Symbol subdirectory to load (e.g. "ES.FUT", "NQ.FUT").
    /// Restricts data loading to `data_dir`/`symbol`/ to avoid mixing instruments.
    #[arg(long, default_value = "ES.FUT")]
    symbol: String,

    /// Transaction cost per trade in basis points (commission + exchange fees).
    /// IBKR ES all-in: ~$2.06/contract on ~$275K notional ≈ 0.08 bps.
    #[arg(long, default_value = "0.1")]
    tx_cost_bps: f64,

    /// Tick size for spread calculation (e.g. 0.25 for ES futures)
    #[arg(long, default_value = "0.25")]
    tick_size: f64,

    /// Typical bid-ask spread in ticks (ES is almost always 1 tick)
    #[arg(long, default_value = "1.0")]
    spread_ticks: f64,

    /// Number of parallel trial evaluations.
    /// Each trial uses ~1 CPU core + shared GPU for forward/backward.
    /// 0 = auto-detect (CPUs - 1; GPU-bound trials need minimal CPU).
    /// 1 = sequential. N = N concurrent trials.
    #[arg(long, default_value = "0")]
    parallel: usize,

    /// Initial trading capital in dollars
    #[arg(long, default_value = "35000")]
    initial_capital: f64,
}

/// Result entry for one model's hyperopt run
fn build_model_result(
    best_objective: f64,
    best_params_json: Value,
    num_trials: usize,
    elapsed_secs: f64,
) -> Value {
    serde_json::json!({
        "best_objective": best_objective,
        "best_params": best_params_json,
        "trials": num_trials,
        "elapsed_secs": elapsed_secs,
    })
}

#[allow(clippy::cognitive_complexity)]
fn run_dqn_hyperopt(args: &Args, parallel: usize, device: &candle_core::Device) -> Result<Value> {
    info!("========================================");
    info!("  DQN Hyperparameter Optimization");
    info!("========================================");

    let run_id = generate_run_id("hyperopt-dqn");
    let training_paths = TrainingPaths::new(&args.base_dir, "dqn", &run_id);

    info!("Run ID: {}", run_id);
    info!("Data directory: {}", args.data_dir.display());
    info!("Epochs per trial: {}", args.epochs);
    info!("Trials: {}", args.trials);

    let symbol_dir = args.data_dir.join(&args.symbol);
    let mut trainer = DQNTrainer::new(&symbol_dir, args.epochs)
        .context("Failed to create DQN trainer")?
        .with_training_paths(training_paths)
        .with_initial_capital(args.initial_capital)
        .with_costs(args.tx_cost_bps, args.tick_size, args.spread_ticks)
        .with_device(device.clone());

    // Preload training data once — all trials reuse via Arc (no per-trial disk I/O)
    if let Err(e) = trainer.preload_data() {
        warn!("Data preload failed ({}), trials will load from disk individually", e);
    } else {
        info!("Training data preloaded and cached for all {} trials", args.trials);
    }

    let optimizer = ArgminOptimizer::builder()
        .max_trials(args.trials)
        .n_initial(args.n_initial)
        .seed(args.seed)
        .build();

    training_metrics::set_hyperopt_trial("dqn", 0.0, args.trials as f64);

    let start = Instant::now();
    let result = if parallel > 1 {
        info!("Using parallel optimization ({} threads)", parallel);
        optimizer
            .optimize_parallel(trainer)
            .context("DQN parallel hyperopt optimization failed")?
    } else {
        optimizer
            .optimize(trainer)
            .context("DQN hyperopt optimization failed")?
    };
    let elapsed = start.elapsed().as_secs_f64();

    training_metrics::set_hyperopt_trial("dqn", result.all_trials.len() as f64, args.trials as f64);
    training_metrics::set_hyperopt_best_objective("dqn", result.best_objective);
    training_metrics::set_hyperopt_elapsed("dqn", elapsed);

    info!("DQN hyperopt complete:");
    info!("  Best objective: {:.6}", result.best_objective);
    info!("  Total trials: {}", result.all_trials.len());
    info!("  Elapsed: {:.1}s", elapsed);

    let best_params_json = serde_json::to_value(&result.best_params)
        .ok()
        .unwrap_or(Value::Null);

    Ok(build_model_result(
        result.best_objective,
        best_params_json,
        result.all_trials.len(),
        elapsed,
    ))
}

#[allow(clippy::cognitive_complexity)]
fn run_ppo_hyperopt(args: &Args, parallel: usize, device: &candle_core::Device) -> Result<Value> {
    info!("========================================");
    info!("  PPO Hyperparameter Optimization");
    info!("========================================");

    let run_id = generate_run_id("hyperopt-ppo");
    let training_paths = TrainingPaths::new(&args.base_dir, "ppo", &run_id);

    info!("Run ID: {}", run_id);
    info!("Data directory: {}", args.data_dir.display());
    info!("Episodes per trial: {}", args.epochs);
    info!("Trials: {}", args.trials);

    let symbol_dir = args.data_dir.join(&args.symbol);
    let mut trainer = PPOTrainer::new(&symbol_dir, args.epochs)
        .context("Failed to create PPO trainer")?
        .with_training_paths(training_paths)
        .with_costs(args.tx_cost_bps, args.tick_size, args.spread_ticks)
        .with_device(device.clone());

    // Preload training data once — all trials reuse via Arc (no per-trial disk I/O)
    if let Err(e) = trainer.preload_data() {
        warn!("Data preload failed ({}), trials will load from disk individually", e);
    } else {
        info!("Training data preloaded and cached for all {} trials", args.trials);
    }

    let optimizer = ArgminOptimizer::builder()
        .max_trials(args.trials)
        .n_initial(args.n_initial)
        .seed(args.seed)
        .build();

    training_metrics::set_hyperopt_trial("ppo", 0.0, args.trials as f64);

    let start = Instant::now();
    let result = if parallel > 1 {
        info!("Using parallel optimization ({} threads)", parallel);
        optimizer
            .optimize_parallel(trainer)
            .context("PPO parallel hyperopt optimization failed")?
    } else {
        optimizer
            .optimize(trainer)
            .context("PPO hyperopt optimization failed")?
    };
    let elapsed = start.elapsed().as_secs_f64();

    training_metrics::set_hyperopt_trial("ppo", result.all_trials.len() as f64, args.trials as f64);
    training_metrics::set_hyperopt_best_objective("ppo", result.best_objective);
    training_metrics::set_hyperopt_elapsed("ppo", elapsed);

    info!("PPO hyperopt complete:");
    info!("  Best objective: {:.6}", result.best_objective);
    info!("  Total trials: {}", result.all_trials.len());
    info!("  Elapsed: {:.1}s", elapsed);

    let best_params_json = serde_json::to_value(&result.best_params)
        .ok()
        .unwrap_or(Value::Null);

    Ok(build_model_result(
        result.best_objective,
        best_params_json,
        result.all_trials.len(),
        elapsed,
    ))
}

#[allow(clippy::cognitive_complexity)]
fn main() -> Result<()> {
    // Initialize tracing with optional OTLP export to Tempo
    let otlp_endpoint = std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").ok();
    if let Err(e) = common::observability::init_observability(
        "hyperopt_baseline_rl",
        otlp_endpoint.as_deref(),
    ) {
        eprintln!("Observability init failed (non-fatal): {e}");
    }

    training_metrics::init();
    metrics_server::start_metrics_server(9094);
    training_metrics::set_active_workers(1.0);

    // Pre-allocate CUBLAS workspace for deterministic + faster tensor core ops.
    // SAFETY: called once at startup before any multi-threading or CUDA work begins.
    unsafe { std::env::set_var("CUBLAS_WORKSPACE_CONFIG", ":4096:8"); }

    let args = Args::parse();

    // Signal hyperopt mode active — will be cleared at exit
    let hyperopt_model_label = args.model.clone();
    training_metrics::set_hyperopt_mode(&hyperopt_model_label, true);

    info!("========================================");
    info!("  Hyperopt Baseline Runner");
    info!("========================================");
    info!("Model: {}", args.model);
    info!("Symbol: {}", args.symbol);
    info!("Trials: {}", args.trials);
    info!("Initial LHS samples: {}", args.n_initial);
    info!("Epochs/episodes per trial: {}", args.epochs);
    info!("Data directory: {}", args.data_dir.display());
    info!("Output: {}", args.output.display());
    info!("Seed: {}", args.seed);
    info!("Tx cost: {} bps + {} tick spread (tick_size={})", args.tx_cost_bps, args.spread_ticks, args.tick_size);
    info!("Initial capital: ${:.0}", args.initial_capital);

    let cpus = std::thread::available_parallelism()
        .map(|n| n.get())
        .unwrap_or(1);

    // Always use GPU — DQN/PPO networks are small enough that all parallel
    // trials fit in VRAM. GPU forward/backward is faster even with CUDA
    // context sharing across threads.
    let device = candle_core::Device::new_cuda(0)
        .map_err(|e| anyhow::anyhow!("CUDA GPU required for hyperopt but unavailable: {}", e))?;
    info!("Device: CUDA GPU");

    // Resolve parallel: 0 = auto-detect.
    // DQN/PPO trials are GPU-bound: each rayon thread submits CUDA kernels
    // for forward/backward passes and spends most of its time waiting on
    // cudaDeviceSynchronize(). CPU usage per thread is minimal (data feeding,
    // reward calculation), so cpus-1 threads is safe — reserve 1 core for
    // the CUDA driver event loop, OS, and the OTLP batch exporter.
    let cpu_threads = if args.parallel == 0 {
        cpus.saturating_sub(1).max(1)
    } else {
        args.parallel
    };

    // Cap by VRAM budget — use DQN trial overhead (7 GB/trial, the most VRAM-heavy
    // model) as conservative estimate. PPO uses less but benefits from the same cap.
    let budget = ml::hyperopt::HardwareBudget::detect();
    let vram_cap = budget.plan_hyperopt(7000.0, 0.0005, 64.0, 4096.0).max_concurrent_trials;
    let num_threads = cpu_threads.min(vram_cap);
    info!("Parallel: {} threads (CPU: {}, CPU cap: {}, VRAM cap: {})", num_threads, cpus, cpu_threads, vram_cap);

    // Configure rayon thread pool for parallel trial evaluation
    if num_threads > 1 {
        rayon::ThreadPoolBuilder::new()
            .num_threads(num_threads)
            .build_global()
            .unwrap_or_else(|e| {
                warn!("Failed to set rayon thread pool size: {}", e);
            });
        info!("Rayon thread pool configured: {} threads", num_threads);
    }

    // Scope data directory to symbol subdirectory to avoid mixing instruments
    let symbol_data_dir = args.data_dir.join(&args.symbol);
    if !symbol_data_dir.exists() {
        // List available subdirectories for helpful error message
        let available: Vec<String> = std::fs::read_dir(&args.data_dir)
            .ok()
            .map(|entries| {
                entries
                    .filter_map(|e| e.ok())
                    .filter(|e| e.path().is_dir())
                    .filter_map(|e| e.file_name().into_string().ok())
                    .collect()
            })
            .unwrap_or_default();
        anyhow::bail!(
            "Symbol directory not found: {}. Available symbols: {:?}. Run `download_baseline` first.",
            symbol_data_dir.display(),
            available
        );
    }

    // Verify trials > n_initial (ArgminOptimizer requirement)
    if args.trials <= args.n_initial {
        anyhow::bail!(
            "trials ({}) must be greater than n_initial ({})",
            args.trials,
            args.n_initial
        );
    }

    // Create output directory
    if let Some(parent) = args.output.parent() {
        std::fs::create_dir_all(parent)
            .with_context(|| format!("Failed to create output directory: {}", parent.display()))?;
    }

    let run_dqn = args.model == "dqn" || args.model == "both";
    let run_ppo = args.model == "ppo" || args.model == "both";

    if !run_dqn && !run_ppo {
        anyhow::bail!(
            "Invalid --model value '{}'. Must be 'dqn', 'ppo', or 'both'.",
            args.model
        );
    }

    let mut results = serde_json::Map::new();

    // Run DQN hyperopt
    if run_dqn {
        match run_dqn_hyperopt(&args, num_threads, &device) {
            Ok(dqn_result) => {
                results.insert("dqn".to_owned(), dqn_result);
            },
            Err(e) => {
                error!("DQN hyperopt failed: {:#}", e);
                warn!("Continuing with remaining models...");
                results.insert(
                    "dqn".to_owned(),
                    serde_json::json!({ "error": format!("{:#}", e) }),
                );
            },
        }
    }

    // Run PPO hyperopt
    if run_ppo {
        match run_ppo_hyperopt(&args, num_threads, &device) {
            Ok(ppo_result) => {
                results.insert("ppo".to_owned(), ppo_result);
            },
            Err(e) => {
                error!("PPO hyperopt failed: {:#}", e);
                warn!("Continuing...");
                results.insert(
                    "ppo".to_owned(),
                    serde_json::json!({ "error": format!("{:#}", e) }),
                );
            },
        }
    }

    // Write results to JSON
    let output_json = Value::Object(results);
    let output_str = serde_json::to_string_pretty(&output_json)
        .context("Failed to serialize results to JSON")?;

    std::fs::write(&args.output, &output_str)
        .with_context(|| format!("Failed to write results to {}", args.output.display()))?;

    info!("========================================");
    info!("  Results saved to: {}", args.output.display());
    info!("========================================");
    info!("{}", output_str);

    training_metrics::set_hyperopt_mode(&hyperopt_model_label, false);
    training_metrics::set_active_workers(0.0);
    Ok(())
}