foxhunt/ml/tests/dqn_c51_shape_validation_test.rs

//! Comprehensive C51 Shape Validation Test
//!
//! **OBJECTIVE**: Validate tensor shapes at EVERY step of the C51 (Distributional RL) pipeline
//! to identify exactly where shape mismatches occur during gradient computation.
//!
//! **INVESTIGATION**: Bug #36 - scatter_add gradient flow issue in distributional RL
//!
//! **ROOT CAUSE HYPOTHESIS**: Shape mismatches between network outputs and loss computation
//! cause gradients to fail during backpropagation through scatter_add operations.
//!
//! This test creates a minimal DQN agent with C51 and logs shapes at every critical point:
//! 1. Network forward pass (online network)
//! 2. Action selection and distribution gathering
//! 3. Target network forward pass
//! 4. Bellman operator application
//! 5. Categorical cross-entropy loss computation
//! 6. Gradient backward pass
//!
//! **EXPECTED SHAPES** (for reference):
//! - States: [batch, state_dim]
//! - Network output (distributions): [batch, num_actions, num_atoms]
//! - Actions: [batch] (integer indices)
//! - Selected distributions: [batch, num_atoms]
//! - Rewards: [batch]
//! - Dones: [batch]
//! - Target distributions: [batch, num_atoms]
//! - Loss: scalar tensor

use candle_core::{Device, Tensor};
use candle_nn::VarMap;
use ml::dqn::{DistributionalDuelingConfig, DistributionalDuelingQNetwork};
use ml::dqn::distributional::{CategoricalDistribution, DistributionalConfig};
use ml::MLError;

/// Helper to print tensor shape with descriptive label
fn log_shape(label: &str, tensor: &Tensor) -> Result<(), MLError> {
    let dims = tensor.dims();
    let dtype = tensor.dtype();
    let device_str = match tensor.device() {
        Device::Cpu => "CPU",
        Device::Cuda(_) => "CUDA",
        _ => "Unknown",
    };
    println!(
        "[SHAPE] {:<40} {:?} (dtype={:?}, device={})",
        label, dims, dtype, device_str
    );
    Ok(())
}

