foxhunt/ml/tests/hyperopt_early_stopping_termination_test.rs

//! Hyperopt Early Stopping Termination Tests
//!
//! **Objective**: Validate that early stopping ACTUALLY TERMINATES training when gradient collapse
//! is detected, not just logs warnings.
//!
//! **Bug Description**:
//! - Detection works: Gradient collapse warnings are logged correctly
//! - Termination BROKEN: Training continues to completion despite exhausted patience
//! - Expected: Training stops at epoch 3-5 (patience=3)
//! - Actual: Training runs to epoch 20 (max_epochs)
//!
//! **Root Cause Hypothesis**:
//! The error returned by `WorkingDQN::train_step()` at line 1704 is being caught/ignored
//! somewhere in the training loop instead of propagating up to terminate training.
//!
//! **Expected Behavior**:
//! 1. Gradient collapse detected for 3 consecutive epochs
//! 2. `MLError::TrainingError` returned from `train_step()`
//! 3. Error propagates through `train_epoch()` → `train_from_parquet()`
//! 4. Training terminates with error message
//! 5. Final epoch count: 3-5 (NOT 20)
//!
//! **Test Strategy**:
//! - Configure DQN with parameters GUARANTEED to cause gradient collapse:
//!   - Very low learning rate (1e-8) → gradients near zero
//!   - Narrow C51 bounds (v_min=-0.001, v_max=0.001) → value range collapse
//!   - Small batch size (8) → high variance, unstable gradients
//! - Set gradient_collapse_patience = 3 (stop after 3 consecutive epochs)
//! - Run training for max 20 epochs
//! - Assert training stops at epoch 3-5 (NOT 20)
//!
//! **Current Status**: EXPECTED TO FAIL (bug exists)
//! **After Fix**: Should PASS (training terminates correctly)

use ml::dqn::{Experience, WorkingDQN, WorkingDQNConfig};
use ml::MLError;

