foxhunt/ml/tests/hyperopt_early_stopping_infrastructure_test.rs

//! Comprehensive test suite for early stopping infrastructure
//!
//! This module tests all components of the early stopping system:
//! - Configuration and defaults
//! - State management and persistence
//! - Strategy implementations (Plateau, MedianPruner, Percentile)
//! - Observer pattern and callbacks
//! - Cross-trial pruning logic

use ml::hyperopt::early_stopping::{
    EarlyStoppingConfig, EarlyStoppingObserver, EarlyStoppingState, EarlyStoppingStrategy,
    EarlyStoppingStrategyTrait, EpochMetrics, MedianPrunerStrategy, ObserverDecision,
    PercentilePrunerStrategy, PlateauDetectionStrategy, TrialObserver, TrialState,
};

#[test]
fn test_early_stopping_config_defaults() {
    let config = EarlyStoppingConfig::default();

    assert_eq!(config.patience_epochs, 10);
    assert_eq!(config.min_delta, 1e-4);
    assert_eq!(config.validation_frequency, 1);
    assert_eq!(config.min_epochs, 20);
    assert!(!config.compare_to_baseline);
    assert!(matches!(config.strategy, EarlyStoppingStrategy::Plateau));
}

#[test]
fn test_early_stopping_config_custom() {
    let config = EarlyStoppingConfig {
        patience_epochs: 15,
        min_delta: 1e-3,
        compare_to_baseline: true,
        validation_frequency: 2,
        min_epochs: 30,
        strategy: EarlyStoppingStrategy::MedianPruner { warmup_steps: 10 },
    };

    assert_eq!(config.patience_epochs, 15);
    assert_eq!(config.min_delta, 1e-3);
    assert!(config.compare_to_baseline);
    assert_eq!(config.validation_frequency, 2);
    assert_eq!(config.min_epochs, 30);
    assert!(matches!(
        config.strategy,
        EarlyStoppingStrategy::MedianPruner { .. }
    ));
}

#[test]
fn test_trial_state_transitions() {
    // Test state lifecycle
    assert_eq!(TrialState::Pending, TrialState::Pending);

    let running = TrialState::Running { current_epoch: 5 };
    assert!(matches!(running, TrialState::Running { .. }));

    let pruned = TrialState::Pruned {
        stopped_at_epoch: 10,
        reason: "Worse than median".to_string(),
    };
    assert!(matches!(pruned, TrialState::Pruned { .. }));

    assert_eq!(TrialState::Completed, TrialState::Completed);
}

#[test]
fn test_early_stopping_state_initialization() {
    let state = EarlyStoppingState::new();

    assert_eq!(state.best_val_loss, f64::INFINITY);
    assert_eq!(state.patience_counter, 0);
    assert!(!state.stopped);
    assert!(state.stopped_at_epoch.is_none());
    assert_eq!(state.epoch_history.len(), 0);
}

#[test]
fn test_early_stopping_state_improvement_detection() {
    let mut state = EarlyStoppingState::new();

    // First update always improves
    let improved = state.update(0.5, 1e-4);
    assert!(improved);
    assert_eq!(state.best_val_loss, 0.5);
    assert_eq!(state.patience_counter, 0);

    // Significant improvement
    let improved = state.update(0.3, 1e-4);
    assert!(improved);
    assert_eq!(state.best_val_loss, 0.3);
    assert_eq!(state.patience_counter, 0);

    // Marginal improvement (< min_delta)
    let improved = state.update(0.29999, 1e-4);
    assert!(!improved);
    assert_eq!(state.best_val_loss, 0.3);
    assert_eq!(state.patience_counter, 1);

    // Worse performance
    let improved = state.update(0.4, 1e-4);
    assert!(!improved);
    assert_eq!(state.best_val_loss, 0.3);
    assert_eq!(state.patience_counter, 2);
}

#[test]
fn test_plateau_detection_strategy_basic() {
    let strategy = PlateauDetectionStrategy::new(5, 1e-3);
    let mut state = EarlyStoppingState::new();

    // Improving performance - should not stop
    for epoch in 0..10 {
        let val_loss = 1.0 - (epoch as f64 * 0.05); // 1.0 → 0.5
        let should_stop = strategy.should_stop(epoch, val_loss, &mut state, &[]);
        assert!(
            !should_stop,
            "Should not stop during improvement at epoch {}",
            epoch
        );
    }
}

