foxhunt/ml/tests/ppo_lstm_architecture_tests.rs

//! LSTM Architecture Tests for Recurrent PPO
//!
//! This module provides comprehensive tests for LSTM-augmented PPO networks:
//! - Policy network with LSTM layers
//! - Value network with LSTM layers
//! - Hidden state propagation across timesteps
//! - Sequence batching for temporal modeling

#![allow(unused_crate_dependencies)]

use candle_core::{Device, Tensor};
use ml::ppo::lstm_networks::{LSTMPolicyNetwork, LSTMValueNetwork};
use ml::ppo::PPOConfig;

// ============================================================================
// LSTM Policy Network Tests
// ============================================================================

#[test]
fn test_lstm_policy_network_output() {
    // Test that LSTM policy network outputs correct shapes for logits and hidden states
    let device = Device::Cpu;
    let input_dim = 64;
    let hidden_dim = 128;
    let num_actions = 45;
    let num_layers = 1;
    let batch_size = 32;

    let network = LSTMPolicyNetwork::new(
        input_dim,
        hidden_dim,
        num_layers,
        num_actions,
        device.clone(),
    )
    .expect("Failed to create LSTM policy network");

    // Create input state: [batch_size, input_dim]
    let state = Tensor::randn(0f32, 1.0, (batch_size, input_dim), &device)
        .expect("Failed to create input tensor");

    // Initial hidden and cell states: [num_layers, batch_size, hidden_dim]
    let h0 = Tensor::zeros((num_layers, batch_size, hidden_dim), candle_core::DType::F32, &device)
        .expect("Failed to create initial hidden state");
    let c0 = Tensor::zeros((num_layers, batch_size, hidden_dim), candle_core::DType::F32, &device)
        .expect("Failed to create initial cell state");

    // Forward pass
    let (logits, h_t, c_t) = network
        .forward(&state, &h0, &c0)
        .expect("Forward pass failed");

    // Verify output shapes
    assert_eq!(logits.dims(), &[batch_size, num_actions], "Logits shape mismatch");
    assert_eq!(h_t.dims(), &[num_layers, batch_size, hidden_dim], "Hidden state shape mismatch");
    assert_eq!(c_t.dims(), &[num_layers, batch_size, hidden_dim], "Cell state shape mismatch");

    println!("✅ LSTM policy network output test passed");
}

#[test]
fn test_lstm_value_network_output() {
    // Test that LSTM value network outputs correct shapes for values and hidden states
    let device = Device::Cpu;
    let input_dim = 64;
    let hidden_dim = 256;
    let num_layers = 2;
    let batch_size = 16;

    let network = LSTMValueNetwork::new(
        input_dim,
        hidden_dim,
        num_layers,
        device.clone(),
    )
    .expect("Failed to create LSTM value network");

    // Create input state: [batch_size, input_dim]
    let state = Tensor::randn(0f32, 1.0, (batch_size, input_dim), &device)
        .expect("Failed to create input tensor");

    // Initial hidden and cell states: [num_layers, batch_size, hidden_dim]
    let h0 = Tensor::zeros((num_layers, batch_size, hidden_dim), candle_core::DType::F32, &device)
        .expect("Failed to create initial hidden state");
    let c0 = Tensor::zeros((num_layers, batch_size, hidden_dim), candle_core::DType::F32, &device)
        .expect("Failed to create initial cell state");

    // Forward pass
    let (value, h_t, c_t) = network
        .forward(&state, &h0, &c0)
        .expect("Forward pass failed");

    // Verify output shapes
    assert_eq!(value.dims(), &[batch_size], "Value shape mismatch");
    assert_eq!(h_t.dims(), &[num_layers, batch_size, hidden_dim], "Hidden state shape mismatch");
    assert_eq!(c_t.dims(), &[num_layers, batch_size, hidden_dim], "Cell state shape mismatch");

    println!("✅ LSTM value network output test passed");
}

// ============================================================================
// Hidden State Propagation Tests
// ============================================================================

