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)

#![allow(unused_crate_dependencies)]

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

mod real_data_helpers;
use real_data_helpers::{load_dqn_states, real_data_available};

// ============================================================================
// 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
}

// ============================================================================
// Real Market Data Tests
// ============================================================================

/// Test PPO trajectory creation with real market data
#[test]
fn test_ppo_trajectory_real_market_data() {
    // Use tokio runtime for async data loading
    let rt = tokio::runtime::Runtime::new().expect("Failed to create runtime");
    let states = rt
        .block_on(async { load_dqn_states(50, 4).await })
        .expect("Failed to load states");

    // Skip if no real data available
    if states.is_empty() || states.len() < 10 {
        eprintln!("Skipping test: insufficient real data");
        return;
    }

    // Create trajectory from real market states
    let mut trajectory = Trajectory::new();

    for i in 0..states.len() - 1 {
        let state = &states[i];
        let next_state = &states[i + 1];

        // Infer reward from price change (first feature is price)
        let price_change = next_state[0] - state[0];
        let reward = price_change.signum();

        trajectory.add_step(TrajectoryStep::new(
            state.clone(),
            TradingAction::Buy,
            -0.5,
            5.0,
            reward,
            i == states.len() - 2,
        ));
    }

    assert!(
        trajectory.steps.len() > 0,
        "Should have created trajectory from real data"
    );
    assert!(trajectory.is_complete(), "Trajectory should be complete");
}

/// Test PPO GAE computation with real market data
#[test]
fn test_ppo_gae_real_market_data() {
    let rt = tokio::runtime::Runtime::new().expect("Failed to create runtime");
    let states = rt
        .block_on(async { load_dqn_states(50, 4).await })
        .expect("Failed to load states");

    if states.is_empty() || states.len() < 20 {
        eprintln!("Skipping test: insufficient real data");
        return;
    }

    // Create trajectory from real market states
    let mut trajectory = Trajectory::new();

    for i in 0..states.len() - 1 {
        let state = &states[i];
        let next_state = &states[i + 1];
        let price_change = next_state[0] - state[0];
        let reward = price_change.signum();

        trajectory.add_step(TrajectoryStep::new(
            state.clone(),
            TradingAction::Buy,
            -0.5,
            5.0,
            reward,
            i == states.len() - 2,
        ));
    }

    let trajectories = vec![trajectory];

    // Compute advantages using GAE
    let gae_method = AdvantageMethod::GAE(GAEConfig {
        gamma: 0.99,
        lambda: 0.95,
        normalize_advantages: true,
    });

    let (advantages, returns) =
        compute_advantages(&trajectories, &gae_method).expect("GAE computation failed");

    // Verify computed advantages and returns
    assert_eq!(advantages.len(), states.len() - 1);
    assert_eq!(returns.len(), states.len() - 1);

    // All advantages and returns should be finite
    for adv in &advantages {
        assert!(adv.is_finite(), "Advantage should be finite with real data");
    }

    for ret in &returns {
        assert!(ret.is_finite(), "Return should be finite with real data");
    }
}

/// Test PPO training with real market data trajectories
#[test]
fn test_ppo_training_real_market_data() {
    let rt = tokio::runtime::Runtime::new().expect("Failed to create runtime");
    let states = rt
        .block_on(async { load_dqn_states(100, 4).await })
        .expect("Failed to load states");

    if states.is_empty() || states.len() < 50 {
        eprintln!("Skipping test: insufficient real data");
        return;
    }

    let config = PPOConfig {
        state_dim: 4,
        num_actions: 3,
        policy_hidden_dims: vec![16, 16],
        value_hidden_dims: vec![16, 16],
        clip_epsilon: 0.2,
        policy_learning_rate: 0.001,
        value_learning_rate: 0.001,
        gae_config: GAEConfig {
            gamma: 0.99,
            lambda: 0.95,
            normalize_advantages: true,
        },
        ..PPOConfig::default()
    };

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

    // Create trajectory from real market states
    let mut trajectory = Trajectory::new();

    for i in 0..states.len() - 1 {
        let state = &states[i];
        let next_state = &states[i + 1];
        let price_change = next_state[0] - state[0];
        let reward = price_change;

        trajectory.add_step(TrajectoryStep::new(
            state.clone(),
            TradingAction::Buy,
            -0.5,
            5.0,
            reward,
            i == states.len() - 2,
        ));
    }

    let trajectories = vec![trajectory];

    // Compute advantages and returns using GAE
    let gae_method = AdvantageMethod::GAE(GAEConfig {
        gamma: 0.99,
        lambda: 0.95,
        normalize_advantages: true,
    });
    let (advantages, returns) =
        compute_advantages(&trajectories, &gae_method).expect("Failed to compute advantages");

    // Train PPO on real market trajectories using TrajectoryBatch
    let mut batch = TrajectoryBatch::from_trajectories(trajectories, advantages, returns);
    let result = ppo.update(&mut batch);
    assert!(
        result.is_ok(),
        "Training should succeed with real market data"
    );

    let (policy_loss, value_loss) = result.unwrap();
    assert!(
        policy_loss.is_finite(),
        "Policy loss should be finite with real data"
    );
    assert!(
        value_loss.is_finite(),
        "Value loss should be finite with real data"
    );
}

/// Test continuous PPO with real market data
#[test]
fn test_continuous_ppo_real_market_data() {
    let rt = tokio::runtime::Runtime::new().expect("Failed to create runtime");
    let states = rt
        .block_on(async { load_dqn_states(50, 4).await })
        .expect("Failed to load states");

    if states.is_empty() || states.len() < 20 {
        eprintln!("Skipping test: insufficient real data");
        return;
    }

    // Create continuous trajectory from real market states
    let mut trajectory = ContinuousTrajectory::new();

    for i in 0..states.len() - 1 {
        let state = &states[i];
        let next_state = &states[i + 1];
        let price_change = next_state[0] - state[0];
        let reward = price_change;

        // Normalize price change to action range [-1, 1]
        let action_value = price_change.clamp(-1.0, 1.0);

        trajectory.add_step(ContinuousTrajectoryStep::new(
            state.clone(),
            ContinuousAction::new(action_value),
            -0.5,
            5.0,
            reward,
            i == states.len() - 2,
        ));
    }

    let trajectories = vec![trajectory];
    let advantages = vec![0.1; states.len() - 1];
    let returns = vec![5.0; states.len() - 1];

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

    assert_eq!(batch.actions.len(), states.len() - 1);
    assert!(
        batch.actions.iter().all(|&a| a.is_finite()),
        "All actions should be finite with real data"
    );
}