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 = "30")]
    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,

    /// Optimizer to use: "tpe" (Tree-Parzen Estimator) or "pso" (Particle Swarm)
    #[arg(long, default_value = "tpe")]
    optimizer: String,

    /// Number of parallel trial evaluations (PSO only; TPE is always sequential).
    /// 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,

    /// Path to MBP-10 order book data for OFI features (order flow imbalance)
    #[arg(long)]
    mbp10_data_dir: Option<PathBuf>,

    /// Path to trades data for VPIN / Kyle's Lambda features
    #[arg(long)]
    trades_data_dir: Option<PathBuf>,
}

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

#[allow(clippy::cognitive_complexity)]
fn run_dqn_hyperopt(args: &Args, parallel: usize, gpu_devices: &[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_devices(gpu_devices.to_vec())
        .with_ofi_data_dirs(
            args.mbp10_data_dir.as_ref().map(|p| p.to_string_lossy().into_owned()),
            args.trades_data_dir.as_ref().map(|p| p.to_string_lossy().into_owned()),
        );

    // 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);
    }

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

    let start = Instant::now();
    let result = if args.optimizer.as_str() == "tpe" {
        info!("Using TPE (Tree-Parzen Estimator) optimizer");
        ml::hyperopt::optimize_with_tpe(trainer, args.trials, args.n_initial, Some(args.seed), Some("dqn"))
            .context("DQN TPE hyperopt optimization failed")?
    } else {
        let optimizer = ArgminOptimizer::builder()
            .max_trials(args.trials)
            .n_initial(args.n_initial)
            .seed(args.seed)
            .model_name("dqn")
            .build();

        if parallel > 1 {
            info!("Using parallel PSO 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 = match serde_json::to_value(&result.best_params) {
        Ok(v) => v,
        Err(e) => {
            warn!("Failed to serialize best hyperopt params: {}", e);
            Value::Null
        }
    };

    // Extract top-5 trials sorted by objective (ascending = best first)
    let mut sorted_trials = result.all_trials.clone();
    sorted_trials.sort_by(|a, b| {
        a.objective
            .partial_cmp(&b.objective)
            .unwrap_or(std::cmp::Ordering::Equal)
    });
    let top_k: Vec<Value> = sorted_trials
        .iter()
        .take(5)
        .filter_map(|t| {
            serde_json::to_value(&t.params).ok().map(|params| {
                serde_json::json!({
                    "params": params,
                    "objective": t.objective,
                    "trial_num": t.trial_num,
                })
            })
        })
        .collect();

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

#[allow(clippy::cognitive_complexity)]
fn run_ppo_hyperopt(args: &Args, parallel: usize, gpu_devices: &[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_devices(gpu_devices.to_vec());

    // 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);
    }

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

    let start = Instant::now();
    let result = if args.optimizer.as_str() == "tpe" {
        info!("Using TPE (Tree-Parzen Estimator) optimizer");
        ml::hyperopt::optimize_with_tpe(trainer, args.trials, args.n_initial, Some(args.seed), Some("ppo"))
            .context("PPO TPE hyperopt optimization failed")?
    } else {
        let optimizer = ArgminOptimizer::builder()
            .max_trials(args.trials)
            .n_initial(args.n_initial)
            .seed(args.seed)
            .model_name("ppo")
            .build();

        if parallel > 1 {
            info!("Using parallel PSO 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 = match serde_json::to_value(&result.best_params) {
        Ok(v) => v,
        Err(e) => {
            warn!("Failed to serialize best hyperopt params: {}", e);
            Value::Null
        }
    };

    // Extract top-5 trials sorted by objective (ascending = best first)
    let mut sorted_trials = result.all_trials.clone();
    sorted_trials.sort_by(|a, b| {
        a.objective
            .partial_cmp(&b.objective)
            .unwrap_or(std::cmp::Ordering::Equal)
    });
    let top_k: Vec<Value> = sorted_trials
        .iter()
        .take(5)
        .filter_map(|t| {
            serde_json::to_value(&t.params).ok().map(|params| {
                serde_json::json!({
                    "params": params,
                    "objective": t.objective,
                    "trial_num": t.trial_num,
                })
            })
        })
        .collect();

    Ok(build_model_result(
        result.best_objective,
        best_params_json,
        top_k,
        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!("Optimizer: {}", args.optimizer);
    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);

    // Detect all available GPUs for multi-GPU trial parallelism.
    // Each trial runs independently on a different GPU (no NCCL needed).
    let gpu_count = ml::cuda_pipeline::multi_gpu::MultiGpuConfig::count_cuda_devices_pub();
    let mut gpu_devices = Vec::with_capacity(gpu_count.max(1));
    for i in 0..gpu_count {
        match candle_core::Device::new_cuda(i) {
            Ok(d) => gpu_devices.push(d),
            Err(e) => {
                warn!("Failed to init CUDA device {}: {}", i, e);
                break;
            }
        }
    }
    if gpu_devices.is_empty() {
        return Err(anyhow::anyhow!("CUDA GPU required for hyperopt but none available"));
    }
    info!("GPUs: {} device(s) available for trial parallelism", gpu_devices.len());

    // Resolve parallel: 0 = auto-detect.
    let cpu_threads = if args.parallel == 0 {
        cpus.saturating_sub(1).max(1)
    } else {
        args.parallel
    };

    // VRAM budget: multiply per-GPU concurrency by number of GPUs.
    let budget = ml::hyperopt::HardwareBudget::detect();
    let per_gpu_cap = budget.plan_hyperopt(7000.0, 0.0005, 64.0, 4096.0).max_concurrent_trials;
    let total_vram_cap = per_gpu_cap * gpu_devices.len();
    let num_threads = cpu_threads.min(total_vram_cap);
    info!("Parallel: {} threads (CPU: {}, CPU cap: {}, per-GPU cap: {}, GPUs: {}, total VRAM cap: {})",
          num_threads, cpus, cpu_threads, per_gpu_cap, gpu_devices.len(), total_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, &gpu_devices) {
            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, &gpu_devices) {
            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);

    // Push final metrics to pushgateway so they persist after pod termination
    if let Err(e) = metrics_server::push_to_gateway(None, "hyperopt_baseline_rl") {
        warn!("Failed to push metrics to gateway (non-fatal): {e}");
    }

    Ok(())
}