foxhunt/ml/tests/dqn_hyperopt_checkpoint_test.rs

//! Test for DQN hyperopt checkpoint saving
//!
//! This test verifies that the DQN hyperopt adapter saves model checkpoints
//! after each trial completes.

use ml::hyperopt::adapters::dqn::{DQNParams, DQNTrainer};
use ml::hyperopt::paths::TrainingPaths;
use ml::hyperopt::traits::HyperparameterOptimizable;
use std::path::PathBuf;

#[test]
fn test_dqn_hyperopt_saves_checkpoint() -> anyhow::Result<()> {
    // Create temporary directory for test
    let temp_dir = tempfile::tempdir()?;
    let base_dir = temp_dir.path();

    // Create training paths
    let training_paths = TrainingPaths::new(base_dir.to_str().unwrap(), "dqn", "test_run_001");

    // Create test data directory with parquet file
    let data_dir = base_dir.join("test_data");
    std::fs::create_dir_all(&data_dir)?;

    // Copy test parquet file
    let test_parquet = PathBuf::from("test_data/ES_FUT_180d.parquet");
    if test_parquet.exists() {
        std::fs::copy(&test_parquet, data_dir.join("ES_FUT_180d.parquet"))?;
    } else {
        eprintln!("⚠️ Test parquet file not found, skipping test");
        return Ok(());
    }

    // Create DQN trainer with minimal epochs for speed
    let mut trainer = DQNTrainer::new(&data_dir, 2)? // Only 2 epochs for fast test
        .with_training_paths(training_paths.clone())
        .with_early_stopping(5, 1); // Allow early stopping after 1 epoch

    // Train with test parameters
    let params = DQNParams {
        learning_rate: 1e-4,
        batch_size: 64,
        gamma: 0.99,
        epsilon_decay: 0.995,
        buffer_size: 10000,
    };

    let metrics = trainer.train_with_params(params)?;

    // CRITICAL TEST: Verify checkpoint file was created
    let checkpoint_path = training_paths
        .checkpoints_dir()
        .join("trial_000_model.safetensors");

    assert!(
        checkpoint_path.exists(),
        "❌ FAILED: Checkpoint file not found at {:?}",
        checkpoint_path
    );

    // Verify checkpoint is not empty
    let metadata = std::fs::metadata(&checkpoint_path)?;
    assert!(
        metadata.len() > 1000,
        "❌ FAILED: Checkpoint file is too small ({} bytes), likely empty",
        metadata.len()
    );

    println!("✅ PASS: Checkpoint saved to {:?}", checkpoint_path);
    println!("✅ PASS: Checkpoint size: {} bytes", metadata.len());
    println!(
        "✅ PASS: Training metrics: train_loss={:.6}, val_loss={:.6}",
        metrics.train_loss, metrics.val_loss
    );

    Ok(())
}

#[test]
fn test_checkpoint_contains_model_weights() -> anyhow::Result<()> {
    // Create temporary directory for test
    let temp_dir = tempfile::tempdir()?;
    let base_dir = temp_dir.path();

    // Create training paths
    let training_paths = TrainingPaths::new(base_dir.to_str().unwrap(), "dqn", "test_run_002");

    // Create test data directory with parquet file
    let data_dir = base_dir.join("test_data");
    std::fs::create_dir_all(&data_dir)?;

    // Copy test parquet file
    let test_parquet = PathBuf::from("test_data/ES_FUT_180d.parquet");
    if test_parquet.exists() {
        std::fs::copy(&test_parquet, data_dir.join("ES_FUT_180d.parquet"))?;
    } else {
        eprintln!("⚠️ Test parquet file not found, skipping test");
        return Ok(());
    }

    // Create DQN trainer
    let mut trainer = DQNTrainer::new(&data_dir, 2)?
        .with_training_paths(training_paths.clone())
        .with_early_stopping(5, 1);

    // Train with test parameters
    let params = DQNParams {
        learning_rate: 1e-4,
        batch_size: 64,
        gamma: 0.99,
        epsilon_decay: 0.995,
        buffer_size: 10000,
    };

    trainer.train_with_params(params)?;

    // Load checkpoint and verify it contains model weights
    let checkpoint_path = training_paths
        .checkpoints_dir()
        .join("trial_000_model.safetensors");

    let device = candle_core::Device::Cpu;
    let tensors = candle_core::safetensors::load(&checkpoint_path, &device)?;

    // Verify checkpoint contains expected layer weights
    assert!(
        !tensors.is_empty(),
        "❌ FAILED: Checkpoint contains no tensors"
    );

    // Check for expected layer names (layer_0, layer_1, output)
    let tensor_names: Vec<_> = tensors.keys().collect();
    println!("✅ PASS: Checkpoint contains {} tensors", tensors.len());
    println!("✅ PASS: Tensor names: {:?}", tensor_names);

    // Verify tensors have reasonable shapes
    for (name, tensor) in tensors.iter() {
        let shape = tensor.shape();
        assert!(
            shape.dims().iter().all(|&d| d > 0),
            "❌ FAILED: Tensor {} has invalid shape: {:?}",
            name,
            shape
        );
    }

    println!("✅ PASS: All tensors have valid shapes");

    Ok(())
}