foxhunt/crates/ml/examples/train_baseline_rl.rs

//! Walk-forward RL training binary for DQN and PPO models.
//!
//! Trains models using expanding walk-forward windows on real OHLCV data loaded
//! from Databento DBN files. Supports early stopping, checkpoint saving, and
//! normalization statistics export for reproducible evaluation.
//!
//! # Usage
//!
//! ```bash
//! SQLX_OFFLINE=true cargo run -p ml --example train_baseline -- \
//!     --model both --epochs 50 --batch-size 128 \
//!     --data-dir test_data/futures-baseline \
//!     --output-dir ml/trained_models
//! ```

#![allow(unused_crate_dependencies)]
#![deny(
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::panic,
    clippy::indexing_slicing
)]

use std::path::{Path, PathBuf};

use anyhow::{Context, Result};
use clap::Parser;
use serde_json::Value;
use tracing::{error, info, warn};

use ml::cuda_pipeline::DqnGpuData;
use ml::prelude::Device;
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
use ml::trainers::ppo::{PpoHyperparameters, PpoTrainer};

#[allow(unreachable_pub)]
mod baseline_common;
use baseline_common::completion::{write_failure_marker, write_success_marker, CompletionMetrics};
use baseline_common::{load_all_bars, spread_cost_bps};
use common::metrics::{server as metrics_server, training_metrics as metrics};
use ml::features::extraction::{extract_ml_features, FeatureVector};
use ml::gpu::memory_profile::{detect_vram_mb, vram_scaled_base_dim};
use ml::types::OHLCVBar;
use ml::walk_forward::{generate_walk_forward_windows, NormStats, WalkForwardConfig};

// ---------------------------------------------------------------------------
// CLI Arguments
// ---------------------------------------------------------------------------

/// Walk-forward training binary for DQN and PPO baseline models.
#[derive(Parser, Debug)]
#[command(name = "train_baseline_rl", about = "Train DQN/PPO with walk-forward RL windows")]
struct Args {
    /// Which model(s) to train: "dqn", "ppo", or "both"
    #[arg(long, default_value = "both")]
    model: String,

    /// Maximum training epochs per fold
    #[arg(long, default_value_t = 50)]
    epochs: usize,

    /// Training batch size
    #[arg(long, default_value_t = 128)]
    batch_size: usize,

    /// Path to directory containing .dbn.zst files
    #[arg(long, default_value = "test_data/futures-baseline")]
    data_dir: PathBuf,

    /// Output directory for trained model checkpoints
    #[arg(long, default_value = "ml/trained_models")]
    output_dir: PathBuf,

    /// Optional path to hyperopt results JSON -- overrides matching config fields
    #[arg(long)]
    hyperopt_params: Option<PathBuf>,

    /// Feature dimension (42 market + 3 portfolio = 45, or 53 with OFI; must match trainer `state_dim`)
    #[arg(long, default_value_t = 43)]
    feature_dim: usize,

    /// Early stopping patience (epochs without improvement)
    #[arg(long, default_value_t = 10)]
    patience: usize,

    /// Number of actions (5 exposure levels for DQN, pass --num-actions 45 for PPO)
    #[arg(long, default_value_t = 5)]
    num_actions: usize,

    /// Walk-forward: initial training window in months
    #[arg(long, default_value_t = 12)]
    train_months: u32,

    /// Walk-forward: validation window in months
    #[arg(long, default_value_t = 3)]
    val_months: u32,

    /// Walk-forward: test window in months
    #[arg(long, default_value_t = 3)]
    test_months: u32,

    /// Walk-forward: step size in months between folds
    #[arg(long, default_value_t = 3)]
    step_months: u32,

    /// Learning rate for optimizer
    #[arg(long, default_value_t = 1e-4)]
    learning_rate: f64,

    /// Max environment steps per epoch (caps trajectory length for OOM safety;
    /// 0 = use all bars, but this can use >1GB RAM per fold on large datasets)
    #[arg(long, default_value_t = 2000)]
    max_steps_per_epoch: usize,

    /// Symbol subdirectory to load (e.g. "ES.FUT", "NQ.FUT")
    #[arg(long, default_value = "ES.FUT")]
    symbol: String,

    /// Maximum absolute per-bar return; larger moves are clamped (contract roll filter)
    #[arg(long, default_value_t = 0.01)]
    max_bar_return: f64,

    /// Round-trip commission cost in basis points (1 bps = 0.01%)
    /// Applied to BUY/SELL rewards; HOLD is free.
    /// Default 1.0 bps covers ~$4 exchange+broker for ES e-mini.
    #[arg(long, default_value_t = 1.0)]
    tx_cost_bps: f64,

    /// Instrument tick size in price units (ES=0.25, NQ=0.25, ZN=1/64)
    #[arg(long, default_value_t = 0.25)]
    tick_size: f64,

    /// Typical bid-ask spread in ticks (ES=1.0, ZN=1.0, 6E=2.0)
    /// Half-spread slippage is added to `tx_cost_bps` per trade.
    #[arg(long, default_value_t = 1.0)]
    spread_ticks: f64,

    /// Number of top hyperopt configs to train as ensemble (default 1 = best only).
    /// When > 1, loads `top_k_params` from the hyperopt JSON and trains a separate
    /// model for each param set, saving as `dqn_ensemble_{k}_fold_{fold}.safetensors`.
    #[arg(long, default_value_t = 1)]
    ensemble_top_k: usize,

    /// Optional path to MBP-10 order book data directory for OFI features.
    /// When set, enables 8 OFI features (OFI L1/L5, depth imbalance, VPIN, etc.)
    /// expanding state dimension from 43 to 51.
    #[arg(long)]
    mbp10_data_dir: Option<PathBuf>,

