foxhunt/ml/tests/rainbow_network_architecture_validation.rs

//! Rainbow DQN Network Architecture Validation Tests
//!
//! Validates that the implementation matches the paper specification:
//! "Rainbow: Combining Improvements in Deep Reinforcement Learning" (Hessel et al., 2017)
//!
//! Key components verified:
//! 1. Noisy Linear layers (not standard Linear)
//! 2. Dueling architecture (value + advantage streams)
//! 3. C51 distributional output (num_actions × num_atoms)
//! 4. Forward pass shapes
//! 5. Softmax over atoms dimension

use anyhow::Result;
use candle_core::{DType, Device, Tensor};
use candle_nn::{Module, VarBuilder, VarMap};
use ml::dqn::{
    CategoricalDistribution, DistributionalConfig, RainbowNetwork, RainbowNetworkConfig,
};
use ml::MLError;

/// Test 1: Network initialization with correct parameter counts
#[test]
fn test_rainbow_network_initialization() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let config = RainbowNetworkConfig {
        input_size: 128,
        hidden_sizes: vec![512, 512],
        num_actions: 3,
        activation: ml::dqn::rainbow_network::ActivationType::ReLU,
        dropout_rate: 0.1,
        distributional: DistributionalConfig {
            num_atoms: 51,
            v_min: -10.0,
            v_max: 10.0,
        },
        use_noisy_layers: true,
        dueling: true,
    };

    let network = RainbowNetwork::new(&vs, config)?;

    // Verify network was created successfully
    assert_eq!(network.config().input_size, 128);
    assert_eq!(network.config().num_actions, 3);
    assert_eq!(network.config().distributional.num_atoms, 51);
    assert!(network.config().use_noisy_layers);
    assert!(network.config().dueling);

    Ok(())
}

/// Test 2: Forward pass shape validation for single sample
#[test]
fn test_forward_pass_single_sample() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let config = RainbowNetworkConfig {
        input_size: 128,
        hidden_sizes: vec![256, 256],
        num_actions: 3,
        activation: ml::dqn::rainbow_network::ActivationType::ReLU,
        dropout_rate: 0.0,
        distributional: DistributionalConfig {
            num_atoms: 51,
            v_min: -10.0,
            v_max: 10.0,
        },
        use_noisy_layers: true,
        dueling: true,
    };

    let network = RainbowNetwork::new(&vs, config)?;

    // Input: [batch=1, state_dim=128]
    let input = Tensor::randn(0.0_f32, 1.0_f32, (1, 128), &device)?;

    // Forward pass
    let output = network
        .forward(&input)
        .map_err(|e| MLError::ModelError(format!("Forward pass failed: {}", e)))?;

    // Output should be [batch=1, num_actions=3, num_atoms=51]
    assert_eq!(output.shape().dims(), &[1, 3, 51]);

    Ok(())
}

/// Test 3: Forward pass shape validation for batched input
#[test]
fn test_forward_pass_batch() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let config = RainbowNetworkConfig {
        input_size: 128,
        hidden_sizes: vec![512, 512],
        num_actions: 3,
        activation: ml::dqn::rainbow_network::ActivationType::ReLU,
        dropout_rate: 0.0,
        distributional: DistributionalConfig {
            num_atoms: 51,
            v_min: -10.0,
            v_max: 10.0,
        },
        use_noisy_layers: true,
        dueling: true,
    };

    let network = RainbowNetwork::new(&vs, config)?;

    // Input: [batch=32, state_dim=128]
    let batch_size = 32;
    let input = Tensor::randn(0.0_f32, 1.0_f32, (batch_size, 128), &device)?;

    // Forward pass
    let output = network
        .forward(&input)
        .map_err(|e| MLError::ModelError(format!("Forward pass failed: {}", e)))?;

    // Output should be [batch=32, num_actions=3, num_atoms=51]
    assert_eq!(output.shape().dims(), &[batch_size, 3, 51]);

    Ok(())
}

