foxhunt/ml/tests/ppo_recurrent_integration_tests.rs

//! Recurrent PPO Integration Tests
//!
//! This module provides comprehensive integration tests for LSTM-enhanced PPO:
//! 1. Single episode training with LSTM
//! 2. Hidden state continuity across timesteps
//! 3. Episode boundary handling (hidden state reset)
//! 4. Recurrent vs feedforward comparison
//! 5. Checkpointing with hidden states
//!
//! These tests verify end-to-end functionality of the recurrent PPO system.

#![allow(unused_crate_dependencies)]

use candle_core::{Device, Tensor, DType};
use ml::dqn::TradingAction;
use ml::ppo::{
    ppo::{PPOConfig, PPO},
    trajectories::{Trajectory, TrajectoryBatch, TrajectoryStep},
    gae::GAEConfig,
};
use std::path::PathBuf;

// ============================================================================
// Test 1: Single Episode Training with LSTM
// ============================================================================

#[test]
fn test_recurrent_ppo_single_episode() {
    // Test that LSTM-enhanced PPO can train on a single episode
    let config = PPOConfig {
        state_dim: 32,
        num_actions: 5,
        policy_hidden_dims: vec![64],
        value_hidden_dims: vec![64],
        use_lstm: true,
        lstm_hidden_dim: 128,
        lstm_num_layers: 1,
        batch_size: 16,
        mini_batch_size: 8,
        num_epochs: 3,
        policy_learning_rate: 1e-4,
        value_learning_rate: 1e-3,
        clip_epsilon: 0.2,
        value_loss_coeff: 1.0,
        entropy_coeff: 0.01,
        gae_config: GAEConfig {
            gamma: 0.99,
            lambda: 0.95,
            normalize_advantages: true,
        },
        max_grad_norm: 0.5,
        early_stopping_enabled: false,
        early_stopping_patience: 10,
        early_stopping_min_delta: 1e-4,
        early_stopping_min_epochs: 20,
        max_position_absolute: 2.0,
        transaction_cost_bps: 0.10,
        cash_reserve_pct: 20.0,
        circuit_breaker_threshold: 5,
        lstm_sequence_length: 32,
        accumulation_steps: 1,
        clip_epsilon_high: None,
    };

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let mut ppo = PPO::with_device(config.clone(), device.clone())
        .expect("Failed to create recurrent PPO");

    // Verify hidden state manager is initialized
    assert!(ppo.hidden_state_manager.is_some(), "Hidden state manager should be initialized when use_lstm=true");

    // Create a simple trajectory (10 timesteps)
    let mut trajectory = Trajectory::new();
    for t in 0..10 {
        let state = vec![t as f32; 32]; // Simple incrementing state
        let action = TradingAction::Buy;
        let log_prob = -1.0;
        let value = 5.0 + t as f32;
        let reward = 1.0;
        let done = t == 9;

        trajectory.add_step(TrajectoryStep::new(state, action, log_prob, value, reward, done));
    }

    // Create batch with simple advantages and returns
    let advantages = vec![1.0; 10];
    let returns = vec![10.0; 10];
    let mut batch = TrajectoryBatch::from_trajectories(vec![trajectory], advantages, returns);

    // Normalize advantages
    batch.normalize_advantages().expect("Failed to normalize advantages");

    // Train for 1 epoch
    let (policy_loss, value_loss) = ppo.update(&mut batch)
        .expect("Training should complete without errors");

    // Verify losses are finite (not NaN or Inf)
    assert!(policy_loss.is_finite(), "Policy loss should be finite, got {}", policy_loss);
    assert!(value_loss.is_finite(), "Value loss should be finite, got {}", value_loss);

    println!("✅ test_recurrent_ppo_single_episode passed: policy_loss={:.4}, value_loss={:.4}",
             policy_loss, value_loss);
}

// ============================================================================
// Test 2: Hidden State Continuity Across Timesteps
// ============================================================================

#[test]
fn test_recurrent_ppo_hidden_state_continuity() {
    // Test that hidden states persist and evolve across timesteps within an episode
    let config = PPOConfig {
        state_dim: 16,
        num_actions: 3,
        policy_hidden_dims: vec![32],
        value_hidden_dims: vec![32],
        use_lstm: true,
        lstm_hidden_dim: 64,
        lstm_num_layers: 1,
        batch_size: 8,
        mini_batch_size: 4,
        num_epochs: 2,
        policy_learning_rate: 1e-4,
        value_learning_rate: 1e-3,
        lstm_sequence_length: 32,
        ..PPOConfig::default()
    };

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let ppo = PPO::with_device(config.clone(), device.clone())
        .expect("Failed to create recurrent PPO");

    // Get hidden state manager
    let manager = ppo.hidden_state_manager.as_ref()
        .expect("Hidden state manager should exist");

    // Extract initial hidden states (should be zeros)
    let (h0_policy, _c0_policy) = manager.get_policy_state();
    let (h0_value, _c0_value) = manager.get_value_state();

    let h0_sum_policy = h0_policy.sum_all().expect("Sum failed").to_scalar::<f32>().expect("To scalar failed");
    let h0_sum_value = h0_value.sum_all().expect("Sum failed").to_scalar::<f32>().expect("To scalar failed");

    // Verify initial states are zeros
    assert_eq!(h0_sum_policy, 0.0, "Initial policy hidden state should be zeros");
    assert_eq!(h0_sum_value, 0.0, "Initial value hidden state should be zeros");

    println!("✅ test_recurrent_ppo_hidden_state_continuity passed: States initialized correctly");
}