#[test]
fn test_c51_shape_flow_validation() -> Result<(), MLError> {
    println!("\n=== C51 SHAPE VALIDATION TEST ===\n");
    println!("Testing distributional RL (C51) tensor shapes through entire pipeline:\n");

    // Configuration
    let batch_size = 32;
    let state_dim = 54; // Production state dimension
    let num_actions = 45; // Production action space (5×3×3)
    let num_atoms = 51; // C51 standard

    let device = Device::cuda_if_available(0)?;
    println!("Device: {:?}\n", device);

    // ============================================
    // STEP 1: Create Distributional Dueling Network
    // ============================================
    println!("STEP 1: Network Creation");
    println!("------------------------");

    let config = DistributionalDuelingConfig::new(
        state_dim,
        num_actions,
        num_atoms,
        vec![256, 128], // shared layers
        64,             // value_hidden_dim
        64,             // advantage_hidden_dim
    );

    let network = DistributionalDuelingQNetwork::new(config.clone(), device.clone())?;
    let target_network = DistributionalDuelingQNetwork::new(config, device.clone())?;

    println!("✓ Networks created (online + target)\n");

    // ============================================
    // STEP 2: Create Categorical Distribution
    // ============================================
    println!("STEP 2: Categorical Distribution");
    println!("---------------------------------");

    let dist_config = DistributionalConfig {
        num_atoms,
        v_min: -2.0,
        v_max: 2.0,
    };
    let categorical_dist = CategoricalDistribution::new(&dist_config, &device)?;

    log_shape("Support tensor", categorical_dist.support())?;
    println!("✓ Categorical distribution initialized\n");

    // ============================================
    // STEP 3: Forward Pass (Online Network)
    // ============================================
    println!("STEP 3: Forward Pass (Online Network)");
    println!("--------------------------------------");

    let states = Tensor::randn(0f32, 1.0, (batch_size, state_dim), &device)?;
    log_shape("Input states", &states)?;

    let distributions = network.forward(&states)?;
    log_shape("Network output (distributions)", &distributions)?;

    // Validate shape
    assert_eq!(
        distributions.dims(),
        &[batch_size, num_actions, num_atoms],
        "Network output shape mismatch"
    );
    println!("✓ Forward pass shape validated\n");

    // ============================================
    // STEP 4: Action Selection & Distribution Gathering
    // ============================================
    println!("STEP 4: Action Selection & Gathering");
    println!("-------------------------------------");

    // Simulate action selection (random actions)
    let actions_vec: Vec<u32> = (0..batch_size)
        .map(|_| rand::random::<u32>() % num_actions as u32)
        .collect();
    let actions = Tensor::from_vec(actions_vec.clone(), batch_size, &device)?;
    log_shape("Actions (indices)", &actions)?;

    // BUG #36 CRITICAL SECTION: Action gathering
    // This is where shape mismatches commonly occur!
    //
    // Goal: Extract distribution for each selected action
    // Input: distributions [batch, num_actions, num_atoms]
    // Actions: [batch] (integer indices)
    // Output: selected_distributions [batch, num_atoms]

    println!("\nAttempting action gathering (CRITICAL SECTION):");
    println!("  Method: gather + squeeze");
    println!("  Input: [batch={}, actions={}, atoms={}]", batch_size, num_actions, num_atoms);
    println!("  Actions: [batch={}]", batch_size);
    println!("  Expected output: [batch={}, atoms={}]", batch_size, num_atoms);

    // Reshape actions for gathering: [batch] → [batch, 1, 1]
    let actions_unsqueezed = actions.unsqueeze(1)?.unsqueeze(2)?;
    log_shape("Actions unsqueezed", &actions_unsqueezed)?;

    // Broadcast to [batch, 1, num_atoms]
    let actions_broadcast = actions_unsqueezed.broadcast_as((batch_size, 1, num_atoms))?.contiguous()?;
    log_shape("Actions broadcast", &actions_broadcast)?;

    // Gather distributions for selected actions
    let distributions_contiguous = distributions.contiguous()?;
    let selected_distributions = distributions_contiguous.gather(&actions_broadcast, 1)?;
    log_shape("Selected distributions (pre-squeeze)", &selected_distributions)?;

    let selected_distributions = selected_distributions.squeeze(1)?;
    log_shape("Selected distributions (post-squeeze)", &selected_distributions)?;

    // Validate shape
    assert_eq!(
        selected_distributions.dims(),
        &[batch_size, num_atoms],
        "Selected distributions shape mismatch"
    );
    println!("✓ Action gathering validated\n");

    // ============================================
    // STEP 5: Target Network Forward Pass
    // ============================================
    println!("STEP 5: Target Network Forward Pass");
    println!("------------------------------------");

    let next_states = Tensor::randn(0f32, 1.0, (batch_size, state_dim), &device)?;
    log_shape("Next states", &next_states)?;

    let next_distributions = target_network.forward(&next_states)?;
    log_shape("Next distributions", &next_distributions)?;

    // For simplicity, select best action based on expected Q-values
    // In production, this would use Double DQN (online network selects, target evaluates)
    let support = categorical_dist.support();
    let support_broadcast = support
        .unsqueeze(0)?
        .unsqueeze(0)?
        .broadcast_as((batch_size, num_actions, num_atoms))?;

    let q_values = (next_distributions.clone() * support_broadcast)?.sum(2)?;
    log_shape("Next Q-values", &q_values)?;

    let best_actions = q_values.argmax_keepdim(1)?;
    log_shape("Best actions", &best_actions)?;

    // Gather next distributions for best actions
    let best_actions_broadcast = best_actions
        .unsqueeze(2)?
        .broadcast_as((batch_size, 1, num_atoms))?.contiguous()?;
    log_shape("Best actions broadcast", &best_actions_broadcast)?;

    let next_distributions_contiguous = next_distributions.contiguous()?;
    let next_selected_distributions = next_distributions_contiguous
        .gather(&best_actions_broadcast, 1)?
        .squeeze(1)?;
    log_shape("Next selected distributions", &next_selected_distributions)?;

    // Validate shape
    assert_eq!(
        next_selected_distributions.dims(),
        &[batch_size, num_atoms],
        "Next selected distributions shape mismatch"
    );
    println!("✓ Target network validated\n");

    // ============================================
    // STEP 6: Bellman Operator Application
    // ============================================
    println!("STEP 6: Bellman Operator");
    println!("-------------------------");

    let rewards = Tensor::randn(0f32, 0.1, batch_size, &device)?;
    log_shape("Rewards", &rewards)?;

    let dones = Tensor::zeros(batch_size, candle_core::DType::F32, &device)?;
    log_shape("Dones", &dones)?;

    let gamma = 0.99f32;
    println!("Gamma: {}\n", gamma);

    let target_distributions = categorical_dist.apply_bellman_operator(
        &rewards,
        &next_selected_distributions,
        &dones,
        gamma,
    )?;
    log_shape("Target distributions (after Bellman)", &target_distributions)?;

    // Validate shape
    assert_eq!(
        target_distributions.dims(),
        &[batch_size, num_atoms],
        "Target distributions shape mismatch"
    );
    println!("✓ Bellman operator validated\n");

    // ============================================
    // STEP 7: Categorical Cross-Entropy Loss
    // ============================================
    println!("STEP 7: Categorical Cross-Entropy Loss");
    println!("---------------------------------------");

    let loss = categorical_dist.categorical_loss(
        &selected_distributions,
        &target_distributions,
    )?;
    log_shape("Loss", &loss)?;

    // Validate scalar
    let loss_dims: Vec<usize> = loss.dims().to_vec();
    assert_eq!(
        loss_dims,
        Vec::<usize>::new(),
        "Loss should be scalar"
    );

    let loss_value: f32 = loss.to_scalar()?;
    println!("Loss value: {:.6}", loss_value);
    assert!(loss_value.is_finite(), "Loss should be finite");
    println!("✓ Loss computation validated\n");

    // ============================================
    // STEP 8: Gradient Backward Pass
    // ============================================
    println!("STEP 8: Gradient Backward Pass");
    println!("-------------------------------");

    let grads = loss.backward()?;
    println!("✓ Backward pass completed");

    // Check if gradients exist
    let all_vars = network.vars().all_vars();
    let grad_count = all_vars.iter().filter(|var| grads.get(var).is_some()).count();
    println!("Gradients computed: {}/{}", grad_count, all_vars.len());

    // Sample one gradient to verify non-zero
    if let Some(grad) = grads.get(&all_vars[0]) {
        let grad_norm = grad.sqr()?.sum_all()?.to_scalar::<f32>()?.sqrt();
        println!("Sample gradient norm: {:.6}", grad_norm);
        assert!(grad_norm > 1e-6, "Gradient should be non-zero");
    }

    println!("✓ Gradients validated\n");

    // ============================================
    // FINAL SUMMARY
    // ============================================
    println!("\n=== SHAPE VALIDATION SUMMARY ===\n");
    println!("✅ ALL SHAPE CHECKS PASSED");
    println!();
    println!("Key Shapes:");
    println!("  - States: [batch={}, state_dim={}]", batch_size, state_dim);
    println!("  - Distributions: [batch={}, actions={}, atoms={}]", batch_size, num_actions, num_atoms);
    println!("  - Selected: [batch={}, atoms={}]", batch_size, num_atoms);
    println!("  - Loss: scalar");
    println!();
    println!("Critical Operations:");
    println!("  ✓ Forward pass (online network)");
    println!("  ✓ Action gathering (gather + squeeze)");
    println!("  ✓ Target network pass");
    println!("  ✓ Bellman operator (project_distribution)");
    println!("  ✓ Categorical loss");
    println!("  ✓ Backward pass (gradient flow)");
    println!();
    println!("Conclusion: C51 pipeline shapes are CORRECT");
    println!("Next: Investigate scatter_add gradient flow in project_distribution");

    Ok(())
}