/// Test 4: C51 distributional output - verify softmax over atoms dimension
#[test]
fn test_c51_output_is_probability_distribution() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let config = RainbowNetworkConfig {
        input_size: 128,
        hidden_sizes: vec![256],
        num_actions: 3,
        activation: ml::dqn::rainbow_network::ActivationType::ReLU,
        dropout_rate: 0.0,
        distributional: DistributionalConfig {
            num_atoms: 51,
            v_min: -10.0,
            v_max: 10.0,
        },
        use_noisy_layers: true,
        dueling: true,
    };

    let network = RainbowNetwork::new(&vs, config)?;

    let input = Tensor::randn(0.0_f32, 1.0_f32, (4, 128), &device)?;

    // Forward pass
    let output = network
        .forward(&input)
        .map_err(|e| MLError::ModelError(format!("Forward pass failed: {}", e)))?;

    // Output shape: [batch=4, actions=3, atoms=51]
    assert_eq!(output.shape().dims(), &[4, 3, 51]);

    // For each (batch, action) pair, the distribution over atoms should sum to 1.0
    for batch_idx in 0..4 {
        for action_idx in 0..3 {
            let dist = output
                .get(batch_idx)
                .map_err(|e| MLError::ModelError(format!("Failed to get batch: {}", e)))?
                .get(action_idx)
                .map_err(|e| MLError::ModelError(format!("Failed to get action: {}", e)))?;

            let sum_tensor = dist
                .sum_all()
                .map_err(|e| MLError::ModelError(format!("Failed to sum: {}", e)))?;

            // Handle both [] and [1] shapes
            let sum: f32 = if sum_tensor.rank() == 0 {
                sum_tensor.to_scalar().map_err(|e| {
                    MLError::ModelError(format!("Failed to convert to scalar: {}", e))
                })?
            } else {
                sum_tensor
                    .squeeze(0)
                    .map_err(|e| MLError::ModelError(format!("Failed to squeeze: {}", e)))?
                    .to_scalar()
                    .map_err(|e| {
                        MLError::ModelError(format!("Failed to convert to scalar: {}", e))
                    })?
            };

            // Check sum is approximately 1.0 (allow small numerical error)
            assert!(
                (sum - 1.0).abs() < 1e-4,
                "Distribution sum should be 1.0, got {}",
                sum
            );

            // Check all probabilities are non-negative
            let min_tensor = dist
                .min_keepdim(0)
                .map_err(|e| MLError::ModelError(format!("Failed to get min: {}", e)))?;

            // Handle both [] and [1] shapes
            let min_val: f32 = if min_tensor.rank() == 0 {
                min_tensor.to_scalar().map_err(|e| {
                    MLError::ModelError(format!("Failed to convert to scalar: {}", e))
                })?
            } else {
                min_tensor
                    .squeeze(0)
                    .map_err(|e| MLError::ModelError(format!("Failed to squeeze: {}", e)))?
                    .to_scalar()
                    .map_err(|e| {
                        MLError::ModelError(format!("Failed to convert to scalar: {}", e))
                    })?
            };

            assert!(
                min_val >= -1e-6,
                "All probabilities should be non-negative, got min={}",
                min_val
            );
        }
    }

    Ok(())
}

/// Test 5: Dueling architecture - verify value and advantage streams are used
#[test]
fn test_dueling_architecture() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create two networks: one with dueling, one without
    let config_dueling = RainbowNetworkConfig {
        input_size: 64,
        hidden_sizes: vec![128],
        num_actions: 3,
        activation: ml::dqn::rainbow_network::ActivationType::ReLU,
        dropout_rate: 0.0,
        distributional: DistributionalConfig {
            num_atoms: 51,
            v_min: -10.0,
            v_max: 10.0,
        },
        use_noisy_layers: false, // Disable noise for deterministic test
        dueling: true,
    };

    let config_no_dueling = RainbowNetworkConfig {
        dueling: false,
        ..config_dueling.clone()
    };

    let network_dueling = RainbowNetwork::new(&vs.pp("dueling"), config_dueling)?;
    let network_standard = RainbowNetwork::new(&vs.pp("standard"), config_no_dueling)?;

    let input = Tensor::randn(0.0_f32, 1.0_f32, (2, 64), &device)?;

    // Both should produce valid outputs
    let output_dueling = network_dueling
        .forward(&input)
        .map_err(|e| MLError::ModelError(format!("Dueling forward failed: {}", e)))?;

    let output_standard = network_standard
        .forward(&input)
        .map_err(|e| MLError::ModelError(format!("Standard forward failed: {}", e)))?;

    // Both should have correct shape [batch=2, actions=3, atoms=51]
    assert_eq!(output_dueling.shape().dims(), &[2, 3, 51]);
    assert_eq!(output_standard.shape().dims(), &[2, 3, 51]);

    // Outputs should be different (dueling combines value+advantage, standard doesn't)
    let diff = output_dueling
        .sub(&output_standard)
        .map_err(|e| MLError::ModelError(format!("Failed to compute diff: {}", e)))?;

    let diff_norm: f32 = diff
        .sqr()
        .map_err(|e| MLError::ModelError(format!("Failed to square: {}", e)))?
        .sum_all()
        .map_err(|e| MLError::ModelError(format!("Failed to sum: {}", e)))?
        .to_scalar()
        .map_err(|e| MLError::ModelError(format!("Failed to convert to scalar: {}", e)))?;

    // Should be significantly different
    assert!(
        diff_norm > 1e-3,
        "Dueling and standard networks should produce different outputs"
    );

    Ok(())
}