// ============================================================================
// Test 3: Episode Boundary Handling (Hidden State Reset)
// ============================================================================

#[test]
fn test_recurrent_ppo_episode_boundaries() {
    // Test that hidden states reset between episodes
    let config = PPOConfig {
        state_dim: 16,
        num_actions: 3,
        policy_hidden_dims: vec![32],
        value_hidden_dims: vec![32],
        use_lstm: true,
        lstm_hidden_dim: 64,
        lstm_num_layers: 1,
        batch_size: 8,
        mini_batch_size: 4,
        num_epochs: 2,
        lstm_sequence_length: 32,
        ..PPOConfig::default()
    };

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let mut ppo = PPO::with_device(config.clone(), device.clone())
        .expect("Failed to create recurrent PPO");

    // Create two episodes
    let mut episode1 = Trajectory::new();
    for t in 0..5 {
        episode1.add_step(TrajectoryStep::new(
            vec![1.0; 16],
            TradingAction::Buy,
            -1.0,
            5.0,
            1.0,
            t == 4, // Done at last step
        ));
    }

    let mut episode2 = Trajectory::new();
    for t in 0..5 {
        episode2.add_step(TrajectoryStep::new(
            vec![2.0; 16],
            TradingAction::Sell,
            -1.0,
            5.0,
            1.0,
            t == 4, // Done at last step
        ));
    }

    // Create batch from both episodes
    let advantages = vec![1.0; 10]; // 5 steps per episode × 2 episodes
    let returns = vec![10.0; 10];
    let mut batch = TrajectoryBatch::from_trajectories(
        vec![episode1, episode2],
        advantages,
        returns,
    );

    batch.normalize_advantages().expect("Failed to normalize advantages");

    // Train on batch (should handle episode boundaries)
    let result = ppo.update(&mut batch);
    assert!(result.is_ok(), "Training should complete without errors across episode boundaries");

    // Verify hidden state manager can reset states
    if let Some(ref mut manager) = ppo.hidden_state_manager {
        // Get actual batch size from hidden state manager
        let (h_policy, _) = manager.get_policy_state();
        let actual_batch_size = h_policy.dims()[1]; // Shape: [num_layers, batch_size, hidden_dim]

        // Create done mask: all environments done (match actual batch size)
        let done_mask = Tensor::ones(&[actual_batch_size], DType::U8, &device)
            .expect("Failed to create done mask");

        // Reset states
        manager.reset_on_done(&done_mask).expect("Reset should succeed");

        // Verify states are zeros after reset
        let (h_policy, _c_policy) = manager.get_policy_state();
        let h_sum = h_policy.sum_all().expect("Sum failed").to_scalar::<f32>().expect("To scalar failed");
        assert_eq!(h_sum, 0.0, "Hidden states should be zeros after reset");
    }

    println!("✅ test_recurrent_ppo_episode_boundaries passed: Episode boundaries handled correctly");
}

// ============================================================================
// Test 4: Recurrent vs Feedforward Comparison
// ============================================================================

