//! Early Stopping Module for DQN Training
//!
//! Implements patience-based early stopping to prevent overfitting.
//! Monitors validation loss and stops training when improvement plateaus.
//!
//! WAVE 24 (Agent 17): Anti-Overfitting Enhancement

use tracing::{debug, info};

/// Early stopping state tracker with patience mechanism
///
/// Monitors validation loss and triggers early stopping when:
/// 1. No improvement for `patience` consecutive epochs
/// 2. Improvement is smaller than `min_delta` threshold
///
/// # Example
/// ```rust,ignore
/// let mut early_stop = EarlyStopping::new(5, 0.001);
///
/// // Training loop
/// for epoch in 0..100 {
///     let val_loss = train_epoch();
///
///     if early_stop.should_stop(val_loss) {
///         println!("Early stopping at epoch {}", epoch);
///         break;
///     }
/// }
/// ```
#[derive(Debug, Clone)]
pub struct EarlyStopping {
    /// Number of epochs with no improvement before stopping
    patience: usize,
    /// Minimum improvement threshold to reset patience counter
    min_delta: f64,
    /// Best validation loss observed so far
    best_val_loss: f64,
    /// Counter for epochs without improvement
    counter: usize,
    /// Epoch when best validation loss was achieved
    best_epoch: usize,
    /// Current epoch number
    current_epoch: usize,
}

impl EarlyStopping {
    /// Create new early stopping tracker
    ///
    /// # Arguments
    /// * `patience` - Number of epochs to wait for improvement (default: 5-10)
    /// * `min_delta` - Minimum improvement threshold (default: 0.001-0.01)
    ///
    /// # Recommended Values
    /// - **Short training (<100 epochs)**: patience=5, min_delta=0.01
    /// - **Medium training (100-500 epochs)**: patience=10, min_delta=0.001
    /// - **Long training (>500 epochs)**: patience=20, min_delta=0.0001
    pub fn new(patience: usize, min_delta: f64) -> Self {
        Self {
            patience,
            min_delta,
            best_val_loss: f64::INFINITY,
            counter: 0,
            best_epoch: 0,
            current_epoch: 0,
        }
    }

    /// Check if training should stop
    ///
    /// Returns `true` if no improvement for `patience` consecutive epochs.
    /// Automatically tracks epochs and best loss internally.
    ///
    /// # Arguments
    /// * `val_loss` - Current epoch validation loss
    ///
    /// # Returns
    /// `true` if early stopping triggered, `false` to continue training
    pub fn should_stop(&mut self, val_loss: f64) -> bool {
        self.current_epoch += 1;

        // Check if validation loss improved by more than min_delta
        let improvement = self.best_val_loss - val_loss;

        if improvement > self.min_delta {
            // Significant improvement detected
            debug!(
                "Early stopping: Improvement detected: {:.6} (val_loss: {:.6} -> {:.6})",
                improvement, self.best_val_loss, val_loss
            );

            self.best_val_loss = val_loss;
            self.best_epoch = self.current_epoch;
            self.counter = 0;
            false
        } else {
            // No significant improvement
            self.counter += 1;

            debug!(
                "Early stopping: No improvement for {} epoch(s) (improvement: {:.6} < min_delta: {:.6})",
                self.counter, improvement, self.min_delta
            );

            if self.counter >= self.patience {
                info!(
                    "🛑 Early stopping triggered at epoch {}! No improvement for {} epochs (best: {:.6} at epoch {})",
                    self.current_epoch, self.patience, self.best_val_loss, self.best_epoch
                );
                true
            } else {
                false
            }
        }
    }

    /// Get the best validation loss observed
    pub fn get_best_val_loss(&self) -> f64 {
        self.best_val_loss
    }

    /// Get the epoch when best validation loss was achieved
    pub fn get_best_epoch(&self) -> usize {
        self.best_epoch
    }

    /// Get current patience counter value
    pub fn get_counter(&self) -> usize {
        self.counter
    }

    /// Reset early stopping state (useful for checkpoint resumption)
    pub fn reset(&mut self) {
        self.best_val_loss = f64::INFINITY;
        self.counter = 0;
        self.best_epoch = 0;
        self.current_epoch = 0;
    }