    /// Optional path to trade data directory (.dbn.zst files with Schema::Trades).
    /// When set, real trade buy/sell classification feeds VPIN and Kyle's Lambda
    /// instead of the tick-rule proxy. Requires --mbp10-data-dir to take effect.
    #[arg(long)]
    trades_data_dir: Option<PathBuf>,

    /// Enable offline RL mode: train exclusively from a pre-collected dataset,
    /// skipping online experience collection. Requires --dataset-path.
    #[arg(long, default_value_t = false)]
    offline: bool,

    /// Path to a pre-collected experience dataset (bincode format).
    /// Used with --offline to load a fixed dataset into the replay buffer.
    #[arg(long)]
    dataset_path: Option<PathBuf>,

    /// Collect and save a dataset from the current policy, then exit.
    /// Use this to generate datasets for offline training.
    #[arg(long)]
    collect_dataset: Option<PathBuf>,
}

// ---------------------------------------------------------------------------
// Hyperopt parameter loading
// ---------------------------------------------------------------------------

/// Load `best_params` from a hyperopt results JSON file.
///
/// Expected format: `{ "model_key": { "best_params": { ... }, ... } }`
/// Returns `None` if the file doesn't exist or can't be parsed.
#[allow(clippy::cognitive_complexity)]
fn load_hyperopt_params(hp_path: &Option<PathBuf>, model_key: &str) -> Option<Value> {
    let file_path = hp_path.as_ref()?;
    if !file_path.exists() {
        info!("Hyperopt params file not found: {}, using defaults", file_path.display());
        return None;
    }
    let contents = match std::fs::read_to_string(file_path) {
        Ok(c) => c,
        Err(e) => {
            warn!("Failed to read hyperopt params {}: {}", file_path.display(), e);
            return None;
        }
    };
    let json: Value = match serde_json::from_str(&contents) {
        Ok(v) => v,
        Err(e) => {
            warn!("Failed to parse hyperopt params JSON: {}", e);
            return None;
        }
    };
    let params = json.get(model_key)
        .and_then(|m| m.get("best_params"))
        .cloned();
    if params.is_some() {
        info!("Loaded hyperopt params for '{}' from {}", model_key, file_path.display());
    } else {
        warn!("No best_params found for '{}' in {}", model_key, file_path.display());
    }
    params
}

/// Load top-K param sets from a hyperopt results JSON file.
///
/// Expected format: `{ "model_key": { "top_k_params": [{ "params": {...}, ... }, ...], "best_params": {...} } }`
/// Falls back to `best_params` as a single entry when `top_k_params` is absent.
/// Returns a vec of length `k` (or fewer if not enough entries), each entry `Some(params)` or `None`.
fn load_top_k_params(hp_path: &Option<PathBuf>, model_key: &str, k: usize) -> Vec<Option<Value>> {
    let file_path = match hp_path.as_ref() {
        Some(p) if p.exists() => p,
        _ => return vec![None; k],
    };
    let Ok(contents) = std::fs::read_to_string(file_path) else {
        return vec![None; k];
    };
    let Ok(json): Result<Value, _> = serde_json::from_str(&contents) else {
        return vec![None; k];
    };
    let top_k = json
        .get(model_key)
        .and_then(|m| m.get("top_k_params"))
        .and_then(|v| v.as_array());
    if let Some(arr) = top_k {
        arr.iter()
            .take(k)
            .map(|entry| entry.get("params").cloned())
            .collect()
    } else {
        // Fallback: just use best_params as the single entry
        let best = json
            .get(model_key)
            .and_then(|m| m.get("best_params"))
            .cloned();
        vec![best]
    }
}

fn hp_f64(params: &Option<Value>, key: &str) -> Option<f64> {
    params.as_ref()?.get(key)?.as_f64()
}

fn hp_usize(params: &Option<Value>, key: &str) -> Option<usize> {
    params.as_ref()?.get(key)?.as_u64().map(|v| v as usize)
}

fn hp_bool(params: &Option<Value>, key: &str) -> Option<bool> {
    params.as_ref()?.get(key)?.as_bool()
}

// ---------------------------------------------------------------------------
// Fold data preparation (extracted for prefetching)
// ---------------------------------------------------------------------------

/// Prepared fold data: (`train_features`, `val_features`, `train_bars`, `val_bars`)
type FoldData = (Vec<FeatureVector>, Vec<FeatureVector>, Vec<OHLCVBar>, Vec<OHLCVBar>);