#[test]
fn test_recurrent_vs_feedforward_ppo() {
    // Compare LSTM vs non-LSTM PPO training behavior
    let base_config = PPOConfig {
        state_dim: 16,
        num_actions: 3,
        policy_hidden_dims: vec![32],
        value_hidden_dims: vec![32],
        batch_size: 16,
        mini_batch_size: 8,
        num_epochs: 5,
        policy_learning_rate: 1e-4,
        value_learning_rate: 1e-3,
        lstm_sequence_length: 32,
        ..PPOConfig::default()
    };

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);

    // Create feedforward PPO
    let mut feedforward_config = base_config.clone();
    feedforward_config.use_lstm = false;
    let mut ppo_feedforward = PPO::with_device(feedforward_config, device.clone())
        .expect("Failed to create feedforward PPO");

    // Create recurrent PPO
    let mut recurrent_config = base_config.clone();
    recurrent_config.use_lstm = true;
    recurrent_config.lstm_hidden_dim = 128;
    recurrent_config.lstm_num_layers = 1;
    let mut ppo_recurrent = PPO::with_device(recurrent_config, device.clone())
        .expect("Failed to create recurrent PPO");

    // Verify configurations
    assert!(ppo_feedforward.hidden_state_manager.is_none(), "Feedforward PPO should not have hidden state manager");
    assert!(ppo_recurrent.hidden_state_manager.is_some(), "Recurrent PPO should have hidden state manager");

    // Create identical trajectory for both
    let mut trajectory = Trajectory::new();
    for t in 0..10 {
        trajectory.add_step(TrajectoryStep::new(
            vec![t as f32; 16],
            TradingAction::Buy,
            -1.0,
            5.0,
            1.0,
            t == 9,
        ));
    }

    let advantages = vec![1.0; 10];
    let returns = vec![10.0; 10];

    // Train feedforward
    let mut batch_ff = TrajectoryBatch::from_trajectories(vec![trajectory.clone()], advantages.clone(), returns.clone());
    batch_ff.normalize_advantages().expect("Failed to normalize");
    let (ff_policy_loss, ff_value_loss) = ppo_feedforward.update(&mut batch_ff)
        .expect("Feedforward training failed");

    // Train recurrent
    let mut batch_rec = TrajectoryBatch::from_trajectories(vec![trajectory], advantages, returns);
    batch_rec.normalize_advantages().expect("Failed to normalize");
    let (rec_policy_loss, rec_value_loss) = ppo_recurrent.update(&mut batch_rec)
        .expect("Recurrent training failed");

    // Both should produce finite losses
    assert!(ff_policy_loss.is_finite(), "Feedforward policy loss should be finite");
    assert!(ff_value_loss.is_finite(), "Feedforward value loss should be finite");
    assert!(rec_policy_loss.is_finite(), "Recurrent policy loss should be finite");
    assert!(rec_value_loss.is_finite(), "Recurrent value loss should be finite");

    println!("✅ test_recurrent_vs_feedforward_ppo passed:");
    println!("   Feedforward: policy_loss={:.4}, value_loss={:.4}", ff_policy_loss, ff_value_loss);
    println!("   Recurrent:   policy_loss={:.4}, value_loss={:.4}", rec_policy_loss, rec_value_loss);
}

// ============================================================================
// Test 5: Checkpointing with Hidden States
// ============================================================================

#[test]
#[ignore = "LSTM checkpoint loading not yet implemented (requires from_varbuilder methods in lstm_networks.rs)"]
fn test_recurrent_ppo_checkpointing() {
    // Test that checkpoints can be saved and loaded with LSTM configuration
    // TODO: Implement from_varbuilder methods in LSTMPolicyNetwork and LSTMValueNetwork
    // to enable loading LSTM checkpoints from safetensors files
    let config = PPOConfig {
        state_dim: 16,
        num_actions: 3,
        policy_hidden_dims: vec![32],
        value_hidden_dims: vec![32],
        use_lstm: true,
        lstm_hidden_dim: 64,
        lstm_num_layers: 1,
        batch_size: 16,
        mini_batch_size: 8,
        num_epochs: 3,
        lstm_sequence_length: 32,
        ..PPOConfig::default()
    };

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let mut ppo = PPO::with_device(config.clone(), device.clone())
        .expect("Failed to create recurrent PPO");

    // Train for a few steps
    let mut trajectory = Trajectory::new();
    for t in 0..10 {
        trajectory.add_step(TrajectoryStep::new(
            vec![t as f32; 16],
            TradingAction::Buy,
            -1.0,
            5.0,
            1.0,
            t == 9,
        ));
    }

    let advantages = vec![1.0; 10];
    let returns = vec![10.0; 10];
    let mut batch = TrajectoryBatch::from_trajectories(vec![trajectory], advantages, returns);
    batch.normalize_advantages().expect("Failed to normalize");

    let (_loss_before, _) = ppo.update(&mut batch)
        .expect("Training failed");

    // Save checkpoint
    let checkpoint_dir = PathBuf::from("/tmp/ppo_recurrent_checkpoint_test");
    std::fs::create_dir_all(&checkpoint_dir).expect("Failed to create checkpoint dir");

    let actor_path = checkpoint_dir.join("actor_epoch_1.safetensors");
    let critic_path = checkpoint_dir.join("critic_epoch_1.safetensors");
    let metadata_path = checkpoint_dir.join("metadata_epoch_1.json");

    let result = ppo.save_checkpoint(
        actor_path.to_str().unwrap(),
        critic_path.to_str().unwrap(),
        metadata_path.to_str().unwrap(),
    );
    assert!(result.is_ok(), "Checkpoint save should succeed for recurrent PPO");

    // Verify checkpoint files exist
    assert!(actor_path.exists(), "Actor checkpoint should exist");
    assert!(critic_path.exists(), "Critic checkpoint should exist");
    assert!(metadata_path.exists(), "Metadata checkpoint should exist");

    // Load checkpoint into new PPO instance
    let ppo_loaded = PPO::load_checkpoint(
        actor_path.to_str().unwrap(),
        critic_path.to_str().unwrap(),
        config.clone(),
        device.clone(),
    );
    assert!(ppo_loaded.is_ok(), "Checkpoint load should succeed for recurrent PPO");

    // Note: load_checkpoint does not restore training_steps from metadata.
    // This is expected behavior - training_steps start from 0 after loading.
    // The important part is that the model weights are restored correctly.

    // Clean up
    std::fs::remove_dir_all(&checkpoint_dir).ok();

    println!("✅ test_recurrent_ppo_checkpointing passed: Checkpoints save/load correctly");
}