foxhunt/ml/tests/test_ppo_checkpoint_loading.rs

//! PPO Checkpoint Loading Production Validation Test
//!
//! Tests WorkingPPO::load_checkpoint() with real trained checkpoints:
//! - Checkpoint existence validation
//! - Actor/critic weight loading
//! - Inference capability
//! - Comparison with random initialization
//!
//! **Agent 170 Mission**: Validate checkpoint loading works with real models

use candle_core::Device;
use ml::ppo::gae::GAEConfig;
use ml::ppo::ppo::{PPOConfig, WorkingPPO};
use std::path::Path;

#[test]
fn test_ppo_checkpoint_existence() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n=== PPO CHECKPOINT EXISTENCE VALIDATION ===\n");

    let checkpoints = vec![
        (
            "ml/trained_models/production/ppo/ppo_actor_epoch_130.safetensors",
            "ml/trained_models/production/ppo/ppo_critic_epoch_130.safetensors",
            130,
        ),
        (
            "ml/trained_models/production/ppo/ppo_actor_epoch_420.safetensors",
            "ml/trained_models/production/ppo/ppo_critic_epoch_420.safetensors",
            420,
        ),
    ];

    for (actor_path, critic_path, epoch) in checkpoints {
        println!("Checking epoch {} checkpoints:", epoch);

        let actor_exists = Path::new(actor_path).exists();
        let critic_exists = Path::new(critic_path).exists();

        println!(
            "  Actor:  {} ({})",
            actor_path,
            if actor_exists { "EXISTS" } else { "MISSING" }
        );
        println!(
            "  Critic: {} ({})",
            critic_path,
            if critic_exists { "EXISTS" } else { "MISSING" }
        );

        // Gracefully skip if checkpoints are missing (CI environment)
        if !actor_exists || !critic_exists {
            println!("SKIP: Checkpoint pair for epoch {} not found (expected in production environment only)", epoch);
            println!("  This is normal in CI/test environments without trained models\n");
            continue;
        }

        println!("  ✓ Both checkpoints found");

        // Check file sizes
        if actor_exists {
            let metadata = std::fs::metadata(actor_path).unwrap();
            println!("  Actor size: {} bytes", metadata.len());
            assert!(metadata.len() > 0, "Actor checkpoint is empty");
        }

        if critic_exists {
            let metadata = std::fs::metadata(critic_path).unwrap();
            println!("  Critic size: {} bytes", metadata.len());
            assert!(metadata.len() > 0, "Critic checkpoint is empty");
        }

        println!("  ✓ Checkpoint pair validated\n");
    }

    Ok(())
}

#[test]
fn test_ppo_checkpoint_loading_epoch_130() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n=== PPO CHECKPOINT LOADING TEST (EPOCH 130) ===\n");

    // Check if checkpoints exist first
    let actor_path = "ml/trained_models/production/ppo/ppo_actor_epoch_130.safetensors";
    let critic_path = "ml/trained_models/production/ppo/ppo_critic_epoch_130.safetensors";

    if !Path::new(actor_path).exists() || !Path::new(critic_path).exists() {
        println!("SKIP: Checkpoint files not found (expected in production environment only)");
        println!("  Actor:  {}", actor_path);
        println!("  Critic: {}", critic_path);
        println!("  This is normal in CI/test environments without trained models\n");
        return Ok(());
    }

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    println!("Using device: {:?}", device);

    // Create PPO config matching training configuration
    let config = PPOConfig {
        state_dim: 16,
        num_actions: 3,
        policy_hidden_dims: vec![128, 64],
        value_hidden_dims: vec![128, 64],
        policy_learning_rate: 3e-4,
        value_learning_rate: 1e-3,
        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,
        },
        num_epochs: 10,
        batch_size: 64,
        mini_batch_size: 32,
        max_grad_norm: 0.5,
    };

    println!("Loading checkpoint...");
    let ppo = WorkingPPO::load_checkpoint(
        actor_path,
        critic_path,
        config.clone(),
        device.clone(),
    )
    .expect("Failed to load PPO checkpoint");

    println!("✓ Checkpoint loaded successfully\n");

    // Test inference with random state
    println!("Testing inference capability...");
    let test_state = vec![
        0.5, -0.3, 1.2, 0.0, -0.5, 0.8, -1.0, 0.3, 0.1, 0.7, -0.2, 0.4, -0.6, 0.9, 0.2, -0.1,
    ];

    let action_probs = ppo.predict(&test_state).expect("Inference failed");
    println!("Action probabilities: {:?}", action_probs);

    // Validate output
    assert_eq!(action_probs.len(), 3, "Should have 3 action probabilities");

    let sum: f32 = action_probs.iter().sum();
    println!("Probability sum: {:.6}", sum);
    assert!(
        (sum - 1.0).abs() < 1e-4,
        "Action probabilities should sum to ~1.0"
    );

    // All probabilities should be valid
    for (i, &prob) in action_probs.iter().enumerate() {
        assert!(
            prob >= 0.0 && prob <= 1.0,
            "Invalid probability at index {}: {}",
            i,
            prob
        );
    }

    println!("✓ Inference validated\n");
    Ok(())
}

