foxhunt/ml/tests/hyperopt_rainbow_search_space_wave6_3_test.rs

//! Wave 6.3: Rainbow DQN Hyperopt Search Space Validation Tests
//!
//! Validates that all 5 Rainbow DQN components are properly exposed
//! in the hyperopt search space with correct dimensionality and bounds.

use ml::hyperopt::adapters::dqn::DQNParams;
use ml::hyperopt::traits::ParameterSpace;
use ml::MLError;

#[test]
fn test_search_space_dimensions() {
    // Validate 14D continuous space
    let continuous_bounds = DQNParams::continuous_bounds();
    assert_eq!(
        continuous_bounds.len(),
        14,
        "Expected 14 continuous parameters (11 base + 3 Rainbow)"
    );

    // Validate parameter names match bounds
    let param_names = DQNParams::param_names();
    assert_eq!(
        param_names.len(),
        14,
        "Parameter names should match continuous bounds count"
    );
}

#[test]
fn test_rainbow_param_bounds() {
    let bounds = DQNParams::continuous_bounds();

    // v_min bounds (index 11): -2000 to -500
    assert_eq!(bounds[11].0, -2000.0, "v_min lower bound should be -2000");
    assert_eq!(bounds[11].1, -500.0, "v_min upper bound should be -500");

    // v_max bounds (index 12): 500 to 2000
    assert_eq!(bounds[12].0, 500.0, "v_max lower bound should be 500");
    assert_eq!(bounds[12].1, 2000.0, "v_max upper bound should be 2000");

    // noisy_sigma_init bounds (index 13): ln(0.1) to ln(1.0) (log scale)
    let expected_min = 0.1_f64.ln();
    let expected_max = 1.0_f64.ln();
    assert!(
        (bounds[13].0 - expected_min).abs() < 1e-6,
        "noisy_sigma_init lower bound should be ln(0.1)"
    );
    assert!(
        (bounds[13].1 - expected_max).abs() < 1e-6,
        "noisy_sigma_init upper bound should be ln(1.0)"
    );
}

#[test]
fn test_param_names_include_rainbow() {
    let names = DQNParams::param_names();

    // Check Rainbow param names are present
    assert!(
        names.contains(&"v_min"),
        "Parameter names should include v_min"
    );
    assert!(
        names.contains(&"v_max"),
        "Parameter names should include v_max"
    );
    assert!(
        names.contains(&"noisy_sigma_init"),
        "Parameter names should include noisy_sigma_init"
    );

    // Check correct indices
    assert_eq!(names[11], "v_min", "v_min should be at index 11");
    assert_eq!(names[12], "v_max", "v_max should be at index 12");
    assert_eq!(
        names[13], "noisy_sigma_init",
        "noisy_sigma_init should be at index 13"
    );
}

#[test]
fn test_from_continuous_14d() -> Result<(), MLError> {
    // Create 14D vector with valid values
    let x = vec![
        (-4.605170_f64), // 0: ln(0.01) = learning_rate
        128.0,           // 1: batch_size
        0.97,            // 2: gamma
        11.512925,       // 3: ln(100000) = buffer_size
        2.0,             // 4: hold_penalty_weight
        5.0,             // 5: max_position_absolute
        0.0,             // 6: ln(1.0) = huber_delta
        0.05,            // 7: entropy_coefficient
        1.0,             // 8: transaction_cost_multiplier
        0.6,             // 9: per_alpha
        0.4,             // 10: per_beta_start
        -1000.0,         // 11: v_min
        1000.0,          // 12: v_max
        (-0.693147),     // 13: ln(0.5) = noisy_sigma_init
    ];

    let params = DQNParams::from_continuous(&x)?;

    // Validate Rainbow params were parsed correctly
    assert_eq!(params.v_min, -1000.0);
    assert_eq!(params.v_max, 1000.0);
    assert!((params.noisy_sigma_init - 0.5).abs() < 0.01);

    // Validate other params still work
    assert!((params.learning_rate - 0.01).abs() < 1e-6);
    assert_eq!(params.batch_size, 128);
    assert_eq!(params.gamma, 0.97);

    Ok(())
}

#[test]
fn test_to_continuous_14d() {
    let params = DQNParams {
        learning_rate: 1e-4,
        batch_size: 128,
        gamma: 0.99,
        buffer_size: 100_000,
        hold_penalty_weight: 2.0,
        max_position_absolute: 2.0,
        huber_delta: 1.0,
        entropy_coefficient: 0.01,
        transaction_cost_multiplier: 1.0,
        use_per: true,
        per_alpha: 0.6,
        per_beta_start: 0.4,
        use_dueling: false,
        dueling_hidden_dim: 128,
        n_steps: 1,
        tau: 0.001,
        use_distributional: false,
        num_atoms: 51,
        v_min: -1000.0,
        v_max: 1000.0,
        use_noisy_nets: false,
        noisy_sigma_init: 0.5,
    };

    let x = params.to_continuous();

    // Validate 14D output
    assert_eq!(x.len(), 14, "to_continuous should return 14 values");

    // Validate Rainbow params are present
    assert_eq!(x[11], -1000.0, "v_min should be at index 11");
    assert_eq!(x[12], 1000.0, "v_max should be at index 12");
    assert!(
        (x[13] - 0.5_f64.ln()).abs() < 1e-6,
        "noisy_sigma_init should be ln(0.5) at index 13"
    );
}