    /// Restore early stopping state from checkpoint
    ///
    /// Useful when resuming training from a saved checkpoint.
    pub fn restore(&mut self, best_val_loss: f64, best_epoch: usize, current_epoch: usize) {
        self.best_val_loss = best_val_loss;
        self.best_epoch = best_epoch;
        self.current_epoch = current_epoch;
        self.counter = 0; // Reset counter on restore
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_early_stopping_triggers_after_patience() {
        let mut early_stop = EarlyStopping::new(3, 0.001);

        // Epoch 1: Initial loss (improves from infinity)
        assert!(!early_stop.should_stop(1.0));
        assert_eq!(early_stop.counter, 0);

        // Epoch 2: Improves significantly
        assert!(!early_stop.should_stop(0.5));
        assert_eq!(early_stop.counter, 0);

        // Epoch 3: No improvement
        assert!(!early_stop.should_stop(0.51));
        assert_eq!(early_stop.counter, 1);

        // Epoch 4: No improvement
        assert!(!early_stop.should_stop(0.52));
        assert_eq!(early_stop.counter, 2);

        // Epoch 5: No improvement - triggers early stop
        assert!(early_stop.should_stop(0.53));
        assert_eq!(early_stop.counter, 3);
    }

    #[test]
    fn test_early_stopping_resets_on_improvement() {
        let mut early_stop = EarlyStopping::new(3, 0.001);

        // Initial loss
        assert!(!early_stop.should_stop(1.0));

        // No improvement for 2 epochs
        assert!(!early_stop.should_stop(1.01));
        assert!(!early_stop.should_stop(1.02));
        assert_eq!(early_stop.counter, 2);

        // Significant improvement - resets counter
        assert!(!early_stop.should_stop(0.5));
        assert_eq!(early_stop.counter, 0);
    }

    #[test]
    fn test_early_stopping_min_delta_threshold() {
        let mut early_stop = EarlyStopping::new(3, 0.01); // 1% min improvement

        // Initial loss
        assert!(!early_stop.should_stop(1.0));

        // Small improvement (0.005 < 0.01) - doesn't reset counter
        assert!(!early_stop.should_stop(0.995));
        assert_eq!(early_stop.counter, 1);

        // Large improvement (0.02 > 0.01) - resets counter
        assert!(!early_stop.should_stop(0.975));
        assert_eq!(early_stop.counter, 0);
    }

    #[test]
    fn test_best_loss_tracking() {
        let mut early_stop = EarlyStopping::new(5, 0.001);

        early_stop.should_stop(1.0);
        assert_eq!(early_stop.get_best_val_loss(), 1.0);
        assert_eq!(early_stop.get_best_epoch(), 1);

        early_stop.should_stop(0.5);
        assert_eq!(early_stop.get_best_val_loss(), 0.5);
        assert_eq!(early_stop.get_best_epoch(), 2);

        early_stop.should_stop(0.6); // Worse - doesn't update
        assert_eq!(early_stop.get_best_val_loss(), 0.5);
        assert_eq!(early_stop.get_best_epoch(), 2);
    }

    #[test]
    fn test_reset() {
        let mut early_stop = EarlyStopping::new(3, 0.001);

        early_stop.should_stop(1.0);
        early_stop.should_stop(1.01);

        early_stop.reset();

        assert_eq!(early_stop.best_val_loss, f64::INFINITY);
        assert_eq!(early_stop.counter, 0);
        assert_eq!(early_stop.best_epoch, 0);
        assert_eq!(early_stop.current_epoch, 0);
    }

    #[test]
    fn test_restore() {
        let mut early_stop = EarlyStopping::new(3, 0.001);

        // Simulate resuming from checkpoint at epoch 50 with best loss 0.5 at epoch 45
        early_stop.restore(0.5, 45, 50);

        assert_eq!(early_stop.best_val_loss, 0.5);
        assert_eq!(early_stop.best_epoch, 45);
        assert_eq!(early_stop.current_epoch, 50);
        assert_eq!(early_stop.counter, 0); // Reset on restore
    }
}