#[test]
fn test_scatter_add_shape_compatibility() -> Result<(), MLError> {
    println!("\n=== SCATTER_ADD SHAPE COMPATIBILITY TEST ===\n");
    println!("Testing scatter_add operation with C51 distribution shapes:\n");

    let batch_size = 32;
    let num_atoms = 51;
    let device = Device::cuda_if_available(0)?;

    // Simulate project_distribution shapes
    println!("CONTEXT: project_distribution scatter operation");
    println!("Goal: Accumulate probabilities into projected distribution\n");

    // Base tensor: [batch, num_atoms]
    let base = Tensor::zeros((batch_size, num_atoms), candle_core::DType::F32, &device)?;
    log_shape("Base tensor (projected)", &base)?;

    // Indices: [batch, num_atoms] (target atom indices)
    // Simulate indices in valid range [0, num_atoms-1]
    let indices_vec: Vec<i64> = (0..batch_size * num_atoms)
        .map(|i| (i % num_atoms) as i64)
        .collect();
    let indices = Tensor::from_vec(indices_vec, (batch_size, num_atoms), &device)?;
    log_shape("Indices", &indices)?;

    // Source: [batch, num_atoms] (weights to scatter)
    let source = Tensor::randn(0f32, 1.0, (batch_size, num_atoms), &device)?;
    log_shape("Source (weights)", &source)?;

    // Scatter add operation
    println!("\nPerforming scatter_add(base, indices, source, dim=1)...");
    let result = base.scatter_add(&indices, &source, 1)?;
    log_shape("Result", &result)?;

    // Validate shape
    assert_eq!(
        result.dims(),
        &[batch_size, num_atoms],
        "Scatter result shape mismatch"
    );

    println!("\n✓ scatter_add operation successful");
    println!("Shape compatibility: CONFIRMED\n");

    Ok(())
}