/// Test 1: Verify that gradient collapse detection actually terminates training
///
/// **Expected**: Training stops at epoch 3-5 with "Gradient collapse detected" error
/// **Actual** (BEFORE FIX): Training runs to epoch 20, warnings logged but ignored
#[test]
fn test_hyperopt_early_stopping_actually_terminates() -> Result<(), MLError> {
    println!("\n🧪 TEST: Hyperopt Early Stopping Termination");
    println!("============================================");
    println!("Objective: Verify training STOPS when gradient collapse detected");
    println!();

    // ========================================
    // STEP 1: Configure DQN for Guaranteed Gradient Collapse
    // ========================================
    println!("📝 STEP 1: Configuring DQN with gradient collapse parameters");

    let mut config = WorkingDQNConfig::conservative();

    // Basic architecture
    config.state_dim = 32;
    config.num_actions = 3;

    // CRITICAL: Parameters guaranteed to cause gradient collapse
    config.learning_rate = 1e-8;  // Extremely low → gradients near zero
    config.batch_size = 8;         // Small → high variance, unstable

    // C51 Distributional RL with narrow bounds (causes value range collapse)
    config.use_distributional = true;
    config.v_min = -0.001;  // Extremely narrow range
    config.v_max = 0.001;   // → distribution collapse → zero gradients
    config.num_atoms = 51;

    // Early stopping configuration
    config.gradient_collapse_patience = 3;      // Stop after 3 consecutive epochs
    config.gradient_collapse_multiplier = 100.0; // Threshold = LR × 100 = 1e-6

    // Replay buffer
    config.replay_buffer_capacity = 500;
    config.min_replay_size = 100;
    config.warmup_steps = 0; // Disable warmup for fast testing

    println!("  ✓ Learning rate: {:.2e} (extremely low → gradient collapse)", config.learning_rate);
    println!("  ✓ Batch size: {} (small → unstable gradients)", config.batch_size);
    println!("  ✓ C51 bounds: [{}, {}] (narrow → value collapse)", config.v_min, config.v_max);
    println!("  ✓ Collapse patience: {} epochs", config.gradient_collapse_patience);
    println!("  ✓ Collapse threshold: {:.2e}", config.learning_rate * config.gradient_collapse_multiplier);
    println!();

    // Save config values before moving
    let threshold = config.learning_rate * config.gradient_collapse_multiplier;

    // ========================================
    // STEP 2: Initialize DQN and Fill Replay Buffer
    // ========================================
    println!("📝 STEP 2: Initializing DQN and filling replay buffer");

    let mut dqn = WorkingDQN::new(config)?;

    // Generate enough synthetic experiences to start training
    let num_experiences = 200;
    for i in 0..num_experiences {
        let state = vec![0.01 * (i as f32); 32];
        let next_state = vec![0.01 * (i as f32 + 1.0); 32];
        let action = (i % 3) as u8;
        let reward = if i % 10 == 0 { 0.1 } else { -0.01 };
        let done = i % 50 == 49;

        let experience = Experience::new(state, action, reward, next_state, done);
        dqn.store_experience(experience)?;
    }

    println!("  ✓ Generated {} experiences", num_experiences);
    println!("  ✓ Replay buffer size: {}", dqn.get_replay_buffer_size()?);
    println!();

    // ========================================
    // STEP 3: Run Training and Monitor for Early Termination
    // ========================================
    println!("📝 STEP 3: Running training (max 20 epochs)");
    println!("Expected: Training stops at epoch 3-5 (patience=3)");
    println!();

    let max_epochs = 20;
    let mut epoch_count = 0;
    let mut gradient_norms = Vec::new();
    let mut losses = Vec::new();
    let mut training_terminated_early = false;
    let mut termination_error_message = String::new();

    for epoch in 0..max_epochs {
        epoch_count = epoch + 1;

        // Run 10 training steps per epoch
        let mut epoch_grad_norm_sum = 0.0;
        let mut epoch_loss_sum = 0.0;
        let steps_per_epoch = 10;

        for step in 0..steps_per_epoch {
            // CRITICAL: This should return an error after patience is exhausted
            let train_result = dqn.train_step(None);

            match train_result {
                Ok((_loss, grad_norm)) => {
                    epoch_grad_norm_sum += grad_norm;
                    epoch_loss_sum += _loss;
                },
                Err(e) => {
                    // SUCCESS: Training terminated with error!
                    training_terminated_early = true;
                    termination_error_message = e.to_string();

                    println!("✅ Training terminated at epoch {}, step {}", epoch_count, step + 1);
                    println!("✅ Error message: {}", termination_error_message);
                    break;
                }
            }
        }

        // If training terminated, break out of epoch loop
        if training_terminated_early {
            break;
        }

        // Calculate epoch averages
        let avg_grad_norm = epoch_grad_norm_sum / steps_per_epoch as f32;
        let avg_loss = epoch_loss_sum / steps_per_epoch as f32;

        gradient_norms.push(avg_grad_norm);
        losses.push(avg_loss);

        println!(
            "  Epoch {}/{}: avg_grad_norm={:.8}, avg_loss={:.6}",
            epoch_count, max_epochs, avg_grad_norm, avg_loss
        );
    }

    println!();

    // ========================================
    // VALIDATION 1: Training Terminated Early
    // ========================================
    println!("✅ VALIDATION 1: Early Termination");
    println!("----------------------------------");

    if training_terminated_early {
        println!("  ✓ Training terminated early at epoch {}", epoch_count);
        println!("  ✓ Error message received");
    } else {
        println!("  ✗ FAILED: Training ran to completion (epoch {})", epoch_count);
        println!("  ✗ Expected: Training should stop at epoch 3-5");
        println!("  ✗ Actual: Training ran to epoch {}", epoch_count);
    }

    assert!(
        training_terminated_early,
        "BUG: Early stopping did NOT terminate training!\n\
         Expected: Training stops at epoch 3-5 (patience=3)\n\
         Actual: Training ran to epoch {}\n\
         This indicates the error is being caught/ignored instead of propagating.",
        epoch_count
    );
    println!();

    // ========================================
    // VALIDATION 2: Stopped at Correct Epoch
    // ========================================
    println!("✅ VALIDATION 2: Termination Timing");
    println!("-----------------------------------");

    // Should stop at epoch 3-5 (patience=3, allow some variance for detection timing)
    let expected_min_epoch = 3;
    let expected_max_epoch = 5;

    println!("  Expected epoch range: {}-{}", expected_min_epoch, expected_max_epoch);
    println!("  Actual epoch: {}", epoch_count);

    assert!(
        epoch_count >= expected_min_epoch && epoch_count <= expected_max_epoch,
        "BUG: Training stopped at wrong epoch!\n\
         Expected: epoch {}-{} (patience=3 + detection lag)\n\
         Actual: epoch {}\n\
         This indicates incorrect patience counting or threshold calculation.",
        expected_min_epoch, expected_max_epoch, epoch_count
    );

    println!("  ✓ Training stopped at epoch {} (within expected range)", epoch_count);
    println!();

    // ========================================
    // VALIDATION 3: Error Message Content
    // ========================================
    println!("✅ VALIDATION 3: Error Message");
    println!("------------------------------");

    println!("  Error: {}", termination_error_message);

    assert!(
        termination_error_message.contains("Gradient collapse detected"),
        "BUG: Error message does not contain 'Gradient collapse detected'!\n\
         Expected: Error message should explain gradient collapse\n\
         Actual: '{}'\n\
         This indicates the wrong error type or incorrect error propagation.",
        termination_error_message
    );

    println!("  ✓ Error message contains 'Gradient collapse detected'");

    // Should also contain helpful suggestions
    let contains_suggestions = termination_error_message.contains("learning rate")
        || termination_error_message.contains("batch size");

    if contains_suggestions {
        println!("  ✓ Error message contains helpful suggestions");
    }
    println!();

    // ========================================
    // VALIDATION 4: Gradient Collapse Observed
    // ========================================
    println!("✅ VALIDATION 4: Gradient Collapse Evidence");
    println!("-------------------------------------------");

    // Gradients should be near zero (below threshold) - use saved value from earlier
    let collapsed_epochs = gradient_norms.iter().filter(|&&g| g < threshold as f32).count();

    println!("  Collapse threshold: {:.8}", threshold);
    println!("  Epochs trained: {}", gradient_norms.len());
    println!("  Epochs with collapsed gradients: {}", collapsed_epochs);

    for (i, &grad_norm) in gradient_norms.iter().enumerate() {
        let status = if grad_norm < threshold as f32 { "✗ COLLAPSED" } else { "✓ OK" };
        println!("    Epoch {}: grad_norm={:.8} {}", i + 1, grad_norm, status);
    }

    // At least 3 consecutive epochs should have collapsed gradients (patience=3)
    assert!(
        collapsed_epochs >= 3,
        "BUG: Insufficient gradient collapse detected!\n\
         Expected: At least 3 collapsed epochs (patience=3)\n\
         Actual: {} collapsed epochs\n\
         This indicates the test configuration didn't trigger collapse correctly.",
        collapsed_epochs
    );

    println!("  ✓ Sufficient gradient collapse detected ({} epochs)", collapsed_epochs);
    println!();

    // ========================================
    // TEST SUMMARY
    // ========================================
    println!("============================================");
    println!("🎉 TEST PASSED: Early Stopping Termination");
    println!("============================================");
    println!("✓ Training terminated early (epoch {})", epoch_count);
    println!("✓ Stopped at correct time ({}-{} epoch range)", expected_min_epoch, expected_max_epoch);
    println!("✓ Error message correct: 'Gradient collapse detected'");
    println!("✓ Gradient collapse observed ({} epochs)", collapsed_epochs);
    println!();

    Ok(())
}

