foxhunt/ml/tests/ppo_hyperopt_validation_split_test.rs

//! PPO Hyperopt Validation Split Tests
//!
//! This test suite verifies that the PPO hyperparameter optimization adapter
//! properly splits data into training and validation sets, preventing overfitting.
//!
//! Critical Requirements:
//! 1. Train/val split must be implemented (80/20)
//! 2. Training must ONLY use train trajectories
//! 3. Validation loss must be computed on held-out val trajectories
//! 4. Optimization metric must be val_loss (not train_loss)
//! 5. Train and val losses should differ (proves separation)

use ml::hyperopt::adapters::ppo::{PPOParams, PPOTrainer};
use ml::hyperopt::traits::HyperparameterOptimizable;

#[test]
fn test_ppo_train_val_separation() {
    // Test that training and validation losses are different,
    // proving that we have separate train/val sets

    let mut trainer = PPOTrainer::new(100).expect("Failed to create PPO trainer");
    let params = PPOParams::default();

    let metrics = trainer.train_with_params(params).expect("Training failed");

    // Verify both train and val metrics exist
    assert!(
        metrics.policy_loss.is_finite(),
        "Training policy loss should be finite"
    );
    assert!(
        metrics.value_loss.is_finite(),
        "Training value loss should be finite"
    );
    assert!(
        metrics.val_policy_loss.is_finite(),
        "Validation policy loss should be finite"
    );
    assert!(
        metrics.val_value_loss.is_finite(),
        "Validation value loss should be finite"
    );

    // Train and val losses should differ (proves separation)
    // With 80/20 split, train loss is computed on 80 trajectories,
    // val loss on 20 trajectories - they will be different
    let policy_diff = (metrics.policy_loss - metrics.val_policy_loss).abs();
    let value_diff = (metrics.value_loss - metrics.val_value_loss).abs();

    println!("Train policy loss: {:.6}", metrics.policy_loss);
    println!("Val policy loss: {:.6}", metrics.val_policy_loss);
    println!("Policy loss difference: {:.6}", policy_diff);
    println!("Train value loss: {:.6}", metrics.value_loss);
    println!("Val value loss: {:.6}", metrics.val_value_loss);
    println!("Value loss difference: {:.6}", value_diff);

    // Losses should be different (not identical) due to different data
    // Allow small differences in case of numerical coincidence
    assert!(
        policy_diff > 1e-6 || value_diff > 1e-6,
        "Train and val losses are identical - no separation! \
         policy_diff={:.6}, value_diff={:.6}",
        policy_diff,
        value_diff
    );
}

#[test]
fn test_ppo_insufficient_trajectories() {
    // Test that training fails gracefully with too few trajectories
    // for 80/20 split (minimum 10 required)

    let mut trainer = PPOTrainer::new(5).expect("Failed to create PPO trainer");
    let params = PPOParams::default();

    let result = trainer.train_with_params(params);

    assert!(
        result.is_err(),
        "Training should fail with insufficient trajectories"
    );

    let error = result.unwrap_err();
    let error_msg = format!("{:?}", error);
    assert!(
        error_msg.contains("Insufficient trajectories") || error_msg.contains("train/val split"),
        "Error should mention insufficient trajectories, got: {}",
        error_msg
    );
}

#[test]
fn test_ppo_edge_case_trajectories() {
    // Test edge case: exactly 10 trajectories (minimum)
    // 80/20 split: 8 train, 2 val

    let mut trainer = PPOTrainer::new(10).expect("Failed to create PPO trainer");
    let params = PPOParams::default();

    let result = trainer.train_with_params(params);

    // Should succeed (10 trajectories is minimum)
    assert!(
        result.is_ok(),
        "Training should succeed with 10 trajectories (minimum)"
    );

    let metrics = result.unwrap();
    assert!(metrics.val_policy_loss.is_finite());
    assert!(metrics.val_value_loss.is_finite());
}