#[test]
fn test_ppo_checkpoint_loading_epoch_420() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n=== PPO CHECKPOINT LOADING TEST (EPOCH 420) ===\n");
    
    // Check if checkpoints exist first
    let actor_path = "ml/trained_models/production/ppo/ppo_actor_epoch_420.safetensors";
    let critic_path = "ml/trained_models/production/ppo/ppo_critic_epoch_420.safetensors";

    if !Path::new(actor_path).exists() || !Path::new(critic_path).exists() {
        println!("SKIP: Checkpoint files not found (expected in production environment only)");
        println!("  Actor:  {}", actor_path);
        println!("  Critic: {}", critic_path);
        println!("  This is normal in CI/test environments without trained models\n");
        return Ok(());
    }

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    println!("Using device: {:?}", device);

    let config = PPOConfig {
        state_dim: 16,
        num_actions: 3,
        policy_hidden_dims: vec![128, 64],
        value_hidden_dims: vec![128, 64],
        policy_learning_rate: 3e-4,
        value_learning_rate: 1e-3,
        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,
        },
        num_epochs: 10,
        batch_size: 64,
        mini_batch_size: 32,
        max_grad_norm: 0.5,
    };

    println!("Loading checkpoint...");
    let ppo = WorkingPPO::load_checkpoint(
        actor_path,
        critic_path,
        config,
        device,
    )
    .expect("Failed to load PPO checkpoint");

    println!("✓ Checkpoint loaded successfully\n");

    // Test inference
    println!("Testing inference capability...");
    let test_state = vec![
        1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
    ];

    let action_probs = ppo.predict(&test_state).expect("Inference failed");
    println!("Action probabilities: {:?}", action_probs);

    assert_eq!(action_probs.len(), 3);
    let sum: f32 = action_probs.iter().sum();
    assert!((sum - 1.0).abs() < 1e-4);

    println!("✓ Inference validated\n");
    Ok(())
}

#[test]
fn test_ppo_loaded_vs_random_initialization() {
    println!("\n=== PPO LOADED VS RANDOM INITIALIZATION ===\n");

    // Check if checkpoints exist first
    let actor_path = "ml/trained_models/production/ppo/ppo_actor_epoch_420.safetensors";
    let critic_path = "ml/trained_models/production/ppo/ppo_critic_epoch_420.safetensors";

    if !Path::new(actor_path).exists() || !Path::new(critic_path).exists() {
        println!("SKIP: Checkpoint files not found (expected in production environment only)");
        println!("  Actor:  {}", actor_path);
        println!("  Critic: {}", critic_path);
        println!("  This is normal in CI/test environments without trained models\n");
        return;
    }

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    println!("Using device: {:?}", device);

    let config = PPOConfig {
        state_dim: 16,
        num_actions: 3,
        policy_hidden_dims: vec![128, 64],
        value_hidden_dims: vec![128, 64],
        policy_learning_rate: 3e-4,
        value_learning_rate: 1e-3,
        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,
        },
        num_epochs: 10,
        batch_size: 64,
        mini_batch_size: 32,
        max_grad_norm: 0.5,
    };

    // Load trained model
    println!("Loading trained checkpoint (epoch 420)...");
    let loaded_ppo = WorkingPPO::load_checkpoint(
        actor_path,
        critic_path,
        config.clone(),
        device.clone(),
    )
    .expect("Failed to load checkpoint");

    // Create random model
    println!("Creating random initialization...");
    let random_ppo = WorkingPPO::with_device(config.clone(), device).expect("Failed to create random PPO");

    // Test with same state
    let test_state = vec![
        0.5, -0.3, 1.2, 0.0, -0.5, 0.8, -1.0, 0.3, 0.1, 0.7, -0.2, 0.4, -0.6, 0.9, 0.2, -0.1,
    ];

    println!("\nTesting inference on same state...");
    let loaded_probs = loaded_ppo
        .predict(&test_state)
        .expect("Loaded inference failed");
    let random_probs = random_ppo
        .predict(&test_state)
        .expect("Random inference failed");

    println!("Loaded model:  {:?}", loaded_probs);
    println!("Random model:  {:?}", random_probs);

    // Compute L2 distance between probability distributions
    let mut l2_distance = 0.0;
    for i in 0..3 {
        let diff = loaded_probs[i] - random_probs[i];
        l2_distance += diff * diff;
    }
    l2_distance = l2_distance.sqrt();

    println!("\nL2 distance between distributions: {:.6}", l2_distance);

    // Loaded model should produce different probabilities than random
    assert!(
        l2_distance > 0.01,
        "Loaded model should differ from random initialization (distance too small: {:.6})",
        l2_distance
    );

    println!("✓ Loaded model differs from random initialization\n");
}