#[allow(clippy::cognitive_complexity)]
/// Prepare a fold's data for training: extract features, normalize, align bars.
///
/// Returns `None` if feature extraction fails or produces empty features.
fn prepare_fold_data(
    window: &ml::walk_forward::WalkForwardWindow,
    output_dir: &Path,
) -> Option<FoldData> {
    let train_feat = match extract_ml_features(&window.train) {
        Ok(f) => f,
        Err(e) => {
            warn!("  Fold {} -- train feature extraction failed: {}", window.fold, e);
            return None;
        }
    };
    let val_feat = match extract_ml_features(&window.val) {
        Ok(f) => f,
        Err(e) => {
            warn!("  Fold {} -- val feature extraction failed: {}", window.fold, e);
            return None;
        }
    };
    if train_feat.is_empty() || val_feat.is_empty() {
        warn!("  Fold {} -- empty features, skipping", window.fold);
        return None;
    }

    let norm_stats = NormStats::from_features(&train_feat);
    let train_norm = norm_stats.normalize_batch(&train_feat);
    let val_norm = norm_stats.normalize_batch(&val_feat);

    // Save NormStats for this fold
    let norm_path = output_dir.join(format!("norm_stats_fold{}.json", window.fold));
    match serde_json::to_string_pretty(&norm_stats) {
        Ok(json) => {
            let norm_tmp = norm_path.with_extension("json.tmp");
            if let Err(e) = std::fs::write(&norm_tmp, &json) {
                error!("  Failed to write NormStats tmp {}: {}", norm_tmp.display(), e);
                return None;
            }
            if let Err(e) = std::fs::rename(&norm_tmp, &norm_path) {
                error!("  Failed to rename NormStats {} -> {}: {}", norm_tmp.display(), norm_path.display(), e);
                drop(std::fs::remove_file(&norm_tmp));
                return None;
            }
            info!("  Saved NormStats to {}", norm_path.display());
        }
        Err(e) => {
            error!("  Failed to serialize NormStats: {}", e);
            return None;
        }
    }

    let train_warmup = window.train.len().saturating_sub(train_norm.len());
    let train_bars_aligned = window.train.get(train_warmup..).unwrap_or(&window.train).to_vec();
    let val_warmup = window.val.len().saturating_sub(val_norm.len());
    let val_bars_aligned = window.val.get(val_warmup..).unwrap_or(&window.val).to_vec();

    Some((train_norm, val_norm, train_bars_aligned, val_bars_aligned))
}

// ---------------------------------------------------------------------------
// DQN Training
// ---------------------------------------------------------------------------

/// Convert pre-aligned features and bars into the format expected by
/// `DQNTrainer::train_with_preloaded_data`: `Vec<(FeatureVector, Vec<f64>)>`.
///
/// The `Vec<f64>` target contains OHLCV prices for the trainer's internal
/// reward computation and portfolio simulation.
fn features_to_trainer_format(
    features: &[FeatureVector],
    bars: &[OHLCVBar],
) -> Vec<(FeatureVector, Vec<f64>)> {
    features
        .iter()
        .zip(bars.iter())
        .map(|(feat, bar)| {
            (*feat, vec![bar.open, bar.high, bar.low, bar.close])
        })
        .collect()
}