/// Test 2: Verify that the gradient collapse counter increments correctly
///
/// **Expected**: Counter increments on collapse, resets on recovery
/// **Tests**: Internal state management of collapse detection
#[test]
fn test_gradient_collapse_counter_increments() -> Result<(), MLError> {
    println!("\n🧪 TEST: Gradient Collapse Counter Logic");
    println!("=========================================");
    println!("Objective: Verify collapse counter increments and resets correctly");
    println!();

    // ========================================
    // STEP 1: Configure DQN with Gradient Collapse Parameters
    // ========================================
    println!("📝 STEP 1: Configuring DQN");

    let mut config = WorkingDQNConfig::conservative();
    config.state_dim = 32;
    config.num_actions = 3;
    config.learning_rate = 1e-8;  // Low enough to trigger collapse
    config.batch_size = 8;
    config.gradient_collapse_patience = 5;  // Higher patience to test counter logic
    config.gradient_collapse_multiplier = 100.0;
    config.replay_buffer_capacity = 500;
    config.min_replay_size = 100;
    config.warmup_steps = 0;

    // Use C51 with narrow bounds to guarantee collapse
    config.use_distributional = true;
    config.v_min = -0.001;
    config.v_max = 0.001;
    config.num_atoms = 51;

    println!("  ✓ Collapse patience: {} epochs", config.gradient_collapse_patience);
    println!("  ✓ Collapse threshold: {:.8}", config.learning_rate * config.gradient_collapse_multiplier);
    println!();

    // Save config values before moving
    let threshold = config.learning_rate * config.gradient_collapse_multiplier;
    let collapse_patience = config.gradient_collapse_patience;

    // ========================================
    // STEP 2: Initialize and Fill Buffer
    // ========================================
    println!("📝 STEP 2: Filling replay buffer");

    let mut dqn = WorkingDQN::new(config)?;

    for i in 0..200 {
        let state = vec![0.01 * (i as f32); 32];
        let next_state = vec![0.01 * (i as f32 + 1.0); 32];
        let action = (i % 3) as u8;
        let reward = if i % 10 == 0 { 0.1 } else { -0.01 };
        let done = i % 50 == 49;

        let experience = Experience::new(state, action, reward, next_state, done);
        dqn.store_experience(experience)?;
    }

    println!("  ✓ Buffer size: {}", dqn.get_replay_buffer_size()?);
    println!();

    // ========================================
    // STEP 3: Run Training and Monitor Counter
    // ========================================
    println!("📝 STEP 3: Running training and monitoring counter");
    println!();

    let mut consecutive_collapsed = 0;
    let mut saw_collapse_warning = false;

    // Train until we see the error (or hit max epochs)
    for epoch in 0..10 {
        let mut epoch_had_collapse = false;

        // Run 5 steps per epoch
        for step in 0..5 {
            match dqn.train_step(None) {
                Ok((_loss, grad_norm)) => {
                    // Check if gradient collapsed
                    if grad_norm < threshold as f32 {
                        epoch_had_collapse = true;
                        saw_collapse_warning = true;

                        println!(
                            "  Epoch {}, Step {}: grad_norm={:.8} < threshold={:.8} ✗ COLLAPSED",
                            epoch + 1, step + 1, grad_norm, threshold
                        );
                    } else {
                        println!(
                            "  Epoch {}, Step {}: grad_norm={:.8} ✓ OK",
                            epoch + 1, step + 1, grad_norm
                        );
                    }
                },
                Err(e) => {
                    println!("  Epoch {}: Training terminated with error", epoch + 1);
                    println!("  Error: {}", e);
                    break;
                }
            }
        }

        // Update consecutive counter
        if epoch_had_collapse {
            consecutive_collapsed += 1;
            println!("  → Consecutive collapsed epochs: {}/{}", consecutive_collapsed, collapse_patience);
        } else {
            if consecutive_collapsed > 0 {
                println!("  → Gradients recovered! Resetting counter (was: {})", consecutive_collapsed);
            }
            consecutive_collapsed = 0;
        }

        println!();

        // If we hit patience, next epoch should terminate
        if consecutive_collapsed >= collapse_patience {
            println!("  → Patience exhausted! Next step should terminate training");
            break;
        }
    }

    // ========================================
    // VALIDATION 1: Collapse Detected
    // ========================================
    println!("✅ VALIDATION 1: Collapse Detection");
    println!("-----------------------------------");

    assert!(
        saw_collapse_warning,
        "BUG: No gradient collapse detected!\n\
         Expected: Gradients < {:.8}\n\
         This indicates test configuration is incorrect.",
        threshold
    );

    println!("  ✓ Gradient collapse detected");
    println!();

    // ========================================
    // VALIDATION 2: Counter Incremented
    // ========================================
    println!("✅ VALIDATION 2: Counter Increment");
    println!("----------------------------------");

    println!("  Final consecutive collapsed epochs: {}", consecutive_collapsed);

    assert!(
        consecutive_collapsed >= 3,
        "BUG: Counter didn't increment sufficiently!\n\
         Expected: At least 3 consecutive collapsed epochs\n\
         Actual: {}\n\
         This indicates the counter logic is broken.",
        consecutive_collapsed
    );

    println!("  ✓ Counter incremented to {}", consecutive_collapsed);
    println!();

    // ========================================
    // TEST SUMMARY
    // ========================================
    println!("=========================================");
    println!("🎉 TEST PASSED: Counter Logic");
    println!("=========================================");
    println!("✓ Gradient collapse detected");
    println!("✓ Counter incremented correctly");
    println!();

    Ok(())
}