#[test]
fn test_existing_failing_test_context() -> Result<(), MLError> {
    println!("\n=== EXISTING FAILING TEST CONTEXT ===\n");
    println!("Analyzing scatter_add_gradient_test.rs:\n");

    let device = Device::cuda_if_available(0)?;

    // Original test uses [2, 5] base and [2, 3] indices
    println!("Original test dimensions:");
    println!("  Base: [2, 5]");
    println!("  Indices: [2, 3]");
    println!("  Source: [2, 3]");
    println!("  Scatter dimension: 1");
    println!();

    println!("Question: Is this representative of production usage?");
    println!();
    println!("Production C51 usage:");
    println!("  Base: [batch=32, atoms=51]");
    println!("  Indices: [batch=32, atoms=51]");
    println!("  Source: [batch=32, atoms=51]");
    println!("  Scatter dimension: 1");
    println!();

    println!("ANALYSIS:");
    println!("  - Original test: Source (3 elements) → Base (5 slots)");
    println!("  - Production: Source (51 elements) → Base (51 slots)");
    println!("  - Difference: Production has SAME source/target size");
    println!("  - Implication: Original test IS representative");
    println!("  - Reason: scatter_add accumulates, size difference doesn't matter");
    println!();

    // Recreate original test scenario
    let varmap = VarMap::new();
    let vb = candle_nn::VarBuilder::from_varmap(&varmap, candle_core::DType::F32, &device);

    let source = vb.get((2, 3), "source")?;
    log_shape("Source (learnable)", &source)?;

    let indices = Tensor::new(&[[0i64, 1, 2], [2, 1, 0]], &device)?;
    log_shape("Indices", &indices)?;

    let base = Tensor::zeros((2, 5), candle_core::DType::F32, &device)?;
    log_shape("Base", &base)?;

    let result = base.scatter_add(&indices, &source, 1)?;
    log_shape("Result", &result)?;

    let loss = result.sum_all()?;
    let grads = loss.backward()?;

    let all_vars = varmap.all_vars();
    if let Some(grad) = grads.get(&all_vars[0]) {
        let grad_norm = grad.sqr()?.sum_all()?.to_scalar::<f32>()?.sqrt();
        println!("\nGradient norm: {:.6}", grad_norm);

        if grad_norm > 1e-6 {
            println!("✓ GRADIENTS FLOW CORRECTLY");
        } else {
            println!("✗ GRADIENT COLLAPSE DETECTED");
        }
    }

    println!("\nConclusion: Original test is representative of production usage");

    Ok(())
}