#[test]
fn test_lstm_hidden_state_propagation() {
    // Test that hidden states correctly propagate across multiple timesteps
    let device = Device::Cpu;
    let input_dim = 32;
    let hidden_dim = 64;
    let num_layers = 1;
    let num_actions = 3;
    let batch_size = 8;
    let seq_len = 5;

    let network = LSTMPolicyNetwork::new(
        input_dim,
        hidden_dim,
        num_layers,
        num_actions,
        device.clone(),
    )
    .expect("Failed to create LSTM policy network");

    // Initialize hidden states
    let mut h_t = Tensor::zeros((num_layers, batch_size, hidden_dim), candle_core::DType::F32, &device)
        .expect("Failed to create initial hidden state");
    let mut c_t = Tensor::zeros((num_layers, batch_size, hidden_dim), candle_core::DType::F32, &device)
        .expect("Failed to create initial cell state");

    // Process sequence of states
    for t in 0..seq_len {
        let state = Tensor::randn(0f32, 1.0, (batch_size, input_dim), &device)
            .expect(&format!("Failed to create state tensor at timestep {}", t));

        let (logits, new_h, new_c) = network
            .forward(&state, &h_t, &c_t)
            .expect(&format!("Forward pass failed at timestep {}", t));

        // Verify shapes remain consistent
        assert_eq!(logits.dims(), &[batch_size, num_actions]);
        assert_eq!(new_h.dims(), &[num_layers, batch_size, hidden_dim]);
        assert_eq!(new_c.dims(), &[num_layers, batch_size, hidden_dim]);

        // Hidden states should NOT be all zeros after first timestep (indicating information flow)
        if t > 0 {
            let h_sum = new_h.sum_all()
                .expect("Failed to sum hidden state")
                .to_scalar::<f32>()
                .expect("Failed to extract scalar");
            assert!(h_sum.abs() > 1e-6, "Hidden state is all zeros at timestep {}", t);
        }

        // Update hidden states for next timestep
        h_t = new_h;
        c_t = new_c;
    }

    println!("✅ LSTM hidden state propagation test passed");
}

// ============================================================================
// Sequence Batching Tests
// ============================================================================

#[test]
fn test_lstm_sequence_batching() {
    // Test that LSTM networks can process sequences in batch format: [seq_len, batch, features]
    let device = Device::Cpu;
    let input_dim = 16;
    let hidden_dim = 32;
    let num_layers = 1;
    let num_actions = 5;
    let seq_len = 10;
    let batch_size = 4;

    let network = LSTMPolicyNetwork::new(
        input_dim,
        hidden_dim,
        num_layers,
        num_actions,
        device.clone(),
    )
    .expect("Failed to create LSTM policy network");

    // Create sequence batch: process each timestep separately (as PPO collects online)
    let mut h_t = Tensor::zeros((num_layers, batch_size, hidden_dim), candle_core::DType::F32, &device)
        .expect("Failed to create initial hidden state");
    let mut c_t = Tensor::zeros((num_layers, batch_size, hidden_dim), candle_core::DType::F32, &device)
        .expect("Failed to create initial cell state");

    let mut all_logits = Vec::new();

    // Process sequence one timestep at a time
    for t in 0..seq_len {
        // Each timestep: [batch_size, input_dim]
        let state = Tensor::randn(0f32, 1.0, (batch_size, input_dim), &device)
            .expect(&format!("Failed to create state at timestep {}", t));

        let (logits, new_h, new_c) = network
            .forward(&state, &h_t, &c_t)
            .expect(&format!("Forward pass failed at timestep {}", t));

        all_logits.push(logits);
        h_t = new_h;
        c_t = new_c;
    }

    // Verify we processed all timesteps
    assert_eq!(all_logits.len(), seq_len);

    // Verify each timestep has correct batch dimension
    for (t, logits) in all_logits.iter().enumerate() {
        assert_eq!(logits.dims(), &[batch_size, num_actions], "Logits shape mismatch at timestep {}", t);
    }

    println!("✅ LSTM sequence batching test passed");
}

// ============================================================================
// Configuration Tests
// ============================================================================

#[test]
fn test_ppo_config_lstm_fields() {
    // Test that PPOConfig includes LSTM configuration fields
    let config = PPOConfig {
        use_lstm: true,
        lstm_hidden_dim: 256,
        lstm_num_layers: 2,
        ..PPOConfig::default()
    };

    assert_eq!(config.use_lstm, true, "use_lstm field not set correctly");
    assert_eq!(config.lstm_hidden_dim, 256, "lstm_hidden_dim field not set correctly");
    assert_eq!(config.lstm_num_layers, 2, "lstm_num_layers field not set correctly");

    println!("✅ PPOConfig LSTM fields test passed");
}

#[test]
fn test_ppo_config_backward_compatible() {
    // Test that default PPOConfig has LSTM disabled for backward compatibility
    let config = PPOConfig::default();

    assert_eq!(config.use_lstm, false, "LSTM should be disabled by default");
    assert!(config.lstm_hidden_dim > 0, "LSTM hidden dim should have valid default");
    assert!(config.lstm_num_layers > 0, "LSTM num layers should have valid default");

    println!("✅ PPOConfig backward compatibility test passed");
}