foxhunt/ml/tests/dqn_hyperopt_test.rs

//! DQN Hyperopt End-to-End Test
//!
//! Proves that the hyperparameter optimizer works end-to-end with real DQN training.
//! This test runs 10 trials of Bayesian optimization (Argmin PSO) and verifies
//! convergence, result structure, and checkpoint artifacts.
//!
//! Run manually:
//!   SQLX_OFFLINE=true cargo test -p ml --test dqn_hyperopt_test -- --ignored --nocapture
//!
//! Expected runtime: 10-20 minutes (GPU), 30-60 minutes (CPU)

#![allow(unused_crate_dependencies)]

use anyhow::{Context, Result};
use std::path::PathBuf;

use ml::hyperopt::adapters::dqn::DQNTrainer;
use ml::hyperopt::ArgminOptimizer;

/// Locate the real training data directory.
///
/// Returns `Ok(path)` if the directory exists, `Err` otherwise so the test
/// can skip gracefully without marking as failed.
fn get_data_dir() -> Result<PathBuf> {
    let workspace_root = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .parent()
        .context("Failed to resolve workspace root")?
        .to_path_buf();

    let data_dir = workspace_root.join("test_data/real/databento/ml_training_small");
    if !data_dir.exists() {
        anyhow::bail!(
            "Training data not found at {}. \
             This test requires real Databento data to run.",
            data_dir.display()
        );
    }
    Ok(data_dir)
}

#[test]
#[ignore]
fn test_dqn_hyperopt_10_trials() -> Result<()> {
    // --- Setup ---
    let data_dir = match get_data_dir() {
        Ok(dir) => dir,
        Err(e) => {
            eprintln!("Skipping test: {e}");
            return Ok(());
        }
    };

    let start = std::time::Instant::now();
    let num_trials: usize = 10;
    let n_initial: usize = 3;
    let epochs_per_trial: usize = 5;

    // --- Create trainer and optimizer ---
    let trainer = DQNTrainer::new(&data_dir, epochs_per_trial)
        .context("Failed to create DQNTrainer")?;

    let optimizer = ArgminOptimizer::builder()
        .max_trials(num_trials)
        .n_initial(n_initial)
        .n_particles(5) // small swarm for test speed
        .seed(42)
        .build();

    // --- Run optimization ---
    let result = optimizer
        .optimize(trainer)
        .context("Hyperopt optimization failed")?;

    let elapsed = start.elapsed();

    // --- Assertions ---

    // 1. All trials should have completed (at least n_initial; PSO may add more)
    let trials_completed = result.all_trials.len();
    assert!(
        trials_completed >= n_initial,
        "Expected at least {n_initial} trials, got {trials_completed}"
    );

    // 2. Best objective should be finite (not NaN or Inf)
    assert!(
        result.best_objective.is_finite(),
        "Best objective is not finite: {}",
        result.best_objective
    );

    // 3. Every trial objective should be finite
    for trial in &result.all_trials {
        assert!(
            trial.objective.is_finite(),
            "Trial {} has non-finite objective: {}",
            trial.trial_num,
            trial.objective
        );
        assert!(
            trial.duration_secs > 0.0,
            "Trial {} has non-positive duration: {}",
            trial.trial_num,
            trial.duration_secs
        );
    }

    // 4. Best objective should be <= first trial (optimizer should not regress)
    if let Some(first_trial) = result.all_trials.first() {
        assert!(
            result.best_objective <= first_trial.objective,
            "Optimizer regressed: best={} > first={}",
            result.best_objective,
            first_trial.objective
        );
    }

    // 5. Convergence plot should have entries matching trial count
    assert_eq!(
        result.convergence_plot_data.len(),
        trials_completed,
        "Convergence plot entries ({}) should match trial count ({trials_completed})",
        result.convergence_plot_data.len()
    );

    // 6. Convergence plot should be monotonically non-increasing
    for window in result.convergence_plot_data.windows(2) {
        let (_, prev_best) = window[0];
        let (_, curr_best) = window[1];
        assert!(
            curr_best <= prev_best + f64::EPSILON,
            "Convergence plot is not monotonically non-increasing: {prev_best} -> {curr_best}"
        );
    }

    // 7. Best params should have sane learning rate (log-scale, typically 1e-5 to 1e-2)
    let best_lr = result.best_params.learning_rate;
    assert!(
        best_lr > 1e-7 && best_lr < 1.0,
        "Best learning rate out of sane range: {best_lr}"
    );

    // 8. Best params batch size should be positive and bounded
    let best_bs = result.best_params.batch_size;
    assert!(
        best_bs > 0 && best_bs <= 512,
        "Best batch size out of range: {best_bs}"
    );

    // --- Report ---
    println!("\n{}", "=".repeat(70));
    println!("  DQN HYPEROPT REPORT");
    println!("{}", "=".repeat(70));
    println!("  Trials completed : {trials_completed}");
    println!("  Best objective   : {:.6}", result.best_objective);
    println!("  Best LR          : {:.2e}", result.best_params.learning_rate);
    println!("  Best batch size  : {}", result.best_params.batch_size);
    println!("  Best gamma       : {:.4}", result.best_params.gamma);
    println!("  Best buffer size : {}", result.best_params.buffer_size);
    println!("  Total time       : {:.1}s", elapsed.as_secs_f64());
    println!(
        "  Avg time/trial   : {:.1}s",
        elapsed.as_secs_f64() / trials_completed as f64
    );
    println!("{}", "-".repeat(70));
    println!("  TRIAL HISTORY");
    println!("{}", "-".repeat(70));
    for trial in &result.all_trials {
        println!(
            "    Trial {:>2} | obj: {:>10.4} | time: {:>6.1}s | lr: {:.2e} | bs: {}",
            trial.trial_num,
            trial.objective,
            trial.duration_secs,
            trial.params.learning_rate,
            trial.params.batch_size,
        );
    }
    println!("{}", "-".repeat(70));
    println!("  CONVERGENCE");
    println!("{}", "-".repeat(70));
    for (trial_num, best_so_far) in &result.convergence_plot_data {
        println!("    Trial {:>2} | best so far: {:>10.4}", trial_num, best_so_far);
    }
    println!("{}", "=".repeat(70));

    Ok(())
}