/// Train a DQN model on a single walk-forward fold using `DQNTrainer`.
///
/// Delegates all GPU optimizations (mixed precision, dynamic batching, gradient
/// accumulation, Rainbow DQN components) to the production trainer. The walk-forward
/// fold structure stays in this binary; only per-fold training delegates to `DQNTrainer`.
///
/// Returns the best validation loss achieved.
#[allow(clippy::cognitive_complexity, clippy::too_many_arguments)]
fn train_dqn_fold(
    fold: usize,
    train_features: &[FeatureVector],
    val_features: &[FeatureVector],
    train_bars: &[OHLCVBar],
    val_bars: &[OHLCVBar],
    args: &Args,
    output_dir: &Path,
    hp: &Option<Value>,
    pre_uploaded_gpu_data: Option<DqnGpuData>,
    checkpoint_prefix: &str,
) -> Result<f64> {
    // Caller must pre-align bars to features (skip warmup period before calling).
    debug_assert_eq!(
        train_bars.len(),
        train_features.len(),
        "train_bars ({}) must be pre-aligned to train_features ({})",
        train_bars.len(),
        train_features.len(),
    );
    debug_assert_eq!(
        val_bars.len(),
        val_features.len(),
        "val_bars ({}) must be pre-aligned to val_features ({})",
        val_bars.len(),
        val_features.len(),
    );
    info!("  [DQN] Fold {} -- {} train, {} val features", fold, train_features.len(), val_features.len());

    // Compute average spread slippage in bps from training bar prices.
    // Same pattern as evaluate_baseline: total_cost = commission + spread.
    let avg_price = if train_bars.is_empty() {
        0.0
    } else {
        train_bars.iter().map(|b| b.close).sum::<f64>() / train_bars.len() as f64
    };
    let avg_spread_bps = spread_cost_bps(avg_price, args.tick_size, args.spread_ticks);
    let total_cost_bps = args.tx_cost_bps + avg_spread_bps;
    info!("  [DQN] Fold {} total tx cost: {:.2} bps (commission {:.1} + spread {:.2})",
        fold, total_cost_bps, args.tx_cost_bps, avg_spread_bps);

    // Build DQNHyperparameters -- hyperopt JSON overrides defaults, CLI args override both.
    // `..DQNHyperparameters::default()` enables all Rainbow components (PER, dueling,
    // distributional C51, noisy nets, gradient accumulation, mixed precision auto-detect).
    // transaction_cost_multiplier scales internal reward cost by total trading friction.
    //
    // hidden_dim_base: hyperopt JSON takes priority. When absent, use VRAM-aware
    // scaling so that large GPUs (L40S 48GB, H100 80GB) get proportionally wider
    // networks instead of the tiny CPU defaults.
    let hp_hidden_base = hp_usize(hp, "hidden_dim_base");
    let dqn_hidden_base = hp_hidden_base.or_else(|| {
        let vram = detect_vram_mb();
        let base = vram_scaled_base_dim(vram, "dqn");
        info!("  [DQN] VRAM-scaled hidden_dim_base: {} (VRAM: {} MB)", base, vram);
        Some(base)
    });

    // When noisy_nets are enabled (the hyperopt default), epsilon-greedy must be
    // low or zero — otherwise epsilon=1.0 forces pure random actions for the entire
    // run, preventing the model from ever learning.  Read exploration params from
    // hyperopt JSON so walk-forward uses the same strategy that won the search.
    let use_noisy = hp_bool(hp, "use_noisy_nets").unwrap_or(true);
    let epsilon_start = if use_noisy {
        // Noisy nets handle exploration; epsilon just adds noise on top.
        // Use hyperopt value or a small default (not 1.0!)
        hp_f64(hp, "epsilon_start").unwrap_or(0.05)
    } else {
        hp_f64(hp, "epsilon_start").unwrap_or(1.0)
    };

    let hyperparams = DQNHyperparameters {
        learning_rate: hp_f64(hp, "learning_rate").unwrap_or(args.learning_rate),
        batch_size: hp_usize(hp, "batch_size").unwrap_or(args.batch_size),
        gamma: hp_f64(hp, "gamma").unwrap_or(0.95),
        epsilon_start,
        epsilon_end: hp_f64(hp, "epsilon_end").unwrap_or(0.01),
        epsilon_decay: hp_f64(hp, "epsilon_decay").unwrap_or(0.995),
        buffer_size: hp_usize(hp, "buffer_size").unwrap_or(50_000),
        min_replay_size: hp_usize(hp, "batch_size").unwrap_or(args.batch_size),
        epochs: args.epochs,
        checkpoint_frequency: 10,
        hidden_dim_base: dqn_hidden_base,
        warmup_steps: 0,
        early_stopping_enabled: true,
        transaction_cost_multiplier: total_cost_bps,
        // Pass through hyperopt architectural choices
        use_per: hp_bool(hp, "use_per").unwrap_or(true),
        per_alpha: hp_f64(hp, "per_alpha").unwrap_or(0.6),
        per_beta_start: hp_f64(hp, "per_beta_start").unwrap_or(0.4),
        use_dueling: hp_bool(hp, "use_dueling").unwrap_or(true),
        dueling_hidden_dim: hp_usize(hp, "dueling_hidden_dim").unwrap_or(128),
        n_steps: hp_usize(hp, "n_steps").unwrap_or(3),
        tau: hp_f64(hp, "tau").unwrap_or(0.005),
        use_distributional: hp_bool(hp, "use_distributional").unwrap_or(true),
        num_atoms: hp_usize(hp, "num_atoms").unwrap_or(51),
        v_min: hp_f64(hp, "v_min").unwrap_or(-2.0),
        v_max: hp_f64(hp, "v_max").unwrap_or(2.0),
        use_noisy_nets: use_noisy,
        noisy_sigma_init: hp_f64(hp, "noisy_sigma_init").unwrap_or(0.5),
        use_qr_dqn: hp_bool(hp, "use_qr_dqn").unwrap_or(false),
        num_quantiles: hp_usize(hp, "num_quantiles").unwrap_or(64),
        noisy_epsilon_floor: hp_f64(hp, "noisy_epsilon_floor").unwrap_or(0.05).into(),
        // Reward-shaping & environment params — critical for trade generation
        hold_penalty_weight: hp_f64(hp, "hold_penalty_weight").unwrap_or(0.01),
        max_position_absolute: hp_f64(hp, "max_position_absolute").unwrap_or(2.0),
        huber_delta: hp_f64(hp, "huber_delta").unwrap_or(10.0),
        entropy_coefficient: hp_f64(hp, "entropy_coefficient").map(Some).unwrap_or(Some(0.01)),
        curiosity_weight: hp_f64(hp, "curiosity_weight").unwrap_or(0.1),
        weight_decay: hp_f64(hp, "weight_decay").unwrap_or(1e-4),
        kelly_fractional: hp_f64(hp, "kelly_fractional").unwrap_or(0.5),
        kelly_max_fraction: hp_f64(hp, "kelly_max_fraction").unwrap_or(0.25),
        mbp10_data_dir: args.mbp10_data_dir.as_ref().map(|p| p.to_string_lossy().into_owned()),
        trades_data_dir: args.trades_data_dir.as_ref().map(|p| p.to_string_lossy().into_owned()),
        offline_mode: args.offline,
        dataset_path: args.dataset_path.as_ref().map(|p| p.to_string_lossy().into_owned()),
        ..DQNHyperparameters::default()
    };

    // Create DQNTrainer -- auto-detects GPU, mixed precision, dynamic batch sizing
    let mut trainer = DQNTrainer::new(hyperparams)
        .context("Failed to create DQNTrainer")?;

    // Inject pre-uploaded GPU data from DoubleBufferedLoader (overlapped with previous fold)
    if let Some(gpu_data) = pre_uploaded_gpu_data {
        info!("  [DQN] Fold {} -- injecting pre-uploaded GPU data ({} bars), skipping lazy upload",
            fold, gpu_data.num_bars);
        trainer.set_gpu_data(gpu_data);
    }

    // Convert features + bars to trainer format
    let training_data = features_to_trainer_format(train_features, train_bars);
    let val_data = features_to_trainer_format(val_features, val_bars);

    // Checkpoint callback: save best model to output directory
    let output_dir_owned = output_dir.to_path_buf();
    let prefix_owned = checkpoint_prefix.to_owned();
    let checkpoint_callback = move |epoch: usize, data: Vec<u8>, is_best: bool| -> Result<String> {
        let suffix = if is_best { "best" } else { &format!("epoch{}", epoch) };
        let ckpt_path = output_dir_owned.join(format!("{}_fold{}_{}.safetensors", prefix_owned, fold, suffix));
        // Atomic write: write to .tmp then rename (POSIX rename is atomic)
        let tmp_path = ckpt_path.with_extension("safetensors.tmp");
        std::fs::write(&tmp_path, &data)
            .with_context(|| format!("Failed to write checkpoint tmp: {}", tmp_path.display()))?;
        if let Err(e) = std::fs::rename(&tmp_path, &ckpt_path) {
            drop(std::fs::remove_file(&tmp_path));
            return Err(e).with_context(|| format!("Failed to rename checkpoint: {} -> {}", tmp_path.display(), ckpt_path.display()));
        }
        info!("  [DQN] Fold {} saved checkpoint: {} (prefix: {})", fold, ckpt_path.display(), prefix_owned);
        Ok(ckpt_path.to_string_lossy().into_owned())
    };

    // Run the async training loop via a tokio runtime.
    // The trainer handles epochs, early stopping, epsilon decay, validation, and
    // all Rainbow DQN components internally.
    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .context("Failed to create tokio runtime for DQN training")?;

    let metrics = rt.block_on(
        trainer.train_with_preloaded_data(training_data, val_data, checkpoint_callback)
    ).context("DQNTrainer training failed")?;

    info!(
        "  [DQN] Fold {} complete -- loss={:.6} epochs_trained={} converged={}",
        fold, metrics.loss, metrics.epochs_trained, metrics.convergence_achieved
    );

    Ok(metrics.loss)
}