#[test]
fn test_plateau_detection_strategy_triggers() {
    let strategy = PlateauDetectionStrategy::new(3, 1e-3);
    let mut state = EarlyStoppingState::new();

    // Set initial best loss
    state.update(0.5, 1e-3);

    // Plateau for 3 epochs (patience threshold)
    for epoch in 1..=3 {
        let should_stop = strategy.should_stop(epoch, 0.5, &mut state, &[]);
        if epoch < 3 {
            assert!(
                !should_stop,
                "Should not stop before patience at epoch {}",
                epoch
            );
        } else {
            assert!(
                should_stop,
                "Should stop after patience exhausted at epoch {}",
                epoch
            );
        }
    }
}

#[test]
fn test_plateau_detection_respects_min_epochs() {
    // Test that observer respects min_epochs (strategy doesn't have this knowledge)
    let config = EarlyStoppingConfig {
        min_epochs: 10,
        patience_epochs: 3,
        min_delta: 1e-3,
        strategy: EarlyStoppingStrategy::Plateau,
        ..Default::default()
    };

    let mut observer = EarlyStoppingObserver::new(config);
    observer.on_trial_start(0, "test");

    // Plateau immediately but before min_epochs
    for epoch in 0..5 {
        let metrics = EpochMetrics {
            epoch,
            train_loss: 0.5,
            val_loss: 0.5,
            timestamp: epoch as f64,
        };
        let decision = observer.on_epoch_complete(0, epoch, &metrics);
        assert_eq!(
            decision,
            ObserverDecision::Continue,
            "Should not stop before min_epochs at epoch {}",
            epoch
        );
    }
}

#[test]
fn test_median_pruner_strategy_warmup() {
    let strategy = MedianPrunerStrategy::new(5);
    let trial_history = vec![0.4, 0.5, 0.6]; // 3 completed trials

    // Should not prune during warmup (epoch < 5)
    for epoch in 0..5 {
        let should_stop = strategy.should_stop(epoch, 0.8, &trial_history);
        assert!(
            !should_stop,
            "Should not prune during warmup at epoch {}",
            epoch
        );
    }
}

#[test]
fn test_median_pruner_strategy_prunes_worse_than_median() {
    let strategy = MedianPrunerStrategy::new(5);
    let trial_history = vec![0.3, 0.4, 0.5, 0.6, 0.7]; // Median = 0.5

    // Current trial worse than median
    let should_stop = strategy.should_stop(10, 0.8, &trial_history);
    assert!(should_stop, "Should prune trial worse than median");

    // Current trial better than median
    let should_stop = strategy.should_stop(10, 0.2, &trial_history);
    assert!(!should_stop, "Should not prune trial better than median");

    // Current trial equal to median
    let should_stop = strategy.should_stop(10, 0.5, &trial_history);
    assert!(!should_stop, "Should not prune trial equal to median");
}

#[test]
fn test_median_pruner_with_insufficient_history() {
    let strategy = MedianPrunerStrategy::new(5);
    let trial_history = vec![0.5]; // Only 1 trial

    // Should not prune with insufficient history
    let should_stop = strategy.should_stop(10, 0.8, &trial_history);
    assert!(
        !should_stop,
        "Should not prune with insufficient trial history"
    );
}

#[test]
fn test_percentile_pruner_strategy_warmup() {
    let strategy = PercentilePrunerStrategy::new(25.0, 5); // Bottom 25%
    let trial_history = vec![0.3, 0.4, 0.5, 0.6];

    // Should not prune during warmup
    for epoch in 0..5 {
        let should_stop = strategy.should_stop(epoch, 0.8, &trial_history);
        assert!(
            !should_stop,
            "Should not prune during warmup at epoch {}",
            epoch
        );
    }
}

#[test]
fn test_percentile_pruner_strategy_prunes_bottom_percentile() {
    let strategy = PercentilePrunerStrategy::new(75.0, 5); // Prune if worse than 75th percentile (top 25% performers)
    let trial_history = vec![0.2, 0.4, 0.6, 0.8]; // 75th percentile = 0.65

    // Current trial worse than 75th percentile (top 25%)
    let should_stop = strategy.should_stop(10, 0.9, &trial_history);
    assert!(should_stop, "Should prune trial worse than 75th percentile");

    // Current trial better than 75th percentile
    let should_stop = strategy.should_stop(10, 0.3, &trial_history);
    assert!(
        !should_stop,
        "Should not prune trial better than 75th percentile"
    );
}

#[test]
fn test_epoch_metrics_creation() {
    let metrics = EpochMetrics {
        epoch: 10,
        train_loss: 0.5,
        val_loss: 0.6,
        timestamp: 1234567890.0,
    };

    assert_eq!(metrics.epoch, 10);
    assert_eq!(metrics.train_loss, 0.5);
    assert_eq!(metrics.val_loss, 0.6);
    assert_eq!(metrics.timestamp, 1234567890.0);
}