/// Test 6: Noisy layers - verify noise sampling changes outputs
#[test]
fn test_noisy_layers_exploration() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let config = RainbowNetworkConfig {
        input_size: 64,
        hidden_sizes: vec![128],
        num_actions: 3,
        activation: ml::dqn::rainbow_network::ActivationType::ReLU,
        dropout_rate: 0.0,
        distributional: DistributionalConfig {
            num_atoms: 51,
            v_min: -10.0,
            v_max: 10.0,
        },
        use_noisy_layers: true,
        dueling: true,
    };

    let network = RainbowNetwork::new(&vs, config)?;

    let input = Tensor::randn(0.0_f32, 1.0_f32, (1, 64), &device)?;

    // First forward pass
    let output1 = network
        .forward(&input)
        .map_err(|e| MLError::ModelError(format!("First forward failed: {}", e)))?;

    // Second forward pass (noise should be different if reset)
    let output2 = network
        .forward(&input)
        .map_err(|e| MLError::ModelError(format!("Second forward failed: {}", e)))?;

    // Both should have correct shape
    assert_eq!(output1.shape().dims(), &[1, 3, 51]);
    assert_eq!(output2.shape().dims(), &[1, 3, 51]);

    // Note: In this test, noise is NOT reset between forward passes,
    // so outputs might be the same. This test verifies that the network
    // CAN produce outputs (noise sampling doesn't crash).
    // A proper test would require access to noise reset functionality.

    Ok(())
}

/// Test 7: Q-value extraction from distributions
#[test]
fn test_q_value_extraction() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let config = RainbowNetworkConfig {
        input_size: 64,
        hidden_sizes: vec![128],
        num_actions: 3,
        activation: ml::dqn::rainbow_network::ActivationType::ReLU,
        dropout_rate: 0.0,
        distributional: DistributionalConfig {
            num_atoms: 51,
            v_min: -10.0,
            v_max: 10.0,
        },
        use_noisy_layers: false,
        dueling: true,
    };

    let network = RainbowNetwork::new(&vs, config)?;

    let input = Tensor::randn(0.0_f32, 1.0_f32, (2, 64), &device)?;

    // Get distributions
    let distributions = network
        .forward(&input)
        .map_err(|e| MLError::ModelError(format!("Forward failed: {}", e)))?;

    // Extract Q-values
    let q_values = network
        .get_q_values(&distributions)
        .map_err(|e| MLError::ModelError(format!("Q-value extraction failed: {}", e)))?;

    // Q-values should have shape [batch=2, actions=3]
    assert_eq!(q_values.shape().dims(), &[2, 3]);

    // Q-values should be within reasonable range (-10.0 to 10.0 given v_min/v_max)
    let q_min_tensor = q_values
        .min_keepdim(1)
        .map_err(|e| MLError::ModelError(format!("Failed to get min: {}", e)))?
        .min_keepdim(0)
        .map_err(|e| MLError::ModelError(format!("Failed to get min: {}", e)))?;

    // Handle both [] and [1, 1] shapes
    let q_min: f32 = if q_min_tensor.rank() == 0 {
        q_min_tensor
            .to_scalar()
            .map_err(|e| MLError::ModelError(format!("Failed to convert to scalar: {}", e)))?
    } else {
        q_min_tensor
            .flatten_all()
            .map_err(|e| MLError::ModelError(format!("Failed to flatten: {}", e)))?
            .get(0)
            .map_err(|e| MLError::ModelError(format!("Failed to get index: {}", e)))?
            .to_scalar()
            .map_err(|e| MLError::ModelError(format!("Failed to convert to scalar: {}", e)))?
    };

    let q_max_tensor = q_values
        .max_keepdim(1)
        .map_err(|e| MLError::ModelError(format!("Failed to get max: {}", e)))?
        .max_keepdim(0)
        .map_err(|e| MLError::ModelError(format!("Failed to get max: {}", e)))?;

    // Handle both [] and [1, 1] shapes
    let q_max: f32 = if q_max_tensor.rank() == 0 {
        q_max_tensor
            .to_scalar()
            .map_err(|e| MLError::ModelError(format!("Failed to convert to scalar: {}", e)))?
    } else {
        q_max_tensor
            .flatten_all()
            .map_err(|e| MLError::ModelError(format!("Failed to flatten: {}", e)))?
            .get(0)
            .map_err(|e| MLError::ModelError(format!("Failed to get index: {}", e)))?
            .to_scalar()
            .map_err(|e| MLError::ModelError(format!("Failed to convert to scalar: {}", e)))?
    };

    assert!(
        q_min >= -10.0 && q_max <= 10.0,
        "Q-values should be within v_min/v_max range, got [{}, {}]",
        q_min,
        q_max
    );

    Ok(())
}