// ---------------------------------------------------------------------------
// PPO Training
// ---------------------------------------------------------------------------

/// Train a PPO model on a single walk-forward fold using `PpoTrainer`.
///
/// Delegates all GPU optimizations (dynamic batch sizing, mixed precision when
/// available, GAE computation, trajectory collection, early stopping) to the
/// production trainer. The walk-forward fold structure stays in this binary;
/// only per-fold training delegates to `PpoTrainer`.
///
/// Returns the best validation loss achieved (approximated by final `value_loss`).
#[allow(clippy::cognitive_complexity)]
fn train_ppo_fold(
    fold: usize,
    train_features: &[FeatureVector],
    val_features: &[FeatureVector],
    train_bars: &[OHLCVBar],
    _val_bars: &[OHLCVBar],
    args: &Args,
    output_dir: &Path,
    hp: &Option<Value>,
) -> Result<f64> {
    info!("  [PPO] Fold {} -- {} train, {} val features", fold, train_features.len(), val_features.len());

    let n_train = train_features.len();
    if n_train < 2 {
        warn!("  [PPO] Fold {} -- insufficient training features ({})", fold, n_train);
        return Ok(f64::MAX);
    }

    // Compute average spread slippage in bps from training bar prices.
    // Same pattern as evaluate_baseline: total_cost = commission + spread.
    let avg_price = if train_bars.is_empty() {
        0.0
    } else {
        train_bars.iter().map(|b| b.close).sum::<f64>() / train_bars.len() as f64
    };
    let avg_spread_bps = spread_cost_bps(avg_price, args.tick_size, args.spread_ticks);
    let total_cost_bps = args.tx_cost_bps + avg_spread_bps;
    info!("  [PPO] Fold {} total tx cost: {:.2} bps (commission {:.1} + spread {:.2})",
        fold, total_cost_bps, args.tx_cost_bps, avg_spread_bps);

    // Build PpoHyperparameters -- hyperopt JSON overrides defaults, CLI args override both.
    // Start from conservative() baseline so all fields have sane values.
    //
    // hidden_dim_base: hyperopt JSON takes priority. When absent, use VRAM-aware
    // scaling so that large GPUs get proportionally wider networks. PPO uses the
    // "ppo_policy" model type for base dim since the trainer internally scales
    // the value network from the same base (value: [4*base, 3*base, 2*base, base, base/2]).
    let hp_ppo_hidden_base = hp_usize(hp, "hidden_dim_base");
    let ppo_hidden_base = hp_ppo_hidden_base.or_else(|| {
        let vram = detect_vram_mb();
        let base = vram_scaled_base_dim(vram, "ppo_policy");
        info!("  [PPO] VRAM-scaled hidden_dim_base: {} (VRAM: {} MB)", base, vram);
        Some(base)
    });

    let hyperparams = PpoHyperparameters {
        learning_rate: hp_f64(hp, "learning_rate").unwrap_or(args.learning_rate),
        actor_learning_rate: Some(hp_f64(hp, "policy_learning_rate").unwrap_or(args.learning_rate)),
        critic_learning_rate: Some(hp_f64(hp, "value_learning_rate").unwrap_or(args.learning_rate * 3.0)),
        batch_size: hp_usize(hp, "batch_size").unwrap_or(args.batch_size.max(64)),
        gamma: hp_f64(hp, "gamma").unwrap_or(0.99),
        clip_epsilon: hp_f64(hp, "clip_epsilon").unwrap_or(0.2) as f32,
        vf_coef: hp_f64(hp, "value_loss_coeff").unwrap_or(0.5) as f32,
        ent_coef: hp_f64(hp, "entropy_coeff").unwrap_or(0.01) as f32,
        gae_lambda: hp_f64(hp, "gae_lambda").unwrap_or(0.95) as f32,
        rollout_steps: hp_usize(hp, "rollout_steps").unwrap_or(
            if args.max_steps_per_epoch > 0 { args.max_steps_per_epoch } else { 2048 }
        ),
        minibatch_size: hp_usize(hp, "minibatch_size").unwrap_or(64),
        epochs: args.epochs,
        early_stopping_enabled: true,
        transaction_cost_bps: total_cost_bps / 100.0, // total (commission + spread) in bps → pct for trainer
        hidden_dim_base: ppo_hidden_base,
        ..PpoHyperparameters::conservative()
    };

    // Checkpoint directory for this fold
    let fold_ckpt_dir = output_dir.join(format!("ppo_fold{}", fold));

    // Create PpoTrainer -- auto-detects GPU, mixed precision, dynamic batch sizing
    let trainer = PpoTrainer::new(
        hyperparams,
        args.feature_dim,
        &fold_ckpt_dir,
        true, // use_gpu = true, trainer auto-detects availability
        None, // num_envs: standard single-env collection
    ).context("Failed to create PpoTrainer")?;

    // Convert features from FeatureVector to Vec<Vec<f32>> for the trainer
    let market_data: Vec<Vec<f32>> = train_features
        .iter()
        .map(|feat| feat.iter().map(|&v| v as f32).collect())
        .collect();

    // Progress callback: log per-epoch metrics
    let progress_callback = move |metrics: ml::trainers::ppo::PpoTrainingMetrics| {
        info!(
            "  [PPO] Fold {} Epoch {}/{} -- policy_loss={:.6} value_loss={:.6} expl_var={:.4} mean_reward={:.4}",
            fold,
            metrics.epoch,
            args.epochs,
            metrics.policy_loss,
            metrics.value_loss,
            metrics.explained_variance,
            metrics.mean_reward,
        );
    };

    // Run the async training loop via a tokio runtime.
    // PpoTrainer handles epochs, early stopping, trajectory collection, GAE,
    // checkpointing, and all GPU optimizations internally.
    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .context("Failed to create tokio runtime for PPO training")?;

    let metrics = rt.block_on(
        trainer.train(market_data, progress_callback)
    ).context("PpoTrainer training failed")?;

    info!(
        "  [PPO] Fold {} complete -- value_loss={:.6} policy_loss={:.6} expl_var={:.4} epochs_trained={}",
        fold, metrics.value_loss, metrics.policy_loss, metrics.explained_variance, metrics.epoch
    );

    Ok(metrics.value_loss as f64)
}