#[test]
fn test_ppo_checkpoint_error_handling() {
    println!("\n=== PPO CHECKPOINT ERROR HANDLING ===\n");

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

    let config = PPOConfig {
        state_dim: 16,
        num_actions: 3,
        policy_hidden_dims: vec![128, 64],
        value_hidden_dims: vec![128, 64],
        policy_learning_rate: 3e-4,
        value_learning_rate: 1e-3,
        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,
        },
        num_epochs: 10,
        batch_size: 64,
        mini_batch_size: 32,
        max_grad_norm: 0.5,
    };

    // Test 1: Missing actor checkpoint
    println!("Test 1: Missing actor checkpoint");
    let result = WorkingPPO::load_checkpoint(
        "nonexistent_actor.safetensors",
        "ml/trained_models/production/ppo/ppo_critic_epoch_130.safetensors",
        config.clone(),
        device.clone(),
    );
    assert!(result.is_err(), "Should fail with missing actor checkpoint");
    println!("  ✓ Correctly rejected missing actor\n");

    // Test 2: Missing critic checkpoint
    println!("Test 2: Missing critic checkpoint");
    let result = WorkingPPO::load_checkpoint(
        "ml/trained_models/production/ppo/ppo_actor_epoch_130.safetensors",
        "nonexistent_critic.safetensors",
        config.clone(),
        device.clone(),
    );
    assert!(
        result.is_err(),
        "Should fail with missing critic checkpoint"
    );
    println!("  ✓ Correctly rejected missing critic\n");

    // Test 3: Both missing
    println!("Test 3: Both checkpoints missing");
    let result = WorkingPPO::load_checkpoint(
        "nonexistent_actor.safetensors",
        "nonexistent_critic.safetensors",
        config,
        device,
    );
    assert!(result.is_err(), "Should fail with both checkpoints missing");
    println!("  ✓ Correctly rejected both missing\n");
}

#[test]
fn test_ppo_checkpoint_batch_inference() {
    println!("\n=== PPO CHECKPOINT BATCH INFERENCE ===\n");

    // Check if checkpoints exist first
    let actor_path = "ml/trained_models/production/ppo/ppo_actor_epoch_420.safetensors";
    let critic_path = "ml/trained_models/production/ppo/ppo_critic_epoch_420.safetensors";

    if !Path::new(actor_path).exists() || !Path::new(critic_path).exists() {
        println!("SKIP: Checkpoint files not found (expected in production environment only)");
        println!("  Actor:  {}", actor_path);
        println!("  Critic: {}", critic_path);
        println!("  This is normal in CI/test environments without trained models\n");
        return;
    }

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    println!("Using device: {:?}", device);

    let config = PPOConfig {
        state_dim: 16,
        num_actions: 3,
        policy_hidden_dims: vec![128, 64],
        value_hidden_dims: vec![128, 64],
        policy_learning_rate: 3e-4,
        value_learning_rate: 1e-3,
        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,
        },
        num_epochs: 10,
        batch_size: 64,
        mini_batch_size: 32,
        max_grad_norm: 0.5,
    };

    println!("Loading checkpoint...");
    let ppo = WorkingPPO::load_checkpoint(
        actor_path,
        critic_path,
        config,
        device,
    )
    .expect("Failed to load checkpoint");

    // Test with multiple diverse states
    let test_states = vec![
        vec![1.0; 16],
        vec![0.0; 16],
        vec![-1.0; 16],
        vec![
            0.5, -0.5, 0.5, -0.5, 0.5, -0.5, 0.5, -0.5, 0.5, -0.5, 0.5, -0.5, 0.5, -0.5, 0.5, -0.5,
        ],
    ];

    println!("\nBatch inference test:");
    for (i, state) in test_states.iter().enumerate() {
        let probs = ppo.predict(state).expect("Inference failed");
        let sum: f32 = probs.iter().sum();

        println!("  State {}: probs={:?}, sum={:.6}", i, probs, sum);

        assert_eq!(probs.len(), 3);
        assert!((sum - 1.0).abs() < 1e-4);
        for prob in &probs {
            assert!(*prob >= 0.0 && *prob <= 1.0);
        }
    }

    println!("\n✓ Batch inference validated\n");
}