#[test]
fn test_early_stopping_observer_initialization() {
    let config = EarlyStoppingConfig::default();
    let observer = EarlyStoppingObserver::new(config.clone());

    // Observer should start with empty state
    assert_eq!(observer.trial_count(), 0);
}

#[test]
fn test_early_stopping_observer_trial_lifecycle() {
    let config = EarlyStoppingConfig {
        patience_epochs: 3,
        min_epochs: 5,
        ..Default::default()
    };
    let mut observer = EarlyStoppingObserver::new(config);

    // Start trial
    observer.on_trial_start(0, "learning_rate=0.001");
    assert_eq!(observer.trial_count(), 1);

    // Improving performance
    for epoch in 0..10 {
        let metrics = EpochMetrics {
            epoch,
            train_loss: 0.5 - (epoch as f64 * 0.01),
            val_loss: 0.6 - (epoch as f64 * 0.01),
            timestamp: epoch as f64,
        };

        let decision = observer.on_epoch_complete(0, epoch, &metrics);
        assert_eq!(decision, ObserverDecision::Continue);
    }

    // Complete trial
    observer.on_trial_complete(0, 0.3);
}

#[test]
fn test_early_stopping_observer_plateau_detection() {
    let config = EarlyStoppingConfig {
        patience_epochs: 3,
        min_epochs: 5,
        min_delta: 1e-3,
        ..Default::default()
    };
    let mut observer = EarlyStoppingObserver::new(config);

    observer.on_trial_start(0, "test_params");

    // Initial improvement
    for epoch in 0..5 {
        let metrics = EpochMetrics {
            epoch,
            train_loss: 0.5,
            val_loss: 0.5,
            timestamp: epoch as f64,
        };
        let decision = observer.on_epoch_complete(0, epoch, &metrics);
        assert_eq!(decision, ObserverDecision::Continue);
    }

    // Plateau for patience_epochs (should trigger stop after min_epochs)
    for epoch in 5..10 {
        let metrics = EpochMetrics {
            epoch,
            train_loss: 0.5,
            val_loss: 0.5,
            timestamp: epoch as f64,
        };
        let decision = observer.on_epoch_complete(0, epoch, &metrics);

        if epoch >= 8 {
            // After 3 epochs of plateau (epochs 5, 6, 7 → patience exhausted at 8)
            assert_eq!(decision, ObserverDecision::StopTrial);
            break;
        }
    }
}

#[test]
fn test_early_stopping_observer_multiple_trials() {
    let config = EarlyStoppingConfig::default();
    let mut observer = EarlyStoppingObserver::new(config);

    // Trial 0
    observer.on_trial_start(0, "trial_0");
    for epoch in 0..5 {
        let metrics = EpochMetrics {
            epoch,
            train_loss: 0.5,
            val_loss: 0.5,
            timestamp: epoch as f64,
        };
        observer.on_epoch_complete(0, epoch, &metrics);
    }
    observer.on_trial_complete(0, 0.5);

    // Trial 1
    observer.on_trial_start(1, "trial_1");
    for epoch in 0..5 {
        let metrics = EpochMetrics {
            epoch,
            train_loss: 0.3,
            val_loss: 0.3,
            timestamp: epoch as f64,
        };
        observer.on_epoch_complete(1, epoch, &metrics);
    }
    observer.on_trial_complete(1, 0.3);

    assert_eq!(observer.trial_count(), 2);
}

#[test]
fn test_early_stopping_observer_failure_handling() {
    let config = EarlyStoppingConfig::default();
    let mut observer = EarlyStoppingObserver::new(config);

    observer.on_trial_start(0, "failing_trial");
    observer.on_trial_failed(0, "OOM error");

    // Failed trial should still be tracked
    assert_eq!(observer.trial_count(), 1);
}

#[test]
fn test_observer_decision_equality() {
    assert_eq!(ObserverDecision::Continue, ObserverDecision::Continue);
    assert_eq!(ObserverDecision::StopTrial, ObserverDecision::StopTrial);
    assert_eq!(ObserverDecision::StopStudy, ObserverDecision::StopStudy);

    assert_ne!(ObserverDecision::Continue, ObserverDecision::StopTrial);
    assert_ne!(ObserverDecision::StopTrial, ObserverDecision::StopStudy);
}