// ---------------------------------------------------------------------------
// Training orchestration
// ---------------------------------------------------------------------------

/// Container for per-model RL results collected during training.
struct RlTrainingResult {
    model_name: String,
    fold_results: Vec<(usize, f64)>,
    total_epochs: usize,
}

/// Run the full walk-forward RL training pipeline.
///
/// Returns per-model results so `main()` can write completion markers.
#[allow(clippy::cognitive_complexity, clippy::too_many_lines)]
fn run_training(args: &Args) -> Result<Vec<RlTrainingResult>> {
    let train_dqn = args.model == "dqn" || args.model == "both";
    let train_ppo = args.model == "ppo" || args.model == "both";

    info!("=== Walk-Forward Baseline Training ===");
    info!("  Model(s): {}", args.model);
    info!("  Symbol: {}", args.symbol);
    info!("  Epochs: {}", args.epochs);
    info!("  Batch size: {}", args.batch_size);
    info!("  Data dir: {}", args.data_dir.display());
    info!("  Output dir: {}", args.output_dir.display());
    info!("  Feature dim: {}", args.feature_dim);
    info!("  Num actions: {}", args.num_actions);
    info!("  Learning rate: {:.1e}", args.learning_rate);
    info!("  Tx cost: {:.1} bps commission + {:.1} tick spread (tick_size={:.4})",
        args.tx_cost_bps, args.spread_ticks, args.tick_size);
    info!("  Patience: {}", args.patience);
    if let Some(ref hp_path) = args.hyperopt_params {
        info!("  Hyperopt params: {}", hp_path.display());
    }
    if args.ensemble_top_k > 1 {
        info!("  Ensemble top-K: {} (training multiple models per fold)", args.ensemble_top_k);
    }

    // 1. Load all OHLCV bars from DBN files
    info!("Step 1/5: Loading OHLCV bars from DBN files...");
    let data_load_start = std::time::Instant::now();
    let bars = load_all_bars(&args.data_dir, &args.symbol)?;
    if bars.is_empty() {
        anyhow::bail!("No bars loaded from {}", args.data_dir.display());
    }
    info!("  Loaded {} bars ({} to {})",
        bars.len(),
        bars.first().map(|b| b.timestamp.to_string()).unwrap_or_default(),
        bars.last().map(|b| b.timestamp.to_string()).unwrap_or_default(),
    );

    // 2. Extract features
    info!("Step 2/5: Extracting {}-dimensional features...", args.feature_dim);
    let all_features = extract_ml_features(&bars)
        .context("Feature extraction failed")?;
    info!("  Extracted {} feature vectors (warmup period consumed {} bars)",
        all_features.len(),
        bars.len().saturating_sub(all_features.len()),
    );

    // Since features skip the warmup period, we need bars aligned to features.
    // Features start at bar index warmup_offset (typically 50).
    let warmup_offset = bars.len().saturating_sub(all_features.len());
    let aligned_bars = bars.get(warmup_offset..).unwrap_or(&bars);

    // 3. Generate walk-forward windows
    info!("Step 3/5: Generating walk-forward windows...");
    let wf_config = WalkForwardConfig {
        initial_train_months: args.train_months,
        val_months: args.val_months,
        test_months: args.test_months,
        step_months: args.step_months,
    };
    let windows = generate_walk_forward_windows(aligned_bars, &wf_config);
    if windows.is_empty() {
        anyhow::bail!(
            "No walk-forward windows generated. Need at least {} months of data.",
            wf_config.initial_train_months + wf_config.val_months + wf_config.test_months
        );
    }
    info!("  Generated {} walk-forward folds", windows.len());

    // Record data loading + feature extraction time
    if train_dqn {
        metrics::record_data_load("dqn", data_load_start.elapsed().as_secs_f64());
    }
    if train_ppo {
        metrics::record_data_load("ppo", data_load_start.elapsed().as_secs_f64());
    }

    // Create output directory
    std::fs::create_dir_all(&args.output_dir)
        .with_context(|| format!("Failed to create output dir: {}", args.output_dir.display()))?;

    // 4. Train each fold
    info!("Step 4/5: Training models on each fold...");
    let mut dqn_results: Vec<(usize, f64)> = Vec::new();
    let mut ppo_results: Vec<(usize, f64)> = Vec::new();

    let mut prefetched_data: Option<FoldData> = None;
    // GPU double-buffer: pre-uploaded GPU data for the next DQN fold.
    // While fold N trains on GPU, fold N+1's data is uploaded in background.
    let mut dqn_gpu_staged: Option<DqnGpuData> = None;

    for (fold_idx, window) in windows.iter().enumerate() {
        info!("--- Fold {} ---", window.fold);
        info!(
            "  Train: {} bars (up to {}), Val: {} bars (up to {}), Test: {} bars (up to {})",
            window.train.len(),
            window.train_end,
            window.val.len(),
            window.val_end,
            window.test.len(),
            window.test_end,
        );

        // Get fold data: from prefetcher (fold 1+) or prepare fresh (fold 0)
        let fold_data = if let Some(data) = prefetched_data.take() {
            info!("  Using prefetched data for fold {}", window.fold);
            Some(data)
        } else {
            prepare_fold_data(window, &args.output_dir)
        };

        // Kick off prefetch for NEXT fold on background thread
        let prefetch_rx = if fold_idx + 1 < windows.len() {
            let next_window = windows.get(fold_idx + 1).cloned();
            if let Some(nw) = next_window {
                let out_dir = args.output_dir.clone();
                let (tx, rx) = std::sync::mpsc::channel::<Option<FoldData>>();
                let _handle = std::thread::Builder::new()
                    .name("fold-prefetch".into())
                    .spawn(move || {
                        info!("  [Prefetch] Loading fold {} data on background thread", nw.fold);
                        let result = prepare_fold_data(&nw, &out_dir);
                        drop(tx.send(result));
                    });
                Some(rx)
            } else {
                None
            }
        } else {
            None
        };

        let Some((train_norm, val_norm, train_bars_aligned, val_bars_aligned)) = fold_data else {
            // Collect prefetch result even if this fold failed
            if let Some(rx) = prefetch_rx {
                if let Ok(data) = rx.recv() {
                    prefetched_data = data;
                }
            }
            continue;
        };

        // Train DQN
        if train_dqn {
            let fold_str = fold_idx.to_string();
            let fold_start = std::time::Instant::now();

            // Take pre-uploaded GPU data from previous fold's background upload
            let mut gpu_data_for_fold = dqn_gpu_staged.take();

            if args.ensemble_top_k > 1 && args.hyperopt_params.is_some() {
                // Ensemble mode: train one model per top-K hyperopt param set
                let param_sets = load_top_k_params(
                    &args.hyperopt_params,
                    "dqn",
                    args.ensemble_top_k,
                );
                for (k, hp) in param_sets.iter().enumerate() {
                    info!(
                        "  [DQN] Training ensemble member {}/{} on fold {}",
                        k + 1,
                        param_sets.len(),
                        window.fold
                    );
                    let prefix = format!("dqn_ensemble_{}", k);
                    // Only the first ensemble member uses the pre-uploaded GPU data
                    let gpu_data = if k == 0 { gpu_data_for_fold.take() } else { None };
                    match train_dqn_fold(
                        window.fold,
                        &train_norm,
                        &val_norm,
                        &train_bars_aligned,
                        &val_bars_aligned,
                        args,
                        &args.output_dir,
                        hp,
                        gpu_data,
                        &prefix,
                    ) {
                        Ok(best_loss) => {
                            info!(
                                "  [DQN] Ensemble member {} fold {} best_loss={:.6}",
                                k, window.fold, best_loss
                            );
                            // Record the best ensemble member's loss as the fold result
                            if k == 0 {
                                let elapsed = fold_start.elapsed().as_secs_f64();
                                metrics::set_epoch("dqn", &fold_str, fold_idx as f64);
                                metrics::set_epoch_loss("dqn", &fold_str, best_loss);
                                metrics::set_validation_loss("dqn", &fold_str, best_loss);
                                metrics::set_iteration_seconds("dqn", &fold_str, elapsed);
                                dqn_results.push((window.fold, best_loss));
                            }
                        }
                        Err(e) => {
                            error!(
                                "  [DQN] Ensemble member {} fold {} failed: {}",
                                k, window.fold, e
                            );
                        }
                    }
                }
            } else {
                // Single-model mode (default)
                let hp = load_hyperopt_params(&args.hyperopt_params, "dqn");
                match train_dqn_fold(
                    window.fold,
                    &train_norm,
                    &val_norm,
                    &train_bars_aligned,
                    &val_bars_aligned,
                    args,
                    &args.output_dir,
                    &hp,
                    gpu_data_for_fold,
                    "dqn",
                ) {
                    Ok(best_loss) => {
                        let elapsed = fold_start.elapsed().as_secs_f64();
                        metrics::set_epoch("dqn", &fold_str, fold_idx as f64);
                        metrics::set_epoch_loss("dqn", &fold_str, best_loss);
                        metrics::set_validation_loss("dqn", &fold_str, best_loss);
                        metrics::set_iteration_seconds("dqn", &fold_str, elapsed);
                        dqn_results.push((window.fold, best_loss));
                    }
                    Err(e) => {
                        error!("  [DQN] Fold {} failed: {}", window.fold, e);
                    }
                }
            }
        }

        // Train PPO
        if train_ppo {
            let hp = load_hyperopt_params(&args.hyperopt_params, "ppo");
            let fold_str = fold_idx.to_string();
            let fold_start = std::time::Instant::now();
            match train_ppo_fold(
                window.fold,
                &train_norm,
                &val_norm,
                &train_bars_aligned,
                &val_bars_aligned,
                args,
                &args.output_dir,
                &hp,
            ) {
                Ok(best_loss) => {
                    let elapsed = fold_start.elapsed().as_secs_f64();
                    metrics::set_epoch("ppo", &fold_str, fold_idx as f64);
                    metrics::set_epoch_loss("ppo", &fold_str, best_loss);
                    metrics::set_validation_loss("ppo", &fold_str, best_loss);
                    metrics::set_iteration_seconds("ppo", &fold_str, elapsed);
                    ppo_results.push((window.fold, best_loss));
                }
                Err(e) => {
                    error!("  [PPO] Fold {} failed: {}", window.fold, e);
                }
            }
        }

        // Collect prefetch result for next iteration
        if let Some(rx) = prefetch_rx {
            match rx.recv() {
                Ok(Some(data)) => {
                    info!("  Prefetched fold {} data ready ({} train features)",
                        windows.get(fold_idx + 1).map_or(0, |w| w.fold), data.0.len());

                    // GPU double-buffer: pre-upload next fold's DQN data to GPU now
                    // while the result is fresh and GPU is idle between folds.
                    if train_dqn {
                        let next_train_data = features_to_trainer_format(&data.0, &data.2);
                        let device = Device::new_cuda(0).unwrap_or(Device::Cpu);
                        if device.is_cuda() {
                            match DqnGpuData::upload(&next_train_data, &device) {
                                Ok(gpu_data) => {
                                    info!("  [DoubleBuffer] Pre-uploaded {} bars to GPU for next fold ({:.1} MB)",
                                        gpu_data.num_bars,
                                        gpu_data.vram_bytes() as f64 / 1_048_576.0);
                                    dqn_gpu_staged = Some(gpu_data);
                                }
                                Err(e) => {
                                    warn!("  [DoubleBuffer] GPU pre-upload failed, will upload lazily: {}", e);
                                }
                            }
                        }
                    }

                    prefetched_data = Some(data);
                }
                Ok(None) => {
                    warn!("  Prefetch for next fold returned None (will load synchronously)");
                    prefetched_data = None;
                }
                Err(_) => {
                    warn!("  Prefetch thread disconnected (will load synchronously)");
                    prefetched_data = None;
                }
            }
        }
    }

    // 5. Summary
    info!("Step 5/5: Training Summary");
    info!("  ===================================");

    let mut all_results = Vec::new();

    if train_dqn {
        info!("  DQN Results ({} folds):", dqn_results.len());
        for (fold, loss) in &dqn_results {
            info!("    Fold {}: best_val_metric = {:.6}", fold, loss);
        }
        if !dqn_results.is_empty() {
            let avg: f64 = dqn_results.iter().map(|(_, l)| l).sum::<f64>()
                / dqn_results.len() as f64;
            info!("    Average: {:.6}", avg);
        }
        all_results.push(RlTrainingResult {
            model_name: "dqn".to_owned(),
            fold_results: dqn_results,
            total_epochs: windows.len() * args.epochs,
        });
    }
    if train_ppo {
        info!("  PPO Results ({} folds):", ppo_results.len());
        for (fold, loss) in &ppo_results {
            info!("    Fold {}: best_val_metric = {:.6}", fold, loss);
        }
        if !ppo_results.is_empty() {
            let avg: f64 = ppo_results.iter().map(|(_, l)| l).sum::<f64>()
                / ppo_results.len() as f64;
            info!("    Average: {:.6}", avg);
        }
        all_results.push(RlTrainingResult {
            model_name: "ppo".to_owned(),
            fold_results: ppo_results,
            total_epochs: windows.len() * args.epochs,
        });
    }
    info!("  Checkpoints saved to: {}", args.output_dir.display());
    info!("  ===================================");

    Ok(all_results)
}