#[test]
fn test_ppo_optimization_uses_val_loss() {
    // Test that extract_objective returns validation loss, not training loss

    let mut trainer = PPOTrainer::new(100).expect("Failed to create PPO trainer");
    let params = PPOParams::default();

    let metrics = trainer.train_with_params(params).expect("Training failed");

    // Get optimization objective
    let objective = PPOTrainer::extract_objective(&metrics);

    // Objective should be based on validation losses
    // extract_objective returns val_policy_loss + val_value_loss
    let expected_objective = metrics.val_policy_loss + metrics.val_value_loss;

    println!("Optimization objective: {:.6}", objective);
    println!("Expected (val losses): {:.6}", expected_objective);
    println!(
        "Train losses sum: {:.6}",
        metrics.policy_loss + metrics.value_loss
    );

    // Verify objective matches validation losses
    assert!(
        (objective - expected_objective).abs() < 1e-6,
        "Optimization objective should be based on validation losses, not training. \
         Got {:.6}, expected {:.6}",
        objective,
        expected_objective
    );

    // Verify objective is NOT equal to training losses
    let train_objective = metrics.policy_loss + metrics.value_loss;
    assert!(
        (objective - train_objective).abs() > 1e-6,
        "Optimization objective should NOT be based on training losses. \
         Objective={:.6}, train_sum={:.6}",
        objective,
        train_objective
    );
}

#[test]
fn test_ppo_val_loss_metrics_exist() {
    // Verify that PPOMetrics struct has val_policy_loss and val_value_loss fields

    let mut trainer = PPOTrainer::new(100).expect("Failed to create PPO trainer");
    let params = PPOParams::default();

    let metrics = trainer.train_with_params(params).expect("Training failed");

    // Access fields to verify they exist (compile-time check)
    let _policy = metrics.policy_loss;
    let _value = metrics.value_loss;
    let _val_policy = metrics.val_policy_loss;
    let _val_value = metrics.val_value_loss;
    let _combined = metrics.combined_loss;
    let _reward = metrics.avg_episode_reward;
    let _episodes = metrics.episodes_completed;

    println!("All required metrics fields exist:");
    println!("  policy_loss: {:.6}", metrics.policy_loss);
    println!("  value_loss: {:.6}", metrics.value_loss);
    println!("  val_policy_loss: {:.6}", metrics.val_policy_loss);
    println!("  val_value_loss: {:.6}", metrics.val_value_loss);
}

#[test]
fn test_ppo_small_val_set_warning() {
    // Test that training succeeds but may warn with small validation set
    // 12 trajectories: 9 train, 3 val (below 5 val threshold)

    let mut trainer = PPOTrainer::new(12).expect("Failed to create PPO trainer");
    let params = PPOParams::default();

    let result = trainer.train_with_params(params);

    // Should succeed (validation set exists, even if small)
    assert!(
        result.is_ok(),
        "Training should succeed with small validation set"
    );

    let metrics = result.unwrap();
    assert!(metrics.val_policy_loss.is_finite());
    assert!(metrics.val_value_loss.is_finite());
}

#[test]
fn test_ppo_hyperopt_prevents_overfitting() {
    // Integration test: Verify that using validation loss for optimization
    // prevents selecting overfitted hyperparameters

    let mut trainer = PPOTrainer::new(100).expect("Failed to create PPO trainer");

    // Test with two different parameter sets
    let params1 = PPOParams {
        policy_learning_rate: 1e-4,
        value_learning_rate: 3e-4,
        clip_epsilon: 0.2,
        value_loss_coeff: 1.0,
        entropy_coeff: 0.01,
    };

    let params2 = PPOParams {
        policy_learning_rate: 5e-5,
        value_learning_rate: 1e-4,
        clip_epsilon: 0.25,
        value_loss_coeff: 0.8,
        entropy_coeff: 0.05,
    };

    let metrics1 = trainer
        .train_with_params(params1.clone())
        .expect("Training 1 failed");

    let metrics2 = trainer
        .train_with_params(params2.clone())
        .expect("Training 2 failed");

    let obj1 = PPOTrainer::extract_objective(&metrics1);
    let obj2 = PPOTrainer::extract_objective(&metrics2);

    println!("\nParams 1:");
    println!(
        "  Train loss: {:.6}",
        metrics1.policy_loss + metrics1.value_loss
    );
    println!("  Val loss: {:.6}", obj1);

    println!("\nParams 2:");
    println!(
        "  Train loss: {:.6}",
        metrics2.policy_loss + metrics2.value_loss
    );
    println!("  Val loss: {:.6}", obj2);

    // Both should produce valid validation losses
    assert!(obj1.is_finite() && obj2.is_finite());

    // Verify objectives are based on validation, not training
    let train_obj1 = metrics1.policy_loss + metrics1.value_loss;
    let train_obj2 = metrics2.policy_loss + metrics2.value_loss;

    assert!(
        (obj1 - train_obj1).abs() > 1e-6 || (obj2 - train_obj2).abs() > 1e-6,
        "At least one objective should differ from training loss"
    );
}