/// Test 3: Verify that training continues when gradients recover
///
/// **Expected**: Counter resets to 0 when gradients recover, training continues
/// **Tests**: Counter reset logic and recovery behavior
#[test]
fn test_gradient_recovery_resets_counter() -> Result<(), MLError> {
    println!("\n🧪 TEST: Gradient Recovery Resets Counter");
    println!("==========================================");
    println!("Objective: Verify counter resets when gradients recover");
    println!();

    // ========================================
    // STEP 1: Configure DQN with MODERATE Parameters
    // ========================================
    println!("📝 STEP 1: Configuring DQN");

    let mut config = WorkingDQNConfig::conservative();
    config.state_dim = 32;
    config.num_actions = 3;

    // Use moderate learning rate (might collapse, might recover)
    config.learning_rate = 5e-6;  // Not as extreme as 1e-8
    config.batch_size = 16;        // Larger than 8 (more stable)

    config.gradient_collapse_patience = 3;
    config.gradient_collapse_multiplier = 100.0;
    config.replay_buffer_capacity = 500;
    config.min_replay_size = 100;
    config.warmup_steps = 0;

    // Don't use C51 narrow bounds (want potential for recovery)
    config.use_distributional = false;

    println!("  ✓ Learning rate: {:.2e} (moderate)", config.learning_rate);
    println!("  ✓ Batch size: {} (moderate)", config.batch_size);
    println!("  ✓ Collapse patience: {} epochs", config.gradient_collapse_patience);
    println!();

    // Save config values before moving
    let threshold = config.learning_rate * config.gradient_collapse_multiplier;

    // ========================================
    // STEP 2: Initialize and Fill Buffer
    // ========================================
    println!("📝 STEP 2: Filling replay buffer with diverse experiences");

    let mut dqn = WorkingDQN::new(config)?;

    // Generate more diverse experiences (helps with potential recovery)
    for i in 0..300 {
        let state = vec![
            0.01 * (i as f32),
            0.02 * ((i % 50) as f32),
            -0.01 * (i as f32),
        ]
        .into_iter()
        .cycle()
        .take(32)
        .collect::<Vec<_>>();

        let next_state = vec![
            0.01 * (i as f32 + 1.0),
            0.02 * ((i % 50 + 1) as f32),
            -0.01 * (i as f32 + 1.0),
        ]
        .into_iter()
        .cycle()
        .take(32)
        .collect::<Vec<_>>();

        let action = (i % 3) as u8;
        let reward = ((i % 20) as f32 - 10.0) / 10.0; // Diverse rewards: -1.0 to +1.0
        let done = i % 50 == 49;

        let experience = Experience::new(state, action, reward, next_state, done);
        dqn.store_experience(experience)?;
    }

    println!("  ✓ Buffer size: {}", dqn.get_replay_buffer_size()?);
    println!();

    // ========================================
    // STEP 3: Monitor for Recovery
    // ========================================
    println!("📝 STEP 3: Training and monitoring for recovery");
    println!();

    let mut saw_recovery = false;

    for epoch in 0..15 {
        let mut epoch_grad_norms = Vec::new();

        // Run 10 steps per epoch
        for _ in 0..10 {
            match dqn.train_step(None) {
                Ok((_, grad_norm)) => {
                    epoch_grad_norms.push(grad_norm);
                },
                Err(e) => {
                    println!("  Epoch {}: Training terminated: {}", epoch + 1, e);
                    break;
                }
            }
        }

        if epoch_grad_norms.is_empty() {
            break;
        }

        let avg_grad_norm = epoch_grad_norms.iter().sum::<f32>() / epoch_grad_norms.len() as f32;
        let is_collapsed = avg_grad_norm < threshold as f32;

        if !is_collapsed {
            saw_recovery = true;
        }

        println!(
            "  Epoch {}: avg_grad_norm={:.8} {} (threshold: {:.8})",
            epoch + 1,
            avg_grad_norm,
            if is_collapsed { "✗ COLLAPSED" } else { "✓ OK" },
            threshold
        );
    }

    println!();

    // ========================================
    // VALIDATION: Recovery Observed
    // ========================================
    println!("✅ VALIDATION: Recovery Behavior");
    println!("--------------------------------");

    // Note: With moderate parameters, we MIGHT see recovery
    // This test is informational - shows counter reset logic exists
    if saw_recovery {
        println!("  ✓ Gradient recovery observed");
        println!("  ℹ Counter reset logic would have been triggered");
    } else {
        println!("  ℹ No recovery observed (parameters too extreme)");
        println!("  ℹ Test configuration note: Use higher LR or larger batch for recovery");
    }
    println!();

    // ========================================
    // TEST SUMMARY
    // ========================================
    println!("==========================================");
    println!("🎉 TEST COMPLETE: Recovery Monitoring");
    println!("==========================================");
    println!("ℹ This test demonstrates recovery detection logic");
    println!("ℹ Adjust parameters (LR, batch size) to observe recovery");
    println!();

    Ok(())
}