foxhunt/crates/ml/tests/ppo_recurrent_integration_tests.rs

#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
)]
//! 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)]

// candle eliminated — test uses native APIs
use ml::dqn::{FactoredAction, ExposureLevel, OrderType, Urgency};
use tracing::info;
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: 63,
        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,
        use_symlog: true,
        use_adaptive_entropy: true,
        use_percentile_scaling: true,
    };

    let mut ppo = PPO::new(config.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 = FactoredAction::new(ExposureLevel::LongFull, OrderType::Market, Urgency::Normal);
        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);

    info!(policy_loss, value_loss, "test_recurrent_ppo_single_episode passed");
}

// ============================================================================
// 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: 63,
        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 ppo = PPO::new(config.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: f32 = h0_policy.iter().sum();
    let h0_sum_value: f32 = h0_value.iter().sum();

    // 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");

    info!("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: 63,
        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 mut ppo = PPO::new(config.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],
            FactoredAction::new(ExposureLevel::LongFull, OrderType::Market, Urgency::Normal),
            -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],
            FactoredAction::new(ExposureLevel::ShortSmall, OrderType::Market, Urgency::Normal),
            -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 actual_batch_size = manager.batch_size();

        // Create done mask: all environments done (match actual batch size)
        let done_mask = vec![true; actual_batch_size];

        // 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: f32 = h_policy.iter().sum();
        assert_eq!(h_sum, 0.0, "Hidden states should be zeros after reset");
    }

    info!("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: 63,
        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()
    };

    // Create feedforward PPO
    let mut feedforward_config = base_config.clone();
    feedforward_config.use_lstm = false;
    let mut ppo_feedforward = PPO::new(feedforward_config)
        .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::new(recurrent_config)
        .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],
            FactoredAction::new(ExposureLevel::LongFull, OrderType::Market, Urgency::Normal),
            -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");

    info!(ff_policy_loss, ff_value_loss, rec_policy_loss, rec_value_loss, "test_recurrent_vs_feedforward_ppo passed");
}

// ============================================================================
// 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.
    // `LSTMPolicyNetwork` / `LSTMValueNetwork` currently lack `from_varbuilder`
    // constructors, so safetensors-based checkpoint loading is not wired; the
    // test is left under #[ignore] until those constructors exist.
    let config = PPOConfig {
        state_dim: 16,
        num_actions: 63,
        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 mut ppo = PPO::new(config.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],
            FactoredAction::new(ExposureLevel::LongFull, OrderType::Market, Urgency::Normal),
            -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 checkpoint_path = checkpoint_dir.join("checkpoint");

    let result = ppo.save_checkpoint(&checkpoint_path);
    assert!(result.is_ok(), "Checkpoint save should succeed for recurrent PPO");

    // Verify checkpoint config file exists (save_checkpoint writes a .json sidecar)
    let config_path = checkpoint_path.with_extension("json");
    assert!(config_path.exists(), "Config file should exist");

    // Load checkpoint into new PPO instance
    let ppo_loaded = PPO::load_checkpoint(&checkpoint_path);
    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();

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