#[test]
fn test_rainbow_params_roundtrip() -> Result<(), MLError> {
    let original = DQNParams {
        learning_rate: 1e-4,
        batch_size: 128,
        gamma: 0.99,
        buffer_size: 100_000,
        hold_penalty_weight: 2.0,
        max_position_absolute: 2.0,
        huber_delta: 1.0,
        entropy_coefficient: 0.01,
        transaction_cost_multiplier: 1.0,
        use_per: true,
        per_alpha: 0.6,
        per_beta_start: 0.4,
        use_dueling: false,
        dueling_hidden_dim: 128,
        n_steps: 1,
        tau: 0.001,
        use_distributional: false,
        num_atoms: 51,
        v_min: -1500.0,
        v_max: 1500.0,
        use_noisy_nets: false,
        noisy_sigma_init: 0.7,
    };

    let continuous = original.to_continuous();
    let recovered = DQNParams::from_continuous(&continuous)?;

    // Validate Rainbow params survived roundtrip
    assert!((recovered.v_min - original.v_min).abs() < 1e-6);
    assert!((recovered.v_max - original.v_max).abs() < 1e-6);
    assert!((recovered.noisy_sigma_init - original.noisy_sigma_init).abs() < 1e-3);

    // Validate base params still work
    assert!((recovered.learning_rate - original.learning_rate).abs() < 1e-6);
    assert_eq!(recovered.batch_size, original.batch_size);
    assert!((recovered.gamma - original.gamma).abs() < 1e-6);

    Ok(())
}

#[test]
fn test_default_rainbow_params() {
    let default = DQNParams::default();

    // Validate Rainbow params have correct defaults
    assert_eq!(
        default.use_distributional, false,
        "use_distributional should default to false"
    );
    assert_eq!(default.num_atoms, 51, "num_atoms should default to 51");
    assert_eq!(default.v_min, -1000.0, "v_min should default to -1000.0");
    assert_eq!(default.v_max, 1000.0, "v_max should default to 1000.0");
    assert_eq!(
        default.use_noisy_nets, false,
        "use_noisy_nets should default to false"
    );
    assert_eq!(
        default.noisy_sigma_init, 0.5,
        "noisy_sigma_init should default to 0.5"
    );
    assert_eq!(default.use_dueling, false, "use_dueling should default to false");
    assert_eq!(
        default.dueling_hidden_dim, 128,
        "dueling_hidden_dim should default to 128"
    );
    assert_eq!(default.n_steps, 1, "n_steps should default to 1");
    assert_eq!(default.tau, 0.001, "tau should default to 0.001");
}

#[test]
fn test_continuous_bounds_count() {
    let bounds = DQNParams::continuous_bounds();

    // Wave 6.3: 14D continuous space (11 base + 3 Rainbow)
    assert_eq!(
        bounds.len(),
        14,
        "Should have 14 continuous parameters in Wave 6.3"
    );
}

#[test]
fn test_all_rainbow_components_present() {
    // This test ensures all 5 Rainbow components have at least some representation
    let params = DQNParams::default();

    // Component 1: Double DQN (already in production, not in params)
    // Component 2: Prioritized Experience Replay (PER)
    assert!(params.use_per, "PER should be enabled by default");
    assert!(params.per_alpha > 0.0 && params.per_alpha <= 1.0);
    assert!(params.per_beta_start >= 0.0 && params.per_beta_start <= 1.0);

    // Component 3: Dueling Networks
    assert!(params.dueling_hidden_dim > 0);

    // Component 4: Multi-Step Learning
    assert!(params.n_steps >= 1);
    assert!(params.tau > 0.0 && params.tau < 1.0);

    // Component 5: Distributional RL (C51)
    assert!(params.num_atoms > 0);
    assert!(params.v_min < params.v_max);

    // Component 6: Noisy Networks
    assert!(params.noisy_sigma_init > 0.0);
}

#[test]
fn test_wave6_3_expansion() {
    // Validate Wave 6.3 expanded the search space correctly
    let bounds = DQNParams::continuous_bounds();
    let names = DQNParams::param_names();

    // Before Wave 6.3: 11D continuous
    // After Wave 6.3: 14D continuous (added v_min, v_max, noisy_sigma_init)
    assert_eq!(
        bounds.len(),
        14,
        "Wave 6.3 should expand to 14D continuous space"
    );
    assert_eq!(
        names.len(),
        14,
        "Wave 6.3 should have 14 parameter names"
    );

    // Validate the 3 new Rainbow params are last
    assert_eq!(names[11], "v_min");
    assert_eq!(names[12], "v_max");
    assert_eq!(names[13], "noisy_sigma_init");
}