//! Unit Tests for Early Stopping Components
//!
//! This module contains comprehensive unit tests for early stopping functionality
//! across all ML adapters (PPO, TFT, MAMBA-2, DQN).
//!
//! Test Categories:
//! 1. Configuration Tests - Default values, custom configs, validation
//! 2. State Management Tests - Loss tracking, patience counters, history
//! 3. Strategy Tests - Plateau detection, pruning algorithms
//! 4. Edge Cases - Zero variance, NaN/Inf handling, boundary conditions

use std::collections::HashMap;
use tracing::info;

// ============================================================================
// TEST CATEGORY 1: CONFIGURATION TESTS
// ============================================================================

/// Test that default early stopping configuration values are sensible
#[test]
fn test_default_early_stopping_config() {
    // DQN early stopping defaults
    let dqn_config = crate::trainers::dqn::DQNHyperparameters::conservative();
    
    assert!(dqn_config.early_stopping_enabled, "Early stopping should be enabled by default");
    assert_eq!(dqn_config.min_epochs_before_stopping, 50, "Minimum epochs should be 50");
    assert_eq!(dqn_config.plateau_window, 30, "Plateau window should be 30");
    assert_eq!(dqn_config.q_value_floor, 0.5, "Q-value floor should be 0.5");
    assert_eq!(dqn_config.min_loss_improvement_pct, 2.0, "Min loss improvement should be 2%");
    
    // Verify minimum epochs prevents premature stopping
    assert!(dqn_config.min_epochs_before_stopping >= 20, 
        "Must allow at least 20 epochs for model to warm up");
    
    // Verify plateau window is reasonable
    assert!(dqn_config.plateau_window >= 10 && dqn_config.plateau_window <= 50,
        "Plateau window should be between 10-50 epochs");
}

/// Test that custom early stopping configurations validate correctly
#[test]
fn test_custom_early_stopping_config_validation() {
    // Valid custom configuration
    let custom_config = crate::trainers::dqn::DQNHyperparameters {
        early_stopping_enabled: true,
        min_epochs_before_stopping: 100,
        plateau_window: 20,
        q_value_floor: 0.3,
        min_loss_improvement_pct: 1.0,
        ..Default::default()
    };
    
    assert_eq!(custom_config.min_epochs_before_stopping, 100);
    assert_eq!(custom_config.plateau_window, 20);
    assert_eq!(custom_config.q_value_floor, 0.3);
    assert_eq!(custom_config.min_loss_improvement_pct, 1.0);
}

/// Test that early stopping can be disabled
#[test]
fn test_early_stopping_disable() {
    let config = crate::trainers::dqn::DQNHyperparameters {
        early_stopping_enabled: false,
        ..Default::default()
    };
    
    assert!(!config.early_stopping_enabled, "Early stopping should be disabled");
}

/// Test configuration boundaries and constraints
#[test]
fn test_early_stopping_config_boundaries() {
    // Test minimum values
    let min_config = crate::trainers::dqn::DQNHyperparameters {
        min_epochs_before_stopping: 1,
        plateau_window: 1,
        q_value_floor: 0.0,
        min_loss_improvement_pct: 0.0,
        ..Default::default()
    };
    
    assert_eq!(min_config.min_epochs_before_stopping, 1);
    assert_eq!(min_config.plateau_window, 1);
    assert_eq!(min_config.q_value_floor, 0.0);
    assert_eq!(min_config.min_loss_improvement_pct, 0.0);
    
    // Test maximum values
    let max_config = crate::trainers::dqn::DQNHyperparameters {
        min_epochs_before_stopping: 500,
        plateau_window: 100,
        q_value_floor: 10.0,
        min_loss_improvement_pct: 50.0,
        ..Default::default()
    };
    
    assert_eq!(max_config.min_epochs_before_stopping, 500);
    assert_eq!(max_config.plateau_window, 100);
    assert_eq!(max_config.q_value_floor, 10.0);
    assert_eq!(max_config.min_loss_improvement_pct, 50.0);
}

// ============================================================================
// TEST CATEGORY 2: STATE MANAGEMENT TESTS
// ============================================================================

/// Test loss history tracking
#[test]
fn test_loss_history_tracking() {
    let losses = vec![1.0, 0.9, 0.8, 0.7, 0.6, 0.5];
    
    // Verify history length
    assert_eq!(losses.len(), 6);
    
    // Verify monotonic decrease (ideal case)
    for i in 1..losses.len() {
        assert!(losses[i] < losses[i-1], "Loss should decrease over time");
    }
    
    // Calculate rolling average (window=3)
    let window_size = 3;
    let recent_avg: f64 = losses[losses.len()-window_size..].iter().sum::<f64>() / window_size as f64;
    let older_avg: f64 = losses[losses.len()-2*window_size..losses.len()-window_size].iter().sum::<f64>() / window_size as f64;
    
    assert!(recent_avg < older_avg, "Recent average should be lower than older average");
}

/// Test loss history with noise
#[test]
fn test_loss_history_with_noise() {
    // Realistic loss history with noise
    let losses = vec![
        1.0, 0.95, 1.02, 0.90, 0.88, 0.93,  // Early training: high variance
        0.85, 0.82, 0.84, 0.80, 0.78, 0.81, // Mid training: decreasing with noise
        0.75, 0.74, 0.76, 0.73, 0.74, 0.75, // Late training: plateau with noise
    ];
    
    // Calculate trend (simple linear regression)
    let n = losses.len() as f64;
    let sum_x: f64 = (0..losses.len()).map(|i| i as f64).sum();
    let sum_y: f64 = losses.iter().sum();
    let sum_xy: f64 = losses.iter().enumerate().map(|(i, &y)| i as f64 * y).sum();
    let sum_x2: f64 = (0..losses.len()).map(|i| (i * i) as f64).sum();
    
    let slope = (n * sum_xy - sum_x * sum_y) / (n * sum_x2 - sum_x * sum_x);
    
    // Overall trend should be decreasing (negative slope)
    assert!(slope < 0.0, "Overall trend should be decreasing despite noise");
    
    // Check variance in different phases
    let early_variance = calculate_variance(&losses[0..6]);
    let late_variance = calculate_variance(&losses[12..18]);
    
    // Early training should have higher variance
    assert!(early_variance > late_variance * 0.5, 
        "Early training variance should be higher");
}

fn calculate_variance(values: &[f64]) -> f64 {
    let mean = values.iter().sum::<f64>() / values.len() as f64;
    let variance = values.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / values.len() as f64;
    variance
}

/// Test patience counter increment and reset
#[test]
fn test_patience_counter_behavior() {
    struct PatienceState {
        counter: usize,
        best_loss: f64,
        patience_limit: usize,
    }
    
    impl PatienceState {
        fn new(patience: usize) -> Self {
            Self {
                counter: 0,
                best_loss: f64::INFINITY,
                patience_limit: patience,
            }
        }
        
        fn update(&mut self, loss: f64) -> bool {
            if loss < self.best_loss {
                // Improvement: reset counter
                self.best_loss = loss;
                self.counter = 0;
                false
            } else {
                // No improvement: increment counter
                self.counter += 1;
                self.counter >= self.patience_limit
            }
        }
    }
    
    let mut state = PatienceState::new(3);
    
    // Test improvement resets counter
    assert!(!state.update(1.0)); // Best: 1.0, counter: 0
    assert!(!state.update(0.9)); // Best: 0.9, counter: 0 (reset)
    assert!(!state.update(0.8)); // Best: 0.8, counter: 0 (reset)
    
    // Test no improvement increments counter
    assert!(!state.update(0.9)); // No improvement, counter: 1
    assert!(!state.update(1.0)); // No improvement, counter: 2
    assert!(state.update(1.1));  // No improvement, counter: 3 -> trigger stop
    
    // Test partial improvement
    let mut state2 = PatienceState::new(5);
    assert!(!state2.update(1.0)); // Best: 1.0
    assert!(!state2.update(1.1)); // Counter: 1
    assert!(!state2.update(1.2)); // Counter: 2
    assert!(!state2.update(0.95)); // Best: 0.95, counter: 0 (reset)
    assert!(!state2.update(1.0)); // Counter: 1
    assert!(!state2.update(1.0)); // Counter: 2
    assert!(!state2.update(1.0)); // Counter: 3
    assert!(!state2.update(1.0)); // Counter: 4
    assert!(state2.update(1.0));  // Counter: 5 -> trigger stop
}

/// Test best loss tracking with floating point precision
#[test]
fn test_best_loss_tracking_precision() {
    let mut best_loss = f64::INFINITY;
    let epsilon = 1e-10;
    
    // Test significant improvement
    let new_loss_1 = 1.0;
    assert!(new_loss_1 < best_loss);
    best_loss = new_loss_1;
    
    // Test tiny improvement (below epsilon)
    let new_loss_2 = best_loss - epsilon / 10.0;
    assert!(new_loss_2 < best_loss, "Even tiny improvements should be detected");
    best_loss = new_loss_2;
    
    // Test no improvement (within epsilon)
    let new_loss_3 = best_loss + epsilon / 10.0;
    assert!(new_loss_3 > best_loss, "Tiny increases should be detected");
}

/// Test epoch history storage
#[test]
fn test_epoch_history_storage() {
    #[derive(Debug, Clone)]
    struct EpochHistory {
        epoch: usize,
        train_loss: f64,
        val_loss: f64,
        metrics: HashMap<String, f64>,
    }
    
    let mut history = Vec::new();
    
    // Simulate 10 epochs
    for epoch in 0..10 {
        let mut metrics = HashMap::new();
        metrics.insert("q_value".to_string(), 0.5 + epoch as f64 * 0.05);
        metrics.insert("grad_norm".to_string(), 1.0 - epoch as f64 * 0.05);
        
        history.push(EpochHistory {
            epoch,
            train_loss: 1.0 - epoch as f64 * 0.08,
            val_loss: 1.05 - epoch as f64 * 0.09,
            metrics,
        });
    }
    
    // Verify history length
    assert_eq!(history.len(), 10);
    
    // Verify chronological order
    for i in 1..history.len() {
        assert!(history[i].epoch > history[i-1].epoch);
    }
    
    // Verify metrics are stored correctly
    assert_eq!(history[5].metrics.get("q_value"), Some(&0.75));
    assert!(history[5].metrics.contains_key("grad_norm"));
}

// ============================================================================
// TEST CATEGORY 3: STRATEGY TESTS
// ============================================================================

/// Test plateau detection algorithm accuracy
#[test]
fn test_plateau_detection_accuracy() {
    // Case 1: Clear plateau (constant loss)
    let plateau_losses = vec![0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5];
    assert!(is_plateau(&plateau_losses, 4, 0.1), "Should detect constant loss as plateau");
    
    // Case 2: Slight improvement (not a plateau)
    let improving_losses = vec![0.5, 0.48, 0.46, 0.44, 0.42, 0.40, 0.38, 0.36];
    assert!(!is_plateau(&improving_losses, 4, 2.0), "Should not detect improving loss as plateau");
    
    // Case 3: Noisy plateau (small oscillations around mean)
    let noisy_plateau = vec![0.50, 0.51, 0.49, 0.50, 0.51, 0.49, 0.50, 0.51];
    assert!(is_plateau(&noisy_plateau, 4, 2.0), "Should detect noisy plateau");
    
    // Case 4: Degrading performance (getting worse)
    let degrading_losses = vec![0.5, 0.52, 0.54, 0.56, 0.58, 0.60, 0.62, 0.64];
    assert!(is_plateau(&degrading_losses, 4, 2.0), "Should detect degrading performance as plateau (no improvement)");
}

fn is_plateau(losses: &[f64], window: usize, min_improvement_pct: f64) -> bool {
    if losses.len() < window * 2 {
        return false;
    }
    
    let recent_avg = losses[losses.len()-window..].iter().sum::<f64>() / window as f64;
    let older_avg = losses[losses.len()-2*window..losses.len()-window].iter().sum::<f64>() / window as f64;
    
    let improvement_pct = if older_avg > 0.0 {
        (older_avg - recent_avg) / older_avg * 100.0
    } else {
        0.0
    };
    
    improvement_pct < min_improvement_pct
}

/// Test median pruner correctness
#[test]
fn test_median_pruner_strategy() {
    // Simulate 5 concurrent trials
    let trial_losses = vec![
        vec![1.0, 0.9, 0.8, 0.7], // Trial 0: Good
        vec![1.0, 0.95, 0.92, 0.90], // Trial 1: Below median
        vec![1.0, 0.85, 0.75, 0.65], // Trial 2: Best
        vec![1.0, 0.98, 0.97, 0.96], // Trial 3: Worst
        vec![1.0, 0.93, 0.88, 0.83], // Trial 4: Median
    ];
    
    // At epoch 3, calculate median loss
    let epoch_3_losses: Vec<f64> = trial_losses.iter().map(|t| t[3]).collect();
    let mut sorted_losses = epoch_3_losses.clone();
    sorted_losses.sort_by(|a, b| a.partial_cmp(b).unwrap());
    let median = sorted_losses[sorted_losses.len() / 2];
    
    assert_eq!(median, 0.83, "Median should be trial 4's loss");
    
    // Trials below median should be pruned
    for (i, trial) in trial_losses.iter().enumerate() {
        let loss = trial[3];
        if loss > median {
            // Trial 1 (0.90) and Trial 3 (0.96) should be pruned
            assert!(i == 1 || i == 3, "Trial {} with loss {} should be pruned", i, loss);
        }
    }
}

/// Test percentile pruner correctness
#[test]
fn test_percentile_pruner_strategy() {
    // Simulate 10 trials
    let trial_losses = vec![
        0.95, 0.90, 0.85, 0.80, 0.75, 0.70, 0.65, 0.60, 0.55, 0.50
    ];
    
    // Keep top 50% (prune bottom 50%)
    let percentile = 50;
    let cutoff_index = trial_losses.len() * percentile / 100;
    
    let mut sorted = trial_losses.clone();
    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
    let cutoff_value = sorted[cutoff_index];
    
    // Trials with loss > 0.75 should be pruned
    let pruned: Vec<f64> = trial_losses.iter().filter(|&&l| l > cutoff_value).copied().collect();
    let kept: Vec<f64> = trial_losses.iter().filter(|&&l| l <= cutoff_value).copied().collect();
    
    assert_eq!(kept.len(), 5, "Should keep top 50%");
    assert_eq!(pruned.len(), 5, "Should prune bottom 50%");
}

/// Test successive halving algorithm
#[test]
fn test_successive_halving_strategy() {
    let initial_trials = 16;
    let min_trials = 1;
    let epochs_per_rung = 10;
    
    // Successive halving schedule
    let mut trials_remaining = initial_trials;
    let mut rung = 0;
    let mut schedule = Vec::new();
    
    while trials_remaining > min_trials {
        schedule.push((rung, trials_remaining, rung * epochs_per_rung));
        trials_remaining /= 2;
        rung += 1;
    }
    schedule.push((rung, trials_remaining, rung * epochs_per_rung));
    
    // Verify schedule
    assert_eq!(schedule.len(), 5, "Should have 5 rungs");
    assert_eq!(schedule[0], (0, 16, 0));   // Rung 0: 16 trials, 0 epochs
    assert_eq!(schedule[1], (1, 8, 10));   // Rung 1: 8 trials, 10 epochs
    assert_eq!(schedule[2], (2, 4, 20));   // Rung 2: 4 trials, 20 epochs
    assert_eq!(schedule[3], (3, 2, 30));   // Rung 3: 2 trials, 30 epochs
    assert_eq!(schedule[4], (4, 1, 40));   // Rung 4: 1 trial, 40 epochs
}

/// Test hyperband bracket logic
#[test]
fn test_hyperband_bracket_logic() {
    let max_iter = 81; // Maximum epochs per trial
    let eta = 3;       // Reduction factor
    
    // Calculate number of brackets (s_max + 1)
    let s_max = (max_iter as f64).log(eta as f64).floor() as usize;
    
    assert_eq!(s_max, 4, "Should have 5 brackets (s_max=4)");
    
    // For each bracket, calculate resource allocation
    for s in (0..=s_max).rev() {
        let n = ((s_max + 1) as f64 / (s + 1) as f64 * eta.pow(s as u32) as f64).ceil() as usize;
        let r = max_iter / eta.pow(s as u32) as usize;
        
        info!(bracket = s, trials = n, epochs_per_trial = r, "Bracket resource allocation");
        
        // Verify reasonable resource allocation
        assert!(n >= 1, "Should have at least 1 trial");
        assert!(r >= 1, "Should have at least 1 epoch");
        assert!(n * r <= max_iter * eta.pow(s_max as u32), "Total resources should be bounded");
    }
}

// ============================================================================
// TEST CATEGORY 4: EDGE CASES
// ============================================================================

/// Test early stopping with zero variance data (constant loss)
#[test]
fn test_early_stopping_zero_variance() {
    let constant_losses = vec![0.5; 100]; // 100 epochs, all same loss
    
    let window = 10;
    let min_improvement = 0.1; // 0.1% improvement required
    
    // Should detect plateau immediately after 2*window epochs
    assert!(constant_losses.len() >= window * 2);
    assert!(is_plateau(&constant_losses[..window*2], window, min_improvement));
}

/// Test early stopping with NaN loss
#[test]
fn test_early_stopping_nan_loss() {
    let losses = vec![1.0, 0.9, 0.8, f64::NAN, 0.6];
    
    // Check if NaN is present
    let has_nan = losses.iter().any(|&l| l.is_nan());
    assert!(has_nan, "Should detect NaN in loss history");
    
    // NaN should cause immediate failure (not counted as improvement)
    for &loss in &losses {
        if loss.is_nan() {
            // In production, this would trigger trial failure
            assert!(loss.is_nan());
        }
    }
}

/// Test early stopping with Inf loss
#[test]
fn test_early_stopping_inf_loss() {
    let losses = vec![1.0, 0.9, 0.8, f64::INFINITY, 0.6];
    
    // Check if Inf is present
    let has_inf = losses.iter().any(|&l| l.is_infinite());
    assert!(has_inf, "Should detect Inf in loss history");
    
    // Inf should cause immediate failure
    for &loss in &losses {
        if loss.is_infinite() {
            assert!(loss.is_infinite());
        }
    }
}

/// Test early stopping with single epoch training
#[test]
fn test_early_stopping_single_epoch() {
    let config = crate::trainers::dqn::DQNHyperparameters {
        epochs: 1,
        min_epochs_before_stopping: 50,
        ..Default::default()
    };
    
    // Should not trigger early stopping (epochs < min_epochs)
    assert!(config.epochs < config.min_epochs_before_stopping);
}

/// Test early stopping with very small improvements
#[test]
fn test_early_stopping_tiny_improvements() {
    let epsilon = 1e-8;
    let losses = vec![
        1.0,
        1.0 - epsilon,
        1.0 - 2.0 * epsilon,
        1.0 - 3.0 * epsilon,
        1.0 - 4.0 * epsilon,
    ];
    
    // Calculate improvement percentage
    let window = 2;
    let recent_avg = losses[losses.len()-window..].iter().sum::<f64>() / window as f64;
    let older_avg = losses[losses.len()-2*window..losses.len()-window].iter().sum::<f64>() / window as f64;
    
    let improvement_pct = ((older_avg - recent_avg) / older_avg * 100.0).abs();
    
    // Improvement should be very small (< 0.001%)
    assert!(improvement_pct < 0.001, "Improvement should be negligible");
}

/// Test early stopping with large batch size edge case
#[test]
fn test_early_stopping_large_batch_size() {
    // Test that batch size doesn't affect early stopping logic
    let config_small_batch = crate::trainers::dqn::DQNHyperparameters {
        batch_size: 32,
        ..Default::default()
    };
    
    let config_large_batch = crate::trainers::dqn::DQNHyperparameters {
        batch_size: 230, // Max for RTX 3050 Ti
        ..Default::default()
    };
    
    // Early stopping parameters should be independent of batch size
    assert_eq!(config_small_batch.min_epochs_before_stopping, 
               config_large_batch.min_epochs_before_stopping);
    assert_eq!(config_small_batch.plateau_window, 
               config_large_batch.plateau_window);
}

/// Test early stopping with extreme learning rates
#[test]
fn test_early_stopping_extreme_learning_rates() {
    // Very low learning rate (slow convergence)
    let config_low_lr = crate::trainers::dqn::DQNHyperparameters {
        learning_rate: 1e-6,
        ..Default::default()
    };
    
    // Very high learning rate (unstable training)
    let config_high_lr = crate::trainers::dqn::DQNHyperparameters {
        learning_rate: 1e-1,
        ..Default::default()
    };
    
    // Early stopping should be same regardless of LR
    assert_eq!(config_low_lr.plateau_window, config_high_lr.plateau_window);
}

// ============================================================================
// HELPER FUNCTIONS
// ============================================================================

/// Helper: Calculate improvement percentage between two values
fn calculate_improvement_pct(old_value: f64, new_value: f64) -> f64 {
    if old_value > 0.0 {
        (old_value - new_value) / old_value * 100.0
    } else {
        0.0
    }
}

/// Helper: Check if value is valid (not NaN or Inf)
fn is_valid_loss(loss: f64) -> bool {
    !loss.is_nan() && !loss.is_infinite()
}

#[test]
fn test_helpers() {
    assert_eq!(calculate_improvement_pct(1.0, 0.9), 10.0);
    assert_eq!(calculate_improvement_pct(1.0, 1.1), -10.0);
    assert_eq!(calculate_improvement_pct(0.0, 0.5), 0.0);
    
    assert!(is_valid_loss(1.0));
    assert!(is_valid_loss(0.0));
    assert!(!is_valid_loss(f64::NAN));
    assert!(!is_valid_loss(f64::INFINITY));
}