foxhunt/ml/tests/dqn_evaluation_shape_test.rs

//! Evaluation Shape Mismatch Bug Test (Wave 10-A1)
//!
//! This test reproduces the shape mismatch bug that occurs during DQN evaluation:
//! "unexpected rank, expected: 0, got: 1 ([1])"
//!
//! The bug manifests after training completes (16,635 steps) during backtest evaluation
//! when the hyperopt adapter runs actions on validation data.
//!
//! Expected behavior:
//! - Tensor operations should return rank-0 scalars for temperature/division
//! - Evaluation should complete without shape errors
//!
//! Current behavior:
//! - FAILS with "unexpected rank, expected: 0, got: 1 ([1])" during action selection
//! - Likely caused by batch dimension [1] not being squeezed before scalar operations

use ml::dqn::dqn::{WorkingDQN, WorkingDQNConfig};

use ml::dqn::dqn::RewardSystem;

/// Helper to create minimal DQN config for testing
fn create_minimal_config() -> WorkingDQNConfig {
    WorkingDQNConfig {
        state_dim: 128,
        num_actions: 45,            // Factored action space
        hidden_dims: vec![128, 64], // Small network for speed
        learning_rate: 0.001,
        gamma: 0.99,
        epsilon_start: 0.0, // Pure greedy for reproducibility
        epsilon_end: 0.0,
        epsilon_decay: 1.0,
        replay_buffer_capacity: 1000,
        batch_size: 32,
        min_replay_size: 32,
        target_update_freq: 100,
        use_double_dqn: true,
        use_huber_loss: false,
        huber_delta: 1.0,
        leaky_relu_alpha: 0.01,
        gradient_clip_norm: 10.0,
        td_error_clip: 10.0,
        tau: 0.001,
        use_soft_updates: false,
        warmup_steps: 0, // No warmup for test
        temperature_start: 0.1,
        temperature_min: 0.01,
        temperature_decay: 0.995,
        target_temperature_fraction: 0.75,
        variance_multiplier: 0.0, // Disable variance adaptation for simplicity
        use_adaptive_temperature: false,
        loss_improvement_threshold: 0.999,
        plateau_window: 10,
        temp_increase_factor: 1.05,
        temperature_slow_decay: 0.998,
        reward_system: RewardSystem::SimplePnL,
        reward_scale: 1.0,
    }
}

#[test]
fn test_evaluation_action_selection_shape_correctness() {
    // Setup: Create minimal DQN with greedy policy (epsilon=0)
    let config = create_minimal_config();
    let mut agent = WorkingDQN::new(config).expect("Failed to create DQN agent");

    // Simulate validation data: single state vector [128 features]
    let state: Vec<f32> = (0..128).map(|i| (i as f32) * 0.01).collect();

    // Test: Select action (this should NOT fail with shape error after Wave 10 fix)
    // Expected: Returns action index 0-44 without error
    // Fixed: No longer fails with "unexpected rank, expected: 0, got: 1"
    let result = agent.select_action(&state);

    match result {
        Ok(action_idx) => {
            // Verify action is in valid range
            assert!(
                action_idx < 45,
                "Action index {} exceeds 45-action space",
                action_idx
            );
            println!(
                "✓ TEST PASSED: Action selection succeeded with action {} (Wave 10 bug fixed)",
                action_idx
            );
        },
        Err(e) => {
            // This should NOT happen after the fix
            panic!(
                "TEST FAILED: Action selection failed with error (bug not fixed): {}",
                e
            );
        },
    }
}

#[test]
fn test_batch_size_one_tensor_shape() {
    // Test that demonstrates the root cause: batch size 1 creates rank-1 tensors
    use candle_core::{Device, Tensor};

    let device = Device::Cpu;

    // Simulate state tensor with batch dimension [1, 128]
    let state_vec: Vec<f32> = (0..128).map(|i| (i as f32) * 0.01).collect();
    let state_tensor = Tensor::from_vec(state_vec.clone(), (1, 128), &device)
        .expect("Failed to create state tensor");

    println!("State tensor shape: {:?}", state_tensor.dims());
    assert_eq!(state_tensor.dims(), &[1, 128]);

    // Simulate temperature scalar (this is where the bug manifests)
    // When we try to divide by temperature, Candle expects rank-0 scalar
    let temperature = 0.1_f32;

    // Attempt 1: Create temperature as rank-0 scalar (correct)
    let temp_scalar = Tensor::new(&[temperature], &device).expect("Failed to create scalar");
    println!("Temperature scalar shape: {:?}", temp_scalar.dims());

    // Attempt 2: What happens if temperature is accidentally [1] instead of []?
    let temp_rank1 =
        Tensor::new(vec![temperature], &device).expect("Failed to create rank-1 tensor");
    println!("Temperature rank-1 shape: {:?}", temp_rank1.dims());

    // The bug: division by rank-1 [1] instead of rank-0 [] causes shape error
    // This is likely happening in select_action() line 715: `q_values / adaptive_temp`
}

#[test]
fn test_softmax_batch_dimension_handling() {
    // Test softmax on batched Q-values to verify dimension handling
    use candle_core::{Device, Tensor};
    use candle_nn::ops::softmax;

    let device = Device::Cpu;

    // Simulate Q-values with batch dimension [1, 45]
    let q_values: Vec<f32> = (0..45).map(|i| i as f32 * 0.1).collect();
    let q_tensor =
        Tensor::from_vec(q_values, (1, 45), &device).expect("Failed to create Q-values tensor");

    println!("Q-values shape: {:?}", q_tensor.dims());

    // Apply softmax along action dimension (dim=1)
    let probs = softmax(&q_tensor, 1).expect("Softmax failed");
    println!("Softmax output shape: {:?}", probs.dims());

    // Extract probabilities - should work without shape error
    let probs_vec = probs
        .flatten_all()
        .expect("Flatten failed")
        .to_vec1::<f32>()
        .expect("to_vec1 failed");

    println!("Extracted {} probabilities", probs_vec.len());
    assert_eq!(probs_vec.len(), 45);

    // Verify probabilities sum to 1.0
    let sum: f32 = probs_vec.iter().sum();
    assert!(
        (sum - 1.0).abs() < 0.001,
        "Probabilities should sum to 1.0, got {}",
        sum
    );
}

#[test]
fn test_temperature_division_shape() {
    // Root cause test: What happens when dividing batched tensor by scalar?
    use candle_core::{Device, Tensor};

    let device = Device::Cpu;

    // Q-values with batch dimension [1, 45]
    let q_values: Vec<f32> = (0..45).map(|i| i as f32).collect();
    let q_tensor = Tensor::from_vec(q_values, (1, 45), &device).expect("Failed to create Q-values");

    // Temperature as f64 (Rust scalar)
    let temperature = 0.1_f64;

    // Attempt division: q_values / temperature
    // This is what happens at line 715 in dqn.rs: `let logits = (q_values / adaptive_temp)?;`
    let result = (q_tensor / temperature);

    match result {
        Ok(logits) => {
            println!("Division succeeded, logits shape: {:?}", logits.dims());
            assert_eq!(
                logits.dims(),
                &[1, 45],
                "Logits should preserve batch dimension"
            );
        },
        Err(e) => {
            println!("Division failed with error: {}", e);
            panic!("Temperature division should not fail for batched tensors");
        },
    }
}