//! PPO GAE (Generalized Advantage Estimation) Tests
//!
//! Tests for PPO implementation with GAE advantage computation.
#![allow(unused_crate_dependencies)]

use candle_core::Device;
use ml::ppo::gae::compute_gae_single_trajectory;
use ml::ppo::{GAEConfig, PPOConfig, ValueNetwork, WorkingPPO};

#[test]
fn test_ppo_config_creation() {
    let config = PPOConfig {
        state_dim: 84 * 84 * 4,
        num_actions: 6,
        policy_hidden_dims: vec![512, 256],
        value_hidden_dims: vec![512, 256],
        policy_learning_rate: 3e-4,
        value_learning_rate: 3e-4,
        clip_epsilon: 0.2,
        value_loss_coeff: 0.5,
        entropy_coeff: 0.01,
        gae_config: GAEConfig {
            gamma: 0.99,
            lambda: 0.95,
            normalize_advantages: true,
        },
        batch_size: 2048,
        mini_batch_size: 64,
        num_epochs: 10,
        max_grad_norm: 0.5,
    };

    assert_eq!(config.state_dim, 84 * 84 * 4);
    assert_eq!(config.num_actions, 6);
    assert_eq!(config.clip_epsilon, 0.2);
    assert_eq!(config.gae_config.lambda, 0.95);
}

#[test]
fn test_ppo_default_config() {
    let config = PPOConfig::default();
    assert_eq!(config.state_dim, 64);
    assert_eq!(config.num_actions, 3);
    assert!(config.clip_epsilon > 0.0);
    assert!(config.gae_config.gamma > 0.0);
    assert!(config.gae_config.gamma < 1.0);
}

#[test]
fn test_gae_config_creation() {
    let config = GAEConfig {
        gamma: 0.99,
        lambda: 0.95,
        normalize_advantages: true,
    };

    assert_eq!(config.gamma, 0.99);
    assert_eq!(config.lambda, 0.95);
    assert!(config.normalize_advantages);
}

#[test]
fn test_gae_default_config() {
    let config = GAEConfig::default();
    assert!(config.gamma > 0.0);
    assert!(config.gamma < 1.0);
    assert!(config.lambda > 0.0);
    assert!(config.lambda < 1.0);
}

#[test]
fn test_gae_single_trajectory_simple() {
    let config = GAEConfig {
        gamma: 0.99,
        lambda: 0.95,
        normalize_advantages: false,
    };

    let rewards = vec![1.0, 2.0, 3.0];
    let values = vec![1.0, 2.0, 3.0];
    let dones = vec![false, false, false];
    let next_value = 4.0;

    let result = compute_gae_single_trajectory(&rewards, &values, &dones, next_value, &config);
    assert!(result.is_ok());

    let (advantages, returns) = result.unwrap();
    assert_eq!(advantages.len(), 3);
    assert_eq!(returns.len(), 3);

    // All values should be finite
    for &adv in &advantages {
        assert!(adv.is_finite());
    }
    for &ret in &returns {
        assert!(ret.is_finite());
    }
}

#[test]
fn test_gae_with_terminal_state() {
    let config = GAEConfig {
        gamma: 0.99,
        lambda: 0.95,
        normalize_advantages: false,
    };

    let rewards = vec![1.0, 2.0, 5.0];
    let values = vec![1.5, 2.5, 3.0];
    let dones = vec![false, false, true];
    let next_value = 0.0;

    let result = compute_gae_single_trajectory(&rewards, &values, &dones, next_value, &config);
    assert!(result.is_ok());

    let (advantages, returns) = result.unwrap();
    assert_eq!(advantages.len(), 3);
    assert_eq!(returns.len(), 3);

    // Final return should be close to final reward since episode terminated
    assert!((returns[2] - rewards[2]).abs() < 0.1);
}

#[test]
fn test_gae_mismatched_lengths() {
    let config = GAEConfig::default();

    let rewards = vec![1.0, 2.0, 3.0];
    let values = vec![1.0, 2.0]; // Wrong length
    let dones = vec![false, false, false];
    let next_value = 4.0;

    let result = compute_gae_single_trajectory(&rewards, &values, &dones, next_value, &config);
    assert!(result.is_err());
}

#[test]
fn test_gae_empty_trajectory() {
    let config = GAEConfig::default();

    let rewards: Vec<f32> = vec![];
    let values: Vec<f32> = vec![];
    let dones: Vec<bool> = vec![];
    let next_value = 0.0;

    let result = compute_gae_single_trajectory(&rewards, &values, &dones, next_value, &config);
    assert!(result.is_ok());

    let (advantages, returns) = result.unwrap();
    assert_eq!(advantages.len(), 0);
    assert_eq!(returns.len(), 0);
}

#[test]
fn test_gae_increasing_rewards() {
    let config = GAEConfig {
        gamma: 0.9,
        lambda: 0.8,
        normalize_advantages: false,
    };

    let rewards = vec![1.0, 2.0, 3.0, 4.0, 5.0];
    let values = vec![0.8, 1.8, 2.8, 3.8, 4.8];
    let dones = vec![false; 5];
    let next_value = 5.8;

    let result = compute_gae_single_trajectory(&rewards, &values, &dones, next_value, &config);
    assert!(result.is_ok());

    let (advantages, returns) = result.unwrap();

    // All values should be finite and positive (rewards are increasing)
    for &adv in &advantages {
        assert!(adv.is_finite());
    }
    for (ret, reward) in returns.into_iter().zip(rewards.into_iter()) {
        assert!(ret.is_finite());
        assert!(ret >= reward); // Returns should be at least as large as immediate reward
    }
}