// ---------------------------------------------------------------------------
// Main
// ---------------------------------------------------------------------------

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(
        "train_baseline_rl",
        otlp_endpoint.as_deref(),
    ) {
        eprintln!("Observability init failed (non-fatal): {e}");
    }

    metrics::init();
    metrics_server::start_metrics_server(9094);
    common::metrics::questdb_sink::init(None);
    metrics::set_active_workers(1.0);

    let args = Args::parse();

    // Ensure output directory exists before training so markers can always be written.
    if let Err(e) = std::fs::create_dir_all(&args.output_dir) {
        error!("Failed to create output dir {}: {}", args.output_dir.display(), e);
    }

    let result = run_training(&args);
    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, "train_baseline_rl") {
        tracing::warn!("Failed to push metrics to gateway (non-fatal): {e}");
    }
    common::metrics::questdb_sink::flush();

    match result {
        Ok(results) => {
            for training_result in &results {
                let best_val = training_result
                    .fold_results
                    .iter()
                    .map(|(_, loss)| *loss)
                    .fold(f64::MAX, f64::min);

                let metrics = CompletionMetrics {
                    model: training_result.model_name.clone(),
                    symbol: args.symbol.clone(),
                    best_val_loss: (best_val < f64::MAX).then_some(best_val),
                    sharpe_ratio: None,
                    epochs_completed: training_result.total_epochs,
                    folds_completed: training_result.fold_results.len(),
                };
                write_success_marker(&args.output_dir, &metrics);
            }
            Ok(())
        }
        Err(e) => {
            let msg = format!("{:#}", e);
            error!("Training failed: {}", msg);
            write_failure_marker(&args.output_dir, &msg);
            Err(e)
        }
    }
}