/// Test 8: Different activation functions
#[test]
fn test_activation_functions() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let activations = vec![
        ml::dqn::rainbow_network::ActivationType::ReLU,
        ml::dqn::rainbow_network::ActivationType::LeakyReLU,
        ml::dqn::rainbow_network::ActivationType::Swish,
        ml::dqn::rainbow_network::ActivationType::ELU,
    ];

    for (idx, activation) in activations.iter().enumerate() {
        let config = RainbowNetworkConfig {
            input_size: 64,
            hidden_sizes: vec![128],
            num_actions: 3,
            activation: *activation,
            dropout_rate: 0.0,
            distributional: DistributionalConfig {
                num_atoms: 51,
                v_min: -10.0,
                v_max: 10.0,
            },
            use_noisy_layers: false,
            dueling: true,
        };

        let network = RainbowNetwork::new(&vs.pp(&format!("net_{}", idx)), config)?;

        let input = Tensor::randn(0.0_f32, 1.0_f32, (2, 64), &device)?;

        let output = network.forward(&input).map_err(|e| {
            MLError::ModelError(format!("Forward failed for {:?}: {}", activation, e))
        })?;

        // Should produce correct shape regardless of activation
        assert_eq!(output.shape().dims(), &[2, 3, 51]);
    }

    Ok(())
}

/// Test 9: Network with different hidden layer configurations
#[test]
fn test_different_hidden_layer_configs() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let hidden_configs = vec![
        vec![128],
        vec![256, 256],
        vec![512, 512, 256],
        vec![64, 128, 256, 128],
    ];

    for (idx, hidden_sizes) in hidden_configs.iter().enumerate() {
        let config = RainbowNetworkConfig {
            input_size: 64,
            hidden_sizes: hidden_sizes.clone(),
            num_actions: 3,
            activation: ml::dqn::rainbow_network::ActivationType::ReLU,
            dropout_rate: 0.0,
            distributional: DistributionalConfig {
                num_atoms: 51,
                v_min: -10.0,
                v_max: 10.0,
            },
            use_noisy_layers: false,
            dueling: true,
        };

        let network = RainbowNetwork::new(&vs.pp(&format!("config_{}", idx)), config)?;

        let input = Tensor::randn(0.0_f32, 1.0_f32, (2, 64), &device)?;

        let output = network.forward(&input).map_err(|e| {
            MLError::ModelError(format!(
                "Forward failed for hidden config {:?}: {}",
                hidden_sizes, e
            ))
        })?;

        // Should produce correct shape regardless of hidden layer configuration
        assert_eq!(output.shape().dims(), &[2, 3, 51]);
    }

    Ok(())
}

/// Test 10: Categorical distribution support values
#[test]
fn test_categorical_distribution_support() -> Result<(), MLError> {
    let config = DistributionalConfig {
        num_atoms: 51,
        v_min: -10.0,
        v_max: 10.0,
    };

    let dist = CategoricalDistribution::new(&config)?;

    // Verify support tensor has correct size
    assert_eq!(dist.num_atoms(), 51);

    let support = dist.support();
    assert_eq!(support.shape().dims(), &[51]);

    // Check first and last values
    let first: f32 = support
        .get(0)
        .map_err(|e| MLError::ModelError(format!("Failed to get first: {}", e)))?
        .to_scalar()
        .map_err(|e| MLError::ModelError(format!("Failed to convert to scalar: {}", e)))?;

    let last: f32 = support
        .get(50)
        .map_err(|e| MLError::ModelError(format!("Failed to get last: {}", e)))?
        .to_scalar()
        .map_err(|e| MLError::ModelError(format!("Failed to convert to scalar: {}", e)))?;

    assert!(
        (first - (-10.0_f32)).abs() < 1e-5,
        "First support value should be v_min=-10.0, got {}",
        first
    );
    assert!(
        (last - 10.0_f32).abs() < 1e-5,
        "Last support value should be v_max=10.0, got {}",
        last
    );

    // Verify atoms are evenly spaced
    let delta_z = (10.0 - (-10.0)) / 50.0; // (v_max - v_min) / (num_atoms - 1)
    for i in 0..50 {
        let expected = -10.0 + i as f32 * delta_z;
        let actual: f32 = support
            .get(i)
            .map_err(|e| MLError::ModelError(format!("Failed to get atom {}: {}", i, e)))?
            .to_scalar()
            .map_err(|e| MLError::ModelError(format!("Failed to convert to scalar: {}", e)))?;

        assert!(
            (actual - expected).abs() < 1e-4,
            "Atom {} should be {}, got {}",
            i,
            expected,
            actual
        );
    }

    Ok(())
}