foxhunt/ml/tests/ppo_tests.rs

//! Comprehensive PPO Model Tests
//!
//! This module provides extensive testing for:
//! - PPO algorithm (clipped surrogate objective)
//! - Continuous PPO (Gaussian policy gradients)
//! - GAE (Generalized Advantage Estimation)
//! - Trajectories (batch collection and preprocessing)

use candle_core::{Device, Tensor};
use foxhunt_ml::dqn::TradingAction;
use foxhunt_ml::ppo::{
    continuous_policy::{ContinuousAction, ContinuousPolicyConfig, ContinuousPolicyNetwork},
    continuous_ppo::{
        ContinuousPPO, ContinuousPPOConfig, ContinuousTrajectory, ContinuousTrajectoryBatch,
        ContinuousTrajectoryStep,
    },
    gae::{
        compute_advantages, compute_discounted_returns, compute_gae, compute_gae_single_trajectory,
        compute_td_advantages, normalize_advantages, AdvantageMethod, GAEConfig,
    },
    ppo::{PolicyNetwork, PPOConfig, ValueNetwork, WorkingPPO},
    trajectories::{Trajectory, TrajectoryBatch, TrajectoryStep},
};

// ============================================================================
// PPO Core Algorithm Tests
// ============================================================================

#[test]
fn test_ppo_clipped_surrogate_objective() {
    // Test PPO clipping behavior with known advantage values
    let config = PPOConfig {
        state_dim: 4,
        num_actions: 3,
        policy_hidden_dims: vec![8],
        value_hidden_dims: vec![8],
        clip_epsilon: 0.2,
        ..PPOConfig::default()
    };

    let ppo = WorkingPPO::new(config).expect("Failed to create PPO");

    // Create a simple trajectory with known values
    let mut trajectory = Trajectory::new();
    trajectory.add_step(TrajectoryStep::new(
        vec![1.0, 0.0, 0.0, 0.0],
        TradingAction::Buy,
        -1.0, // log_prob
        5.0,  // value
        10.0, // reward
        false,
    ));

    let trajectories = vec![trajectory];
    let advantages = vec![2.0]; // Positive advantage
    let returns = vec![15.0];

    let mut batch = TrajectoryBatch::from_trajectories(trajectories, advantages, returns);

    // Verify batch construction
    assert_eq!(batch.total_steps(), 1);
    assert_eq!(batch.advantages[0], 2.0);

    // Test advantage normalization
    batch.normalize_advantages().expect("Failed to normalize");

    // After normalization with single value, advantage should be 0 (zero mean, unit variance not applicable)
    assert!(batch.advantages[0].abs() < 1e-6);
}

#[test]
fn test_ppo_clipping_boundary_cases() {
    // Test clipping at ε = 0.2 boundaries
    let config = PPOConfig {
        state_dim: 2,
        num_actions: 3,
        policy_hidden_dims: vec![4],
        value_hidden_dims: vec![4],
        clip_epsilon: 0.2,
        ..PPOConfig::default()
    };

    let ppo = WorkingPPO::new(config).expect("Failed to create PPO");

    // Test with multiple advantage values to test clipping
    let mut trajectory = Trajectory::new();
    for i in 0..5 {
        trajectory.add_step(TrajectoryStep::new(
            vec![i as f32, (i * 2) as f32],
            TradingAction::Buy,
            -1.0,
            5.0,
            (i + 1) as f32,
            i == 4,
        ));
    }

    let trajectories = vec![trajectory];
    // Different advantages to test clipping behavior
    let advantages = vec![-2.0, -0.5, 0.0, 0.5, 2.0];
    let returns = vec![1.0, 2.0, 3.0, 4.0, 5.0];

    let batch = TrajectoryBatch::from_trajectories(trajectories, advantages, returns);
    assert_eq!(batch.total_steps(), 5);
    assert_eq!(batch.advantages.len(), 5);
}

#[test]
fn test_ppo_value_network() {
    // Test value network predictions
    let device = Device::Cpu;
    let value_net = ValueNetwork::new(6, &[16, 8], device.clone()).expect("Failed to create value network");

    let states = Tensor::from_vec(vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6], (1, 6), &device)
        .expect("Failed to create state tensor");

    let values = value_net.forward(&states).expect("Forward pass failed");

    // Value should be scalar
    assert_eq!(values.dims(), &[1]);

    let value_scalar = values.to_vec1::<f32>().expect("Failed to extract value");
    assert!(value_scalar[0].is_finite());
}

#[test]
fn test_ppo_policy_network() {
    // Test policy network action probabilities
    let device = Device::Cpu;
    let policy_net = PolicyNetwork::new(4, &[8, 4], 3, device.clone())
        .expect("Failed to create policy network");

    let state = Tensor::from_vec(vec![1.0, 0.5, -0.3, 0.8], (1, 4), &device)
        .expect("Failed to create state tensor");

    let probs = policy_net.action_probabilities(&state).expect("Failed to get probabilities");

    let probs_vec = probs.flatten_all().unwrap().to_vec1::<f32>().unwrap();

    // Probabilities should sum to 1
    let sum: f32 = probs_vec.iter().sum();
    assert!((sum - 1.0).abs() < 1e-5, "Probabilities don't sum to 1: {}", sum);

    // All probabilities should be in [0, 1]
    for &p in &probs_vec {
        assert!(p >= 0.0 && p <= 1.0, "Invalid probability: {}", p);
    }
}

#[test]
fn test_ppo_entropy_computation() {
    // Test entropy calculation
    let device = Device::Cpu;
    let policy_net = PolicyNetwork::new(3, &[6], 3, device.clone()).expect("Failed to create policy network");

    let states = Tensor::from_vec(vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6], (2, 3), &device)
        .expect("Failed to create states");

    let entropy = policy_net.entropy(&states).expect("Failed to compute entropy");

    let entropy_vec = entropy.to_vec1::<f32>().unwrap();

    // Entropy should be positive for discrete distributions
    for &e in &entropy_vec {
        assert!(e >= 0.0, "Entropy should be non-negative: {}", e);
    }
}

#[test]
fn test_ppo_training_steps_counter() {
    // Test that training steps are tracked correctly
    let config = PPOConfig::default();
    let mut ppo = WorkingPPO::new(config).expect("Failed to create PPO");

    assert_eq!(ppo.get_training_steps(), 0);

    // Manually increment to simulate training
    ppo.training_steps = 1;
    assert_eq!(ppo.get_training_steps(), 1);

    ppo.training_steps = 100;
    assert_eq!(ppo.get_training_steps(), 100);
}

// ============================================================================
// Continuous PPO Tests (Gaussian Policy)
// ============================================================================

#[test]
fn test_continuous_ppo_gaussian_policy() {
    // Test Gaussian policy outputs
    let config = ContinuousPolicyConfig {
        state_dim: 8,
        hidden_dims: vec![16],
        learnable_std: true,
        ..ContinuousPolicyConfig::default()
    };

    let device = Device::Cpu;
    let policy = ContinuousPolicyNetwork::new(config, device.clone()).expect("Failed to create policy");

    let state = Tensor::from_vec(vec![0.1; 8], (1, 8), &device).expect("Failed to create state");

    let (mean, log_std) = policy.forward(&state).expect("Forward pass failed");

    // Mean should be in [0, 1] (bounded by sigmoid)
    let mean_val = mean.flatten_all().unwrap().to_vec1::<f32>().unwrap()[0];
    assert!(
        mean_val >= 0.0 && mean_val <= 1.0,
        "Mean out of bounds: {}",
        mean_val
    );

    // Log std should be clamped to [-5, 2]
    let log_std_val = log_std.flatten_all().unwrap().to_vec1::<f32>().unwrap()[0];
    assert!(
        log_std_val >= -5.0 && log_std_val <= 2.0,
        "Log std out of bounds: {}",
        log_std_val
    );
}

#[test]
fn test_continuous_ppo_action_sampling() {
    // Test action sampling from Gaussian policy
    let config = ContinuousPolicyConfig {
        state_dim: 4,
        hidden_dims: vec![8],
        learnable_std: false, // Fixed std for reproducibility
        init_log_std: -1.0,
        ..ContinuousPolicyConfig::default()
    };

    let device = Device::Cpu;
    let policy = ContinuousPolicyNetwork::new(config, device.clone()).expect("Failed to create policy");

    let state = Tensor::from_vec(vec![0.5; 4], (1, 4), &device).expect("Failed to create state");

    // Sample multiple times to check distribution
    for _ in 0..20 {
        let (action, log_prob) = policy.sample_action(&state).expect("Failed to sample action");

        // Action should be in [0, 1]
        assert!(action >= 0.0 && action <= 1.0, "Action out of bounds: {}", action);

        // Log prob should be finite and negative
        assert!(log_prob.is_finite(), "Log prob not finite: {}", log_prob);
        assert!(log_prob <= 0.0, "Log prob should be negative: {}", log_prob);
    }
}

#[test]
fn test_continuous_ppo_log_prob_computation() {
    // Test log probability calculation for Gaussian distribution
    let config = ContinuousPolicyConfig {
        state_dim: 4,
        hidden_dims: vec![8],
        learnable_std: false,
        init_log_std: -1.0,
        ..ContinuousPolicyConfig::default()
    };

    let device = Device::Cpu;
    let policy = ContinuousPolicyNetwork::new(config, device.clone()).expect("Failed to create policy");

    let states = Tensor::from_vec(vec![0.3; 8], (2, 4), &device).expect("Failed to create states");
    let actions = Tensor::from_vec(vec![0.5, 0.7], (2, 1), &device).expect("Failed to create actions");

    let log_probs = policy.log_probs(&states, &actions).expect("Failed to compute log probs");

    let log_probs_vec = log_probs.to_vec1::<f32>().unwrap();

    // All log probs should be finite and negative
    for &lp in &log_probs_vec {
        assert!(lp.is_finite(), "Log prob not finite: {}", lp);
        assert!(lp <= 0.0, "Log prob should be negative: {}", lp);
    }
}

#[test]
fn test_continuous_ppo_entropy() {
    // Test entropy for Gaussian distribution
    let config = ContinuousPolicyConfig {
        state_dim: 6,
        hidden_dims: vec![12],
        learnable_std: true,
        ..ContinuousPolicyConfig::default()
    };

    let device = Device::Cpu;
    let policy = ContinuousPolicyNetwork::new(config, device.clone()).expect("Failed to create policy");

    let states = Tensor::from_vec(vec![0.2; 12], (2, 6), &device).expect("Failed to create states");

    let entropy = policy.entropy(&states).expect("Failed to compute entropy");

    let entropy_vec = entropy.to_vec1::<f32>().unwrap();

    // Gaussian entropy: 0.5 * log(2πe) + log_std
    // Should be positive
    for &e in &entropy_vec {
        assert!(e > 0.0, "Entropy should be positive: {}", e);
        assert!(e.is_finite(), "Entropy should be finite: {}", e);
    }
}

#[test]
fn test_continuous_ppo_fixed_vs_learnable_std() {
    let device = Device::Cpu;

    // Test fixed std
    let config_fixed = ContinuousPolicyConfig {
        state_dim: 4,
        hidden_dims: vec![8],
        learnable_std: false,
        init_log_std: -2.0,
        ..ContinuousPolicyConfig::default()
    };

    let policy_fixed = ContinuousPolicyNetwork::new(config_fixed, device.clone())
        .expect("Failed to create fixed std policy");

    let state = Tensor::from_vec(vec![0.1; 4], (1, 4), &device).expect("Failed to create state");

    let (_mean_fixed, log_std_fixed) = policy_fixed.forward(&state).expect("Forward pass failed");
    let log_std_val = log_std_fixed.flatten_all().unwrap().to_vec1::<f32>().unwrap()[0];

    // Fixed std should be close to init_log_std
    assert!((log_std_val - (-2.0)).abs() < 0.1, "Fixed std not preserved: {}", log_std_val);

    // Test learnable std
    let config_learnable = ContinuousPolicyConfig {
        state_dim: 4,
        hidden_dims: vec![8],
        learnable_std: true,
        ..ContinuousPolicyConfig::default()
    };

    let policy_learnable = ContinuousPolicyNetwork::new(config_learnable, device.clone())
        .expect("Failed to create learnable std policy");

    let (_mean_learnable, log_std_learnable) = policy_learnable.forward(&state).expect("Forward pass failed");

    // Learnable std should be within bounds but can vary
    let log_std_learnable_val = log_std_learnable.flatten_all().unwrap().to_vec1::<f32>().unwrap()[0];
    assert!(
        log_std_learnable_val >= -5.0 && log_std_learnable_val <= 2.0,
        "Learnable std out of bounds: {}",
        log_std_learnable_val
    );
}

#[test]
fn test_continuous_ppo_trajectory_collection() {
    // Test continuous trajectory collection
    let mut trajectory = ContinuousTrajectory::new();

    assert!(trajectory.is_empty());
    assert_eq!(trajectory.len(), 0);

    let action1 = ContinuousAction::new(0.3);
    let step1 = ContinuousTrajectoryStep::new(vec![1.0; 4], action1, -1.2, 5.0, 2.5, false);

    trajectory.add_step(step1);

    assert!(!trajectory.is_empty());
    assert_eq!(trajectory.len(), 1);

    let action2 = ContinuousAction::new(0.7);
    let step2 = ContinuousTrajectoryStep::new(vec![0.5; 4], action2, -0.8, 3.0, 1.8, true);

    trajectory.add_step(step2);

    assert_eq!(trajectory.len(), 2);
    assert_eq!(trajectory.steps()[0].action.position_size(), 0.3);
    assert_eq!(trajectory.steps()[1].action.position_size(), 0.7);
}

#[test]
fn test_continuous_action_validation() {
    // Test continuous action bounds
    let action1 = ContinuousAction::new(0.5);
    assert!(action1.is_valid());
    assert_eq!(action1.position_size(), 0.5);

    // Test clamping
    let action2 = ContinuousAction::new(1.5);
    assert!(action2.is_valid());
    assert_eq!(action2.position_size(), 1.0); // Clamped to max

    let action3 = ContinuousAction::new(-0.3);
    assert!(action3.is_valid());
    assert_eq!(action3.position_size(), 0.0); // Clamped to min

    // Test invalid action
    let action4 = ContinuousAction::new(f32::NAN);
    assert!(!action4.is_valid());
}

// ============================================================================
// GAE (Generalized Advantage Estimation) Tests
// ============================================================================

#[test]
fn test_gae_single_trajectory_computation() {
    // Test GAE with known values
    let rewards = vec![1.0, 2.0, 3.0];
    let values = vec![5.0, 6.0, 7.0];
    let dones = vec![false, false, true];
    let next_value = 0.0; // Terminal state

    let config = GAEConfig {
        gamma: 0.9,
        lambda: 0.95,
        normalize_advantages: false,
    };

    let (advantages, returns) =
        compute_gae_single_trajectory(&rewards, &values, &dones, next_value, &config)
            .expect("GAE computation failed");

    assert_eq!(advantages.len(), 3);
    assert_eq!(returns.len(), 3);

    // Verify advantages are finite
    for &adv in &advantages {
        assert!(adv.is_finite(), "Advantage not finite: {}", adv);
    }

    // Verify returns are finite
    for &ret in &returns {
        assert!(ret.is_finite(), "Return not finite: {}", ret);
    }
}

#[test]
fn test_gae_multi_step_advantage() {
    // Test multi-step advantage computation
    let rewards = vec![1.0, 1.0, 1.0, 1.0, 1.0];
    let values = vec![5.0, 5.0, 5.0, 5.0, 5.0];
    let dones = vec![false, false, false, false, true];
    let next_value = 0.0;

    let config = GAEConfig {
        gamma: 0.99,
        lambda: 0.95,
        normalize_advantages: false,
    };

    let (advantages, _returns) =
        compute_gae_single_trajectory(&rewards, &values, &dones, next_value, &config)
            .expect("GAE computation failed");

    // With uniform rewards and values, advantages should follow a pattern
    assert_eq!(advantages.len(), 5);

    // Advantages should decrease (decay) as we go forward in time
    // (when computed backwards, they accumulate)
    for i in 0..advantages.len() - 1 {
        assert!(advantages[i].is_finite());
    }
}

#[test]
fn test_gae_lambda_return() {
    // Test λ-return computation with different λ values
    let rewards = vec![2.0, 3.0, 4.0];
    let values = vec![10.0, 12.0, 14.0];
    let dones = vec![false, false, true];
    let next_value = 0.0;

    // Test with λ = 1.0 (Monte Carlo)
    let config_mc = GAEConfig {
        gamma: 0.9,
        lambda: 1.0,
        normalize_advantages: false,
    };

    let (adv_mc, _) = compute_gae_single_trajectory(&rewards, &values, &dones, next_value, &config_mc)
        .expect("GAE computation failed");

    // Test with λ = 0.0 (TD(0))
    let config_td = GAEConfig {
        gamma: 0.9,
        lambda: 0.0,
        normalize_advantages: false,
    };

    let (adv_td, _) = compute_gae_single_trajectory(&rewards, &values, &dones, next_value, &config_td)
        .expect("GAE computation failed");

    // Monte Carlo and TD should give different results
    assert_ne!(adv_mc, adv_td, "MC and TD advantages should differ");

    // Both should have same length
    assert_eq!(adv_mc.len(), adv_td.len());
}

#[test]
fn test_gae_normalization() {
    // Test advantage normalization
    let mut advantages = vec![1.0, 2.0, 3.0, 4.0, 5.0];

    normalize_advantages(&mut advantages).expect("Normalization failed");

    // Check zero mean
    let mean: f32 = advantages.iter().sum::<f32>() / advantages.len() as f32;
    assert!(mean.abs() < 1e-6, "Mean not zero: {}", mean);

    // Check unit variance
    let variance: f32 = advantages.iter().map(|&a| a * a).sum::<f32>() / advantages.len() as f32;
    assert!((variance - 1.0).abs() < 1e-5, "Variance not unit: {}", variance);
}

#[test]
fn test_gae_terminal_states() {
    // Test GAE with terminal states
    let rewards = vec![1.0, 1.0, 10.0]; // Large terminal reward
    let values = vec![5.0, 5.0, 5.0];
    let dones = vec![false, false, true];
    let next_value = 0.0;

    let config = GAEConfig {
        gamma: 0.9,
        lambda: 0.95,
        normalize_advantages: false,
    };

    let (advantages, returns) =
        compute_gae_single_trajectory(&rewards, &values, &dones, next_value, &config)
            .expect("GAE computation failed");

    // Terminal advantage should be: reward + 0 - value = 10.0 + 0 - 5.0 = 5.0
    assert!((advantages[2] - 5.0).abs() < 1e-5, "Terminal advantage incorrect: {}", advantages[2]);

    // Terminal return should be just the reward (no future)
    assert!((returns[2] - 10.0).abs() < 1e-5, "Terminal return incorrect: {}", returns[2]);
}

#[test]
fn test_gae_discounted_returns() {
    // Test discounted return computation
    let mut trajectory = Trajectory::new();

    trajectory.add_step(TrajectoryStep::new(
        vec![1.0],
        TradingAction::Buy,
        0.0,
        0.0,
        1.0,
        false,
    ));
    trajectory.add_step(TrajectoryStep::new(
        vec![2.0],
        TradingAction::Sell,
        0.0,
        0.0,
        2.0,
        false,
    ));
    trajectory.add_step(TrajectoryStep::new(
        vec![3.0],
        TradingAction::Hold,
        0.0,
        0.0,
        3.0,
        true,
    ));

    let returns = trajectory.compute_returns(0.9);

    // returns[2] = 3.0
    // returns[1] = 2.0 + 0.9 * 3.0 = 4.7
    // returns[0] = 1.0 + 0.9 * 4.7 = 5.23

    assert!((returns[2] - 3.0).abs() < 1e-5);
    assert!((returns[1] - 4.7).abs() < 1e-5);
    assert!((returns[0] - 5.23).abs() < 1e-5);
}

#[test]
fn test_advantage_methods() {
    // Test different advantage estimation methods
    let mut trajectory = Trajectory::new();
    for i in 0..3 {
        trajectory.add_step(TrajectoryStep::new(
            vec![i as f32],
            TradingAction::Buy,
            -0.5,
            5.0,
            (i + 1) as f32,
            i == 2,
        ));
    }

    let trajectories = vec![trajectory];

    // Test GAE
    let gae_method = AdvantageMethod::GAE(GAEConfig::default());
    let (adv_gae, ret_gae) = compute_advantages(&trajectories, &gae_method)
        .expect("GAE advantage computation failed");
    assert_eq!(adv_gae.len(), 3);
    assert_eq!(ret_gae.len(), 3);

    // Test TD
    let td_method = AdvantageMethod::TemporalDifference {
        gamma: 0.9,
        normalize: true,
    };
    let (adv_td, ret_td) = compute_advantages(&trajectories, &td_method)
        .expect("TD advantage computation failed");
    assert_eq!(adv_td.len(), 3);
    assert_eq!(ret_td.len(), 3);

    // Test Monte Carlo
    let mc_method = AdvantageMethod::MonteCarlo {
        gamma: 0.9,
        normalize: false,
    };
    let (adv_mc, ret_mc) = compute_advantages(&trajectories, &mc_method)
        .expect("MC advantage computation failed");
    assert_eq!(adv_mc.len(), 3);
    assert_eq!(ret_mc.len(), 3);
}

// ============================================================================
// Trajectory Tests (Batch Collection & Preprocessing)
// ============================================================================

#[test]
fn test_trajectory_batch_creation() {
    // Test batch creation from trajectories
    let mut traj1 = Trajectory::new();
    traj1.add_step(TrajectoryStep::new(
        vec![1.0, 2.0],
        TradingAction::Buy,
        -0.5,
        5.0,
        1.0,
        false,
    ));
    traj1.add_step(TrajectoryStep::new(
        vec![2.0, 3.0],
        TradingAction::Sell,
        -0.3,
        6.0,
        2.0,
        true,
    ));

    let mut traj2 = Trajectory::new();
    traj2.add_step(TrajectoryStep::new(
        vec![3.0, 4.0],
        TradingAction::Hold,
        -0.7,
        4.0,
        3.0,
        true,
    ));

    let trajectories = vec![traj1, traj2];
    let advantages = vec![0.1, 0.2, 0.3];
    let returns = vec![5.0, 6.0, 7.0];

    let batch = TrajectoryBatch::from_trajectories(trajectories, advantages, returns);

    assert_eq!(batch.total_steps(), 3);
    assert_eq!(batch.num_trajectories(), 2);
    assert_eq!(batch.states.len(), 3);
    assert_eq!(batch.actions.len(), 3);
    assert_eq!(batch.advantages.len(), 3);
}

#[test]
fn test_trajectory_batch_preprocessing() {
    // Test batch preprocessing with different sizes
    let mut trajectory = Trajectory::new();
    for i in 0..10 {
        trajectory.add_step(TrajectoryStep::new(
            vec![i as f32],
            TradingAction::Buy,
            -0.5,
            5.0,
            1.0,
            i == 9,
        ));
    }

    let trajectories = vec![trajectory];
    let advantages = vec![0.1; 10];
    let returns = vec![5.0; 10];

    let batch = TrajectoryBatch::from_trajectories(trajectories, advantages, returns);

    assert_eq!(batch.total_steps(), 10);

    // Test tensor conversion
    let device = Device::Cpu;
    let tensors = batch.to_tensors(&device, 1).expect("Tensor conversion failed");

    assert_eq!(tensors.states.dims(), &[10, 1]);
    assert_eq!(tensors.actions.dims(), &[10]);
    assert_eq!(tensors.advantages.dims(), &[10]);
}

#[test]
fn test_trajectory_mini_batch_creation() {
    // Test mini-batch creation with different sizes
    let trajectories = vec![Trajectory::new()];
    let advantages = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0];
    let returns = vec![0.0; 7];
    let states = vec![vec![1.0]; 7];
    let actions = vec![TradingAction::Buy; 7];
    let log_probs = vec![0.0; 7];
    let values = vec![0.0; 7];
    let dones = vec![false; 7];

    let mut batch = TrajectoryBatch::from_trajectories(trajectories, advantages, returns);
    batch.states = states;
    batch.actions = actions;
    batch.log_probs = log_probs;
    batch.values = values;
    batch.dones = dones;

    // Create mini-batches of size 3
    let mini_batches = batch.create_mini_batches(3);

    assert_eq!(mini_batches.len(), 3); // 7 steps / 3 = 3 batches (3, 3, 1)
    assert_eq!(mini_batches[0].states.len(), 3);
    assert_eq!(mini_batches[1].states.len(), 3);
    assert_eq!(mini_batches[2].states.len(), 1); // Remainder
}

#[test]
fn test_continuous_trajectory_batch() {
    // Test continuous trajectory batching
    let action1 = ContinuousAction::new(0.3);
    let action2 = ContinuousAction::new(0.7);
    let action3 = ContinuousAction::new(0.5);

    let step1 = ContinuousTrajectoryStep::new(vec![1.0; 4], action1, -1.0, 10.0, 5.0, false);
    let step2 = ContinuousTrajectoryStep::new(vec![2.0; 4], action2, -0.8, 15.0, 7.0, false);
    let step3 = ContinuousTrajectoryStep::new(vec![3.0; 4], action3, -1.2, 20.0, 10.0, true);

    let mut trajectory = ContinuousTrajectory::new();
    trajectory.add_step(step1);
    trajectory.add_step(step2);
    trajectory.add_step(step3);

    let trajectories = vec![trajectory];
    let advantages = vec![0.1, 0.2, 0.3];
    let returns = vec![15.0, 22.0, 30.0];

    let batch = ContinuousTrajectoryBatch::from_trajectories(trajectories, advantages, returns);

    assert_eq!(batch.actions.len(), 3);
    assert_eq!(batch.actions[0], 0.3);
    assert_eq!(batch.actions[1], 0.7);
    assert_eq!(batch.actions[2], 0.5);
}

#[test]
fn test_continuous_batch_normalization() {
    // Test advantage normalization for continuous trajectories
    let action = ContinuousAction::new(0.5);
    let step1 = ContinuousTrajectoryStep::new(vec![1.0; 2], action, -1.0, 5.0, 2.0, false);
    let step2 = ContinuousTrajectoryStep::new(vec![2.0; 2], action, -1.0, 5.0, 2.0, false);
    let step3 = ContinuousTrajectoryStep::new(vec![3.0; 2], action, -1.0, 5.0, 2.0, true);

    let mut trajectory = ContinuousTrajectory::new();
    trajectory.add_step(step1);
    trajectory.add_step(step2);
    trajectory.add_step(step3);

    let trajectories = vec![trajectory];
    let advantages = vec![1.0, 3.0, 5.0];
    let returns = vec![0.0; 3];

    let mut batch = ContinuousTrajectoryBatch::from_trajectories(trajectories, advantages, returns);

    batch.normalize_advantages().expect("Normalization failed");

    // Check zero mean
    let mean: f32 = batch.advantages.iter().sum::<f32>() / batch.advantages.len() as f32;
    assert!(mean.abs() < 1e-6, "Mean not zero: {}", mean);

    // Check unit variance
    let variance: f32 = batch.advantages.iter().map(|&a| a * a).sum::<f32>() / batch.advantages.len() as f32;
    assert!((variance - 1.0).abs() < 1e-5, "Variance not unit: {}", variance);
}

#[test]
fn test_continuous_mini_batch_creation() {
    // Test mini-batch creation for continuous actions
    let action = ContinuousAction::new(0.5);
    let mut trajectory = ContinuousTrajectory::new();

    for i in 0..8 {
        trajectory.add_step(ContinuousTrajectoryStep::new(
            vec![(i as f32) * 0.1; 3],
            action,
            -1.0,
            5.0,
            1.0,
            i == 7,
        ));
    }

    let trajectories = vec![trajectory];
    let advantages = vec![0.0; 8];
    let returns = vec![0.0; 8];

    let batch = ContinuousTrajectoryBatch::from_trajectories(trajectories, advantages, returns);

    // Create mini-batches of size 3
    let mini_batches = batch.create_mini_batches(3);

    assert_eq!(mini_batches.len(), 3); // 8 steps / 3 = 3 batches (3, 3, 2)
    assert_eq!(mini_batches[0].states.len(), 3);
    assert_eq!(mini_batches[1].states.len(), 3);
    assert_eq!(mini_batches[2].states.len(), 2); // Remainder
}

#[test]
fn test_trajectory_completeness() {
    // Test trajectory completeness detection
    let mut trajectory = Trajectory::new();

    assert!(!trajectory.is_complete()); // Empty trajectory not complete

    trajectory.add_step(TrajectoryStep::new(
        vec![1.0],
        TradingAction::Buy,
        0.0,
        0.0,
        1.0,
        false,
    ));

    assert!(!trajectory.is_complete()); // Not done yet

    trajectory.add_step(TrajectoryStep::new(
        vec![2.0],
        TradingAction::Sell,
        0.0,
        0.0,
        2.0,
        true,
    ));

    assert!(trajectory.is_complete()); // Now done
}