#[test]
fn test_value_network_creation() {
    let device = Device::Cpu;
    let result = ValueNetwork::new(64, &[128, 64], device.clone());
    assert!(result.is_ok());
}

#[test]
fn test_working_ppo_creation() {
    let config = PPOConfig {
        state_dim: 32,
        num_actions: 4,
        policy_hidden_dims: vec![64, 32],
        value_hidden_dims: vec![64, 32],
        policy_learning_rate: 3e-4,
        value_learning_rate: 3e-4,
        clip_epsilon: 0.2,
        value_loss_coeff: 0.5,
        entropy_coeff: 0.01,
        gae_config: GAEConfig::default(),
        batch_size: 256,
        mini_batch_size: 32,
        num_epochs: 4,
        max_grad_norm: 0.5,
    };

    let result = WorkingPPO::new(config);
    assert!(result.is_ok());
}

#[test]
fn test_gae_different_gamma_values() {
    let rewards = vec![1.0; 5];
    let values = vec![0.9; 5];
    let dones = vec![false; 5];
    let next_value = 0.9;

    // Test with different gamma values
    for gamma in [0.9, 0.95, 0.99] {
        let config = GAEConfig {
            gamma,
            lambda: 0.95,
            normalize_advantages: false,
        };

        let result = compute_gae_single_trajectory(&rewards, &values, &dones, next_value, &config);
        assert!(result.is_ok(), "Failed with gamma={}", gamma);

        let (advantages, returns) = result.unwrap();
        assert_eq!(advantages.len(), 5);
        assert_eq!(returns.len(), 5);

        for &adv in &advantages {
            assert!(adv.is_finite());
        }
    }
}

#[test]
fn test_gae_different_lambda_values() {
    let rewards = vec![1.0; 5];
    let values = vec![0.9; 5];
    let dones = vec![false; 5];
    let next_value = 0.9;

    // Test with different lambda values
    for lambda in [0.8, 0.9, 0.95, 0.99] {
        let config = GAEConfig {
            gamma: 0.99,
            lambda,
            normalize_advantages: false,
        };

        let result = compute_gae_single_trajectory(&rewards, &values, &dones, next_value, &config);
        assert!(result.is_ok(), "Failed with lambda={}", lambda);

        let (advantages, _returns) = result.unwrap();
        assert_eq!(advantages.len(), 5);

        for &adv in &advantages {
            assert!(adv.is_finite());
        }
    }
}

#[test]
fn test_ppo_config_validation() {
    let config = PPOConfig {
        state_dim: 16,
        num_actions: 2,
        policy_hidden_dims: vec![32, 16],
        value_hidden_dims: vec![32, 16],
        policy_learning_rate: 1e-3,
        value_learning_rate: 1e-3,
        clip_epsilon: 0.15,
        value_loss_coeff: 0.5,
        entropy_coeff: 0.01,
        gae_config: GAEConfig {
            gamma: 0.95,
            lambda: 0.9,
            normalize_advantages: true,
        },
        batch_size: 128,
        mini_batch_size: 16,
        num_epochs: 4,
        max_grad_norm: 1.0,
    };

    // Validate config properties
    assert!(config.state_dim > 0);
    assert!(config.num_actions > 0);
    assert!(config.policy_learning_rate > 0.0);
    assert!(config.value_learning_rate > 0.0);
    assert!(config.clip_epsilon > 0.0);
    assert!(config.clip_epsilon < 1.0);
    assert!(config.batch_size >= config.mini_batch_size);
    assert!(config.mini_batch_size > 0);
    assert!(config.num_epochs > 0);
    assert!(config.max_grad_norm > 0.0);
}

#[test]
fn test_gae_multiple_episodes() {
    let config = GAEConfig::default();

    // Test multiple short episodes
    for episode_length in [3, 5, 10, 20] {
        let rewards = vec![1.0; episode_length];
        let values = vec![0.9; episode_length];
        let dones = vec![false; episode_length];
        let next_value = 0.9;

        let result = compute_gae_single_trajectory(&rewards, &values, &dones, next_value, &config);
        assert!(
            result.is_ok(),
            "Failed with episode_length={}",
            episode_length
        );

        let (advantages, returns) = result.unwrap();
        assert_eq!(advantages.len(), episode_length);
        assert_eq!(returns.len(), episode_length);
    }
}

#[test]
fn test_gae_negative_rewards() {
    let config = GAEConfig {
        gamma: 0.99,
        lambda: 0.95,
        normalize_advantages: false,
    };

    let rewards = vec![-1.0, -2.0, -3.0];
    let values = vec![0.0, 0.0, 0.0];
    let dones = vec![false, false, true];
    let next_value = 0.0;

    let result = compute_gae_single_trajectory(&rewards, &values, &dones, next_value, &config);
    assert!(result.is_ok());

    let (advantages, returns) = result.unwrap();

    // All values should be finite (even with negative rewards)
    for &adv in &advantages {
        assert!(adv.is_finite());
    }
    for &ret in &returns {
        assert!(ret.is_finite());
    }
}

#[test]
fn test_gae_zero_gamma() {
    let config = GAEConfig {
        gamma: 0.0,
        lambda: 0.95,
        normalize_advantages: false,
    };

    let rewards = vec![1.0, 2.0, 3.0];
    let values = vec![0.5, 1.5, 2.5];
    let dones = vec![false, false, false];
    let next_value = 3.5;

    let result = compute_gae_single_trajectory(&rewards, &values, &dones, next_value, &config);
    assert!(result.is_ok());

    let (_advantages, returns) = result.unwrap();

    // With gamma=0, returns should equal rewards
    for (ret, reward) in returns.into_iter().zip(rewards.into_iter()) {
        assert!((ret - reward).abs() < 1e-6);
    }
}