#[test]
fn test_early_stopping_state_records_epoch_history() {
    let mut state = EarlyStoppingState::new();

    state.update(0.5, 1e-4);
    state.record_epoch_metrics(EpochMetrics {
        epoch: 0,
        train_loss: 0.5,
        val_loss: 0.5,
        timestamp: 0.0,
    });

    state.update(0.3, 1e-4);
    state.record_epoch_metrics(EpochMetrics {
        epoch: 1,
        train_loss: 0.3,
        val_loss: 0.3,
        timestamp: 1.0,
    });

    assert_eq!(state.epoch_history.len(), 2);
    assert_eq!(state.epoch_history[0].epoch, 0);
    assert_eq!(state.epoch_history[1].epoch, 1);
}

#[test]
fn test_strategy_enum_matching() {
    let plateau = EarlyStoppingStrategy::Plateau;
    assert!(matches!(plateau, EarlyStoppingStrategy::Plateau));

    let median = EarlyStoppingStrategy::MedianPruner { warmup_steps: 5 };
    if let EarlyStoppingStrategy::MedianPruner { warmup_steps } = median {
        assert_eq!(warmup_steps, 5);
    }

    let percentile = EarlyStoppingStrategy::PercentilePruner {
        percentile: 25.0,
        warmup_steps: 10,
    };
    if let EarlyStoppingStrategy::PercentilePruner {
        percentile,
        warmup_steps,
    } = percentile
    {
        assert_eq!(percentile, 25.0);
        assert_eq!(warmup_steps, 10);
    }
}

#[test]
fn test_serde_serialization() {
    use serde_json;

    let config = EarlyStoppingConfig {
        patience_epochs: 15,
        min_delta: 1e-3,
        compare_to_baseline: true,
        validation_frequency: 2,
        min_epochs: 30,
        strategy: EarlyStoppingStrategy::MedianPruner { warmup_steps: 10 },
    };

    // Serialize
    let json = serde_json::to_string(&config).unwrap();

    // Deserialize
    let deserialized: EarlyStoppingConfig = serde_json::from_str(&json).unwrap();

    assert_eq!(deserialized.patience_epochs, 15);
    assert_eq!(deserialized.min_delta, 1e-3);
}

#[test]
fn test_integration_plateau_then_improvement() {
    let config = EarlyStoppingConfig {
        patience_epochs: 5,
        min_epochs: 10,
        min_delta: 1e-3,
        ..Default::default()
    };
    let mut observer = EarlyStoppingObserver::new(config);

    observer.on_trial_start(0, "test");

    // Plateau for 4 epochs (just below patience)
    for epoch in 0..4 {
        let metrics = EpochMetrics {
            epoch,
            train_loss: 0.5,
            val_loss: 0.5,
            timestamp: epoch as f64,
        };
        let decision = observer.on_epoch_complete(0, epoch, &metrics);
        assert_eq!(decision, ObserverDecision::Continue);
    }

    // Improvement resets patience
    let metrics = EpochMetrics {
        epoch: 4,
        train_loss: 0.3,
        val_loss: 0.3,
        timestamp: 4.0,
    };
    let decision = observer.on_epoch_complete(0, 4, &metrics);
    assert_eq!(decision, ObserverDecision::Continue);

    // Continue training - will plateau again at 0.3
    // After improvement at epoch 4, patience resets to 0
    // Epochs 5-9: no improvement, patience builds (1,2,3,4,5)
    // Epoch 10: min_epochs satisfied, patience=5 hits limit, should stop
    for epoch in 5..20 {
        let metrics = EpochMetrics {
            epoch,
            train_loss: 0.3,
            val_loss: 0.3,
            timestamp: epoch as f64,
        };
        let decision = observer.on_epoch_complete(0, epoch, &metrics);

        // Should continue until epoch 10 (when patience=5 and min_epochs=10 both satisfied)
        if epoch < 10 {
            assert_eq!(
                decision,
                ObserverDecision::Continue,
                "Should continue at epoch {}",
                epoch
            );
        } else if epoch == 10 {
            // At epoch 10: patience counter should be 5 (epochs 5,6,7,8,9 without improvement)
            // and min_epochs=10 is satisfied, so should stop
            assert_eq!(
                decision,
                ObserverDecision::StopTrial,
                "Should stop at epoch {} after patience exhausted",
                epoch
            );
            break;
        }
    }
}

/// Test that the module exports are accessible
#[test]
fn test_module_exports() {
    use ml::hyperopt::early_stopping::*;

    // Verify all main types are exported
    let _config: EarlyStoppingConfig = EarlyStoppingConfig::default();
    let _state: EarlyStoppingState = EarlyStoppingState::new();
    let _strategy: EarlyStoppingStrategy = EarlyStoppingStrategy::Plateau;
    let _decision: ObserverDecision = ObserverDecision::Continue;
    let _trial_state: TrialState = TrialState::Pending;
}