//! Early Stopping Infrastructure for Hyperparameter Optimization
//!
//! This module provides production-ready early stopping strategies for hyperparameter
//! optimization trials. It implements multiple stopping criteria including:
//!
//! - **Plateau Detection**: Stop when loss plateaus (no improvement for N epochs)
//! - **Median Pruner**: Stop trials performing worse than median
//! - **Percentile Pruner**: Stop trials in bottom N percentile
//! - **Observer Pattern**: Extensible callback system for custom strategies
//!
//! ## Design Principles
//!
//! 1. **Orthogonal**: Works independently of optimizer internals
//! 2. **Zero-overhead**: Disabled by default, no cost when not used
//! 3. **Type-safe**: Generic over parameter types
//! 4. **Production-ready**: Comprehensive error handling and logging
//!
//! ## Architecture
//!
//! ```text
//! ┌─────────────────────────────────────────────────────────────┐
//! │                    TrialObserver Trait                      │
//! │  ┌──────────────────┐    ┌────────────────────────────┐    │
//! │  │ on_trial_start   │    │  on_epoch_complete         │    │
//! │  └──────────────────┘    └────────────────────────────┘    │
//! └─────────────────────────────────────────────────────────────┘
//!                              ↓
//!                 ┌─────────────────────────────┐
//!                 │  EarlyStoppingObserver      │
//!                 │  ┌─────────────────────┐   │
//!                 │  │ EarlyStoppingConfig │   │
//!                 │  └─────────────────────┘   │
//!                 │  ┌─────────────────────┐   │
//!                 │  │ Per-trial State     │   │
//!                 │  │ HashMap<id, State>  │   │
//!                 │  └─────────────────────┘   │
//!                 └─────────────────────────────┘
//!                              ↓
//!                 ┌─────────────────────────────┐
//!                 │  EarlyStoppingStrategy      │
//!                 │  ┌───────────────────────┐  │
//!                 │  │ Plateau               │  │
//!                 │  │ MedianPruner          │  │
//!                 │  │ PercentilePruner      │  │
//!                 │  │ SuccessiveHalving     │  │
//!                 │  │ Hyperband             │  │
//!                 │  └───────────────────────┘  │
//!                 └─────────────────────────────┘
//! ```
//!
//! ## Usage Example
//!
//! ```rust,no_run
//! use ml::hyperopt::early_stopping::{
//!     EarlyStoppingConfig, EarlyStoppingObserver, EarlyStoppingStrategy,
//!     TrialObserver, ObserverDecision, EpochMetrics,
//! };
//!
//! // Create observer with plateau detection
//! let config = EarlyStoppingConfig {
//!     patience_epochs: 10,
//!     min_delta: 1e-4,
//!     strategy: EarlyStoppingStrategy::Plateau,
//!     ..Default::default()
//! };
//! let mut observer = EarlyStoppingObserver::new(config);
//!
//! // In training loop
//! observer.on_trial_start(0, "learning_rate=0.001");
//!
//! for epoch in 0..100 {
//!     // ... training code ...
//!     let val_loss = 0.5; // From validation
//!     
//!     let metrics = EpochMetrics {
//!         epoch,
//!         train_loss: 0.4,
//!         val_loss,
//!         timestamp: epoch as f64,
//!     };
//!     
//!     match observer.on_epoch_complete(0, epoch, &metrics) {
//!         ObserverDecision::Continue => continue,
//!         ObserverDecision::StopTrial => {
//!             println!("Trial stopped early at epoch {}", epoch);
//!             break;
//!         }
//!         ObserverDecision::StopStudy => {
//!             println!("Study stopped early");
//!             return;
//!         }
//!     }
//! }
//!
//! observer.on_trial_complete(0, 0.5);
//! ```
//!
//! ## Performance Characteristics
//!
//! - **Memory**: O(trials × epochs) for history tracking
//! - **Per-epoch overhead**: ~1-10μs for decision logic
//! - **Cross-trial pruning**: O(trials) for median computation
//!
//! ## Integration with Adapters
//!
//! Adapters can integrate early stopping by:
//!
//! 1. Creating observer: `let observer = EarlyStoppingObserver::new(config)`
//! 2. Calling `on_trial_start()` before training
//! 3. Calling `on_epoch_complete()` after each epoch
//! 4. Breaking training loop on `StopTrial` decision
//! 5. Calling `on_trial_complete()` or `on_trial_failed()` at end

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use tracing::{info, warn};

/// Configuration for early stopping criteria
///
/// Defines when to stop training trials based on validation loss behavior.
/// All configurations support serialization for experiment tracking.
///
/// ## Parameter Guidelines
///
/// - **patience_epochs**: 5-20 epochs depending on model convergence speed
/// - **min_delta**: 1e-4 for stable models, 1e-3 for noisy training
/// - **min_epochs**: 20-50 epochs to avoid premature stopping
/// - **validation_frequency**: 1 for small models, 2-5 for large models
///
/// ## Example
///
/// ```rust
/// use ml::hyperopt::early_stopping::{EarlyStoppingConfig, EarlyStoppingStrategy};
///
/// let config = EarlyStoppingConfig {
///     patience_epochs: 15,
///     min_delta: 5e-4,
///     min_epochs: 30,
///     strategy: EarlyStoppingStrategy::Plateau,
///     ..Default::default()
/// };
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EarlyStoppingConfig {
    /// Number of epochs with no improvement before stopping
    ///
    /// Training stops if validation loss doesn't improve by at least
    /// `min_delta` for this many consecutive epochs.
    pub patience_epochs: usize,

    /// Minimum improvement to reset patience (absolute value)
    ///
    /// An improvement is considered significant if:
    /// `old_loss - new_loss > min_delta`
    pub min_delta: f64,

    /// Whether to compare to best trial across all trials
    ///
    /// When true, uses study-wide best loss as baseline instead of
    /// trial's own best loss. Useful for aggressive pruning.
    pub compare_to_baseline: bool,

    /// Frequency of validation checks (every N epochs)
    ///
    /// Set to 1 for epoch-by-epoch monitoring, or higher values
    /// to reduce overhead for expensive validation.
    pub validation_frequency: usize,

    /// Minimum epochs before early stopping can trigger
    ///
    /// Prevents premature stopping during initial exploration.
    /// Recommended: 20-50 epochs depending on model complexity.
    pub min_epochs: usize,

    /// Strategy to use for early stopping decision
    ///
    /// Available strategies:
    /// - `Plateau`: Loss plateau detection (single-trial)
    /// - `MedianPruner`: Prune if worse than median (cross-trial)
    /// - `PercentilePruner`: Prune bottom N% (cross-trial)
    /// - `SuccessiveHalving`: Aggressive halving (cross-trial)
    /// - `Hyperband`: Optimal resource allocation (cross-trial)
    pub strategy: EarlyStoppingStrategy,
}

impl Default for EarlyStoppingConfig {
    fn default() -> Self {
        Self {
            patience_epochs: 10,
            min_delta: 1e-4,
            compare_to_baseline: false,
            validation_frequency: 1,
            min_epochs: 20,
            strategy: EarlyStoppingStrategy::Plateau,
        }
    }
}

/// State of a trial in the optimization process
///
/// Tracks the lifecycle of a single trial from initialization through
/// completion or early termination. Used for result analysis and debugging.
///
/// ## State Transitions
///
/// ```text
/// Pending → Running → {EarlyStopped*, Completed, Failed}
///                ↓
///            Pruned*
/// ```
///
/// *Early termination states preserve the epoch where stopping occurred
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum TrialState {
    /// Trial is waiting to start
    Pending,

    /// Trial is currently running
    Running {
        /// Current epoch number
        current_epoch: usize,
    },

    /// Trial stopped early due to loss plateau
    ///
    /// Triggered when validation loss fails to improve for
    /// `patience_epochs` consecutive epochs.
    EarlyStoppedPlateau {
        /// Epoch where stopping occurred
        stopped_at_epoch: usize,
    },

    /// Trial stopped early due to threshold violation
    ///
    /// Triggered when validation loss exceeds a predefined
    /// maximum threshold (e.g., divergence detection).
    EarlyStoppedThreshold {
        /// Epoch where stopping occurred
        stopped_at_epoch: usize,
    },

    /// Trial stopped early by pruner (median/percentile)
    ///
    /// Triggered by cross-trial comparison strategies when trial
    /// performs worse than the study median or target percentile.
    Pruned {
        /// Epoch where pruning occurred
        stopped_at_epoch: usize,
        /// Human-readable reason (e.g., "Worse than median")
        reason: String,
    },

    /// Trial completed successfully
    Completed,

    /// Trial failed with error
    Failed {
        /// Error message describing failure cause
        error: String,
    },
}

/// Per-trial early stopping state
///
/// Maintains all state needed to make early stopping decisions for a single trial.
/// Includes best loss tracking, patience counter, and epoch history for analysis.
///
/// ## Memory Usage
///
/// - Fixed: ~48 bytes (best_loss, counters, flags)
/// - Variable: ~48 bytes × epochs (epoch history)
/// - Total: ~48 + 48N bytes for N epochs
///
/// ## Example
///
/// ```rust
/// use ml::hyperopt::early_stopping::EarlyStoppingState;
///
/// let mut state = EarlyStoppingState::new();
///
/// // Update with validation loss
/// let improved = state.update(0.5, 1e-4);
/// assert!(improved); // First update always improves
///
/// let improved = state.update(0.3, 1e-4);
/// assert!(improved); // Significant improvement
///
/// let improved = state.update(0.299, 1e-4);
/// assert!(!improved); // Marginal improvement (< min_delta)
/// ```
#[derive(Debug, Clone)]
pub struct EarlyStoppingState {
    /// Best validation loss seen so far
    pub best_val_loss: f64,

    /// Consecutive epochs without improvement
    pub patience_counter: usize,

    /// Whether trial was stopped early
    pub stopped: bool,

    /// Epoch where stopping occurred (if stopped)
    pub stopped_at_epoch: Option<usize>,

    /// Complete epoch history for this trial
    pub epoch_history: Vec<EpochMetrics>,
}

impl EarlyStoppingState {
    /// Create a new early stopping state
    ///
    /// Initializes with infinite loss (no observations yet) and
    /// zero patience counter.
    pub fn new() -> Self {
        Self {
            best_val_loss: f64::INFINITY,
            patience_counter: 0,
            stopped: false,
            stopped_at_epoch: None,
            epoch_history: Vec::new(),
        }
    }

    /// Update state with new validation loss
    ///
    /// Compares new loss to best loss and updates patience counter.
    /// Returns true if this represents a significant improvement.
    ///
    /// # Arguments
    ///
    /// * `val_loss` - Current validation loss
    /// * `min_delta` - Minimum improvement threshold
    ///
    /// # Returns
    ///
    /// True if `best_val_loss - val_loss > min_delta`, false otherwise
    ///
    /// # Example
    ///
    /// ```rust
    /// # use ml::hyperopt::early_stopping::EarlyStoppingState;
    /// let mut state = EarlyStoppingState::new();
    ///
    /// assert!(state.update(0.5, 1e-4)); // First update
    /// assert!(state.update(0.3, 1e-4)); // Big improvement
    /// assert!(!state.update(0.299, 1e-4)); // Marginal (< 1e-4)
    /// ```
    pub fn update(&mut self, val_loss: f64, min_delta: f64) -> bool {
        // Check if this is a significant improvement
        let improvement = self.best_val_loss - val_loss;

        // Use epsilon comparison to avoid floating point precision issues
        // e.g., 0.5 - 0.49 = 0.010000000000000009 (not exactly 0.01)
        const EPSILON: f64 = 1e-10;
        if improvement > min_delta + EPSILON {
            // Significant improvement - update best and reset patience
            self.best_val_loss = val_loss;
            self.patience_counter = 0;
            true
        } else {
            // No significant improvement - increment patience
            self.patience_counter += 1;
            false
        }
    }

    /// Record epoch metrics for history tracking
    ///
    /// Stores metrics for later analysis, convergence plots, and debugging.
    ///
    /// # Arguments
    ///
    /// * `metrics` - Epoch metrics to record
    pub fn record_epoch_metrics(&mut self, metrics: EpochMetrics) {
        self.epoch_history.push(metrics);
    }

    /// Mark trial as stopped early
    ///
    /// # Arguments
    ///
    /// * `epoch` - Epoch where stopping occurred
    pub fn mark_stopped(&mut self, epoch: usize) {
        self.stopped = true;
        self.stopped_at_epoch = Some(epoch);
    }
}

impl Default for EarlyStoppingState {
    fn default() -> Self {
        Self::new()
    }
}

/// Metrics for a single training epoch
///
/// Contains all metrics needed for early stopping decisions and
/// result analysis. Serializable for experiment tracking.
///
/// ## Usage
///
/// Create after each epoch and pass to `TrialObserver::on_epoch_complete()`:
///
/// ```rust
/// use ml::hyperopt::early_stopping::EpochMetrics;
/// use std::time::{SystemTime, UNIX_EPOCH};
///
/// let metrics = EpochMetrics {
///     epoch: 10,
///     train_loss: 0.45,
///     val_loss: 0.52,
///     timestamp: SystemTime::now()
///         .duration_since(UNIX_EPOCH)
///         .unwrap()
///         .as_secs_f64(),
/// };
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EpochMetrics {
    /// Epoch number (0-indexed)
    pub epoch: usize,

    /// Training loss for this epoch
    pub train_loss: f64,

    /// Validation loss for this epoch (used for early stopping decisions)
    pub val_loss: f64,

    /// Unix timestamp when epoch completed
    pub timestamp: f64,
}

/// Available early stopping strategies
///
/// Each strategy implements a different approach to trial termination:
///
/// - **Single-trial**: Decisions based only on current trial (Plateau)
/// - **Cross-trial**: Decisions based on comparison with other trials (Median, Percentile, etc.)
///
/// ## Strategy Comparison
///
/// | Strategy | Type | Warmup | Aggressiveness | Use Case |
/// |----------|------|--------|----------------|----------|
/// | Plateau | Single | None | Low | Stable convergence |
/// | MedianPruner | Cross | Required | Medium | Balanced exploration |
/// | PercentilePruner | Cross | Required | High | Aggressive pruning |
/// | SuccessiveHalving | Cross | None | Very High | Known good ranges |
/// | Hyperband | Cross | None | Optimal | Large-scale search |
///
/// ## Example
///
/// ```rust
/// use ml::hyperopt::early_stopping::EarlyStoppingStrategy;
///
/// // Conservative - only stop on clear plateau
/// let plateau = EarlyStoppingStrategy::Plateau;
///
/// // Balanced - prune bottom half after warmup
/// let median = EarlyStoppingStrategy::MedianPruner { warmup_steps: 5 };
///
/// // Aggressive - prune bottom 25% after warmup
/// let percentile = EarlyStoppingStrategy::PercentilePruner {
///     percentile: 25.0,
///     warmup_steps: 10,
/// };
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum EarlyStoppingStrategy {
    /// Stop if no improvement for N epochs (single-trial)
    ///
    /// Most conservative strategy. Only considers current trial's
    /// own performance history. Good for:
    /// - Small-scale optimization (< 30 trials)
    /// - Unstable training dynamics
    /// - When you want every trial to reach natural convergence
    Plateau,

    /// Stop if worse than median of all trials (cross-trial)
    ///
    /// Compares current trial's loss to the median loss across all
    /// completed and running trials at the same epoch. Prunes trials
    /// in the bottom 50%.
    ///
    /// **Parameters**:
    /// - `warmup_steps`: Minimum epochs before pruning can occur
    ///
    /// **Best for**: Balanced exploration vs. exploitation
    MedianPruner {
        /// Minimum epochs before pruning can trigger
        warmup_steps: usize,
    },

    /// Stop if in bottom N percentile (cross-trial)
    ///
    /// More aggressive than median pruning. Prunes trials performing
    /// worse than the Nth percentile across all trials.
    ///
    /// **Parameters**:
    /// - `percentile`: Target percentile (0-100). Lower = more aggressive
    /// - `warmup_steps`: Minimum epochs before pruning
    ///
    /// **Example**: `percentile=25.0` prunes bottom 25% of trials
    ///
    /// **Best for**: Large-scale search (100+ trials) with tight budgets
    PercentilePruner {
        /// Target percentile threshold (0-100)
        percentile: f64,
        /// Minimum epochs before pruning
        warmup_steps: usize,
    },

    /// Successive halving algorithm (cross-trial)
    ///
    /// Aggressively halves the number of trials at each bracket.
    /// Allocates resources to best-performing trials only.
    ///
    /// **Parameters**:
    /// - `reduction_factor`: How many trials to keep (default: 2 = keep half)
    ///
    /// **Best for**: Known good parameter ranges, want to zoom in fast
    SuccessiveHalving {
        /// Reduction factor (2 = keep half, 3 = keep third, etc.)
        reduction_factor: usize,
    },

    /// Hyperband algorithm (cross-trial)
    ///
    /// Runs multiple successive halving brackets with different
    /// resource allocations. Provably optimal for any smooth function.
    ///
    /// **Parameters**:
    /// - `max_resource`: Maximum epochs per trial
    /// - `reduction_factor`: Halving rate (typically 2-4)
    ///
    /// **Best for**: Large-scale search with unknown convergence behavior
    Hyperband {
        /// Maximum resource allocation (epochs)
        max_resource: usize,
        /// Reduction factor for successive halving
        reduction_factor: usize,
    },
}

/// Trait for early stopping strategies
///
/// Implement this trait to create custom stopping logic. The trait is
/// called by `EarlyStoppingObserver` after each epoch to make decisions.
///
/// ## Contract
///
/// - Must be deterministic (same inputs → same output)
/// - Should be fast (called every epoch for every trial)
/// - Can access cross-trial history for comparison-based pruning
///
/// ## Example Implementation
///
/// ```rust,ignore
/// struct CustomStrategy {
///     threshold: f64,
/// }
///
/// impl EarlyStoppingStrategyTrait for CustomStrategy {
///     fn should_stop(
///         &self,
///         epoch: usize,
///         val_loss: f64,
///         state: &mut EarlyStoppingState,
///         trial_history: &[f64],
///     ) -> bool {
///         val_loss > self.threshold
///     }
///
///     fn name(&self) -> &'static str {
///         "custom_threshold"
///     }
/// }
/// ```
pub trait EarlyStoppingStrategyTrait {
    /// Check if trial should stop
    ///
    /// # Arguments
    ///
    /// * `epoch` - Current epoch number
    /// * `val_loss` - Current validation loss
    /// * `state` - Mutable trial state for updates
    /// * `trial_history` - Best losses from other trials (for cross-trial strategies)
    ///
    /// # Returns
    ///
    /// True if trial should stop, false to continue
    fn should_stop(
        &self,
        epoch: usize,
        val_loss: f64,
        state: &mut EarlyStoppingState,
        trial_history: &[f64],
    ) -> bool;

    /// Name of the strategy for logging
    fn name(&self) -> &'static str;
}

/// Plateau detection strategy implementation
///
/// Stops training when validation loss plateaus for `patience_epochs`
/// consecutive epochs without improvement >= `min_delta`.
///
/// ## Algorithm
///
/// ```text
/// for each epoch:
///     if val_loss improves by >= min_delta:
///         reset patience_counter = 0
///     else:
///         patience_counter += 1
///     
///     if patience_counter >= patience_epochs:
///         return STOP
/// ```
///
/// ## Performance
///
/// - Time: O(1) per epoch
/// - Memory: O(1)
/// - Overhead: ~1μs per decision
#[derive(Debug, Clone)]
pub struct PlateauDetectionStrategy {
    patience_epochs: usize,
    min_delta: f64,
}

impl PlateauDetectionStrategy {
    /// Create new plateau detection strategy
    ///
    /// # Arguments
    ///
    /// * `patience_epochs` - Number of epochs to wait for improvement
    /// * `min_delta` - Minimum improvement threshold
    pub fn new(patience_epochs: usize, min_delta: f64) -> Self {
        Self {
            patience_epochs,
            min_delta,
        }
    }
}

impl EarlyStoppingStrategyTrait for PlateauDetectionStrategy {
    fn should_stop(
        &self,
        _epoch: usize,
        val_loss: f64,
        state: &mut EarlyStoppingState,
        _trial_history: &[f64],
    ) -> bool {
        // Update state and check if improved
        state.update(val_loss, self.min_delta);

        // Stop if patience exhausted
        state.patience_counter >= self.patience_epochs
    }

    fn name(&self) -> &'static str {
        "plateau_detection"
    }
}

/// Median pruner strategy implementation
///
/// Prunes trials performing worse than the median across all trials
/// at the same epoch. Requires warmup period to accumulate statistics.
///
/// ## Algorithm
///
/// ```text
/// if epoch < warmup_steps:
///     return CONTINUE
///
/// median = compute_median(trial_history)
/// if val_loss > median:
///     return STOP
/// else:
///     return CONTINUE
/// ```
///
/// ## Performance
///
/// - Time: O(T log T) per decision, where T = number of trials
/// - Memory: O(1)
/// - Overhead: ~10-50μs per decision (depends on trial count)
///
/// ## Example
///
/// ```rust
/// use ml::hyperopt::early_stopping::MedianPrunerStrategy;
///
/// let pruner = MedianPrunerStrategy::new(5);
/// let trial_history = vec![0.3, 0.4, 0.5, 0.6, 0.7];
///
/// // Current trial worse than median (0.5)
/// assert!(pruner.should_stop(10, 0.8, &trial_history));
///
/// // Current trial better than median
/// assert!(!pruner.should_stop(10, 0.2, &trial_history));
/// ```
#[derive(Debug, Clone)]
pub struct MedianPrunerStrategy {
    warmup_steps: usize,
}

impl MedianPrunerStrategy {
    /// Create new median pruner strategy
    ///
    /// # Arguments
    ///
    /// * `warmup_steps` - Minimum epochs before pruning can occur
    pub fn new(warmup_steps: usize) -> Self {
        Self { warmup_steps }
    }

    /// Check if trial should stop based on median comparison
    ///
    /// # Arguments
    ///
    /// * `epoch` - Current epoch number
    /// * `val_loss` - Current validation loss
    /// * `trial_history` - Best losses from other trials
    ///
    /// # Returns
    ///
    /// True if current loss is worse than median, false otherwise
    pub fn should_stop(&self, epoch: usize, val_loss: f64, trial_history: &[f64]) -> bool {
        // Wait for warmup
        if epoch < self.warmup_steps {
            return false;
        }

        // Need at least 2 trials for meaningful comparison
        if trial_history.len() < 2 {
            return false;
        }

        // Compute median
        let mut sorted = trial_history.to_vec();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
        let median = if sorted.len() % 2 == 0 {
            let mid = sorted.len() / 2;
            (sorted[mid - 1] + sorted[mid]) / 2.0
        } else {
            sorted[sorted.len() / 2]
        };

        // Prune if worse than median
        val_loss > median
    }
}

/// Percentile pruner strategy implementation
///
/// Prunes trials in the bottom N percentile across all trials.
/// More aggressive than median pruning.
///
/// ## Algorithm
///
/// ```text
/// if epoch < warmup_steps:
///     return CONTINUE
///
/// threshold = compute_percentile(trial_history, percentile)
/// if val_loss > threshold:
///     return STOP
/// else:
///     return CONTINUE
/// ```
///
/// ## Performance
///
/// - Time: O(T log T) per decision
/// - Memory: O(1)
/// - Overhead: ~10-50μs per decision
///
/// ## Example
///
/// ```rust
/// use ml::hyperopt::early_stopping::PercentilePrunerStrategy;
///
/// // Prune bottom 25%
/// let pruner = PercentilePrunerStrategy::new(25.0, 5);
/// let trial_history = vec![0.2, 0.4, 0.6, 0.8];
///
/// // 25th percentile ≈ 0.35
/// assert!(pruner.should_stop(10, 0.9, &trial_history)); // Bottom 25%
/// assert!(!pruner.should_stop(10, 0.5, &trial_history)); // Above 25th
/// ```
#[derive(Debug, Clone)]
pub struct PercentilePrunerStrategy {
    percentile: f64,
    warmup_steps: usize,
}

impl PercentilePrunerStrategy {
    /// Create new percentile pruner strategy
    ///
    /// # Arguments
    ///
    /// * `percentile` - Target percentile (0-100). Lower = more aggressive
    /// * `warmup_steps` - Minimum epochs before pruning
    ///
    /// # Panics
    ///
    /// Panics if `percentile` is not in range [0, 100]
    pub fn new(percentile: f64, warmup_steps: usize) -> Self {
        assert!(
            (0.0..=100.0).contains(&percentile),
            "Percentile must be in [0, 100]"
        );
        Self {
            percentile,
            warmup_steps,
        }
    }

    /// Check if trial should stop based on percentile comparison
    ///
    /// # Arguments
    ///
    /// * `epoch` - Current epoch number
    /// * `val_loss` - Current validation loss
    /// * `trial_history` - Best losses from other trials
    ///
    /// # Returns
    ///
    /// True if current loss is in bottom percentile, false otherwise
    pub fn should_stop(&self, epoch: usize, val_loss: f64, trial_history: &[f64]) -> bool {
        // Wait for warmup
        if epoch < self.warmup_steps {
            return false;
        }

        // Need sufficient trials for meaningful percentile
        if trial_history.len() < 3 {
            return false;
        }

        // Compute percentile threshold
        let mut sorted = trial_history.to_vec();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

        let index = ((self.percentile / 100.0) * (sorted.len() as f64 - 1.0)).round() as usize;
        let threshold = sorted[index.min(sorted.len() - 1)];

        // Prune if worse than threshold
        val_loss > threshold
    }
}

/// Decision from observer callbacks
///
/// Returned by `TrialObserver::on_epoch_complete()` to control trial execution.
///
/// ## Decision Flow
///
/// ```text
/// on_epoch_complete() → ObserverDecision
///                            ↓
///    ┌──────────────────────┼──────────────────────┐
///    ↓                      ↓                      ↓
/// Continue           StopTrial               StopStudy
/// (keep going)     (stop this trial)     (stop all trials)
/// ```
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ObserverDecision {
    /// Continue training this trial
    Continue,

    /// Stop this trial early (save resources)
    StopTrial,

    /// Stop entire optimization study (found optimal or budget exhausted)
    StopStudy,
}

/// Trait for trial lifecycle observers
///
/// Implement this trait to receive callbacks during trial execution.
/// Used by adapters to integrate early stopping logic.
///
/// ## Callback Order
///
/// ```text
/// on_trial_start(0)
///   ↓
/// on_epoch_complete(0, 0) → Continue
///   ↓
/// on_epoch_complete(0, 1) → Continue
///   ↓
/// on_epoch_complete(0, 2) → StopTrial
///   ↓
/// on_trial_complete(0, final_loss)
/// ```
///
/// ## Example Implementation
///
/// See `EarlyStoppingObserver` for full implementation.
pub trait TrialObserver {
    /// Called when trial starts
    ///
    /// # Arguments
    ///
    /// * `trial_num` - Trial number (0-indexed)
    /// * `params` - Human-readable parameter string
    fn on_trial_start(&mut self, trial_num: usize, params: &str);

    /// Called after each training epoch
    ///
    /// # Arguments
    ///
    /// * `trial_num` - Trial number
    /// * `epoch` - Epoch number (0-indexed)
    /// * `metrics` - Epoch metrics (train/val loss, timestamp)
    ///
    /// # Returns
    ///
    /// Decision on whether to continue, stop trial, or stop study
    fn on_epoch_complete(
        &mut self,
        trial_num: usize,
        epoch: usize,
        metrics: &EpochMetrics,
    ) -> ObserverDecision;

    /// Called when trial completes successfully
    ///
    /// # Arguments
    ///
    /// * `trial_num` - Trial number
    /// * `final_loss` - Final validation loss
    fn on_trial_complete(&mut self, trial_num: usize, final_loss: f64);

    /// Called when trial fails with error
    ///
    /// # Arguments
    ///
    /// * `trial_num` - Trial number
    /// * `error` - Error message
    fn on_trial_failed(&mut self, trial_num: usize, error: &str);
}

/// Default early stopping observer implementation
///
/// Manages early stopping state for multiple trials and applies
/// configured stopping strategy. Implements `TrialObserver` trait
/// for integration with training loops.
///
/// ## State Management
///
/// - Maintains per-trial state in `HashMap<trial_id, State>`
/// - Tracks global best loss for cross-trial strategies
/// - Records complete epoch history for analysis
///
/// ## Memory Usage
///
/// - Fixed: ~128 bytes (observer overhead)
/// - Per-trial: ~48 + 48N bytes (N = epochs)
/// - Total for 30 trials × 100 epochs: ~144KB
///
/// ## Example
///
/// ```rust
/// use ml::hyperopt::early_stopping::{
///     EarlyStoppingConfig, EarlyStoppingObserver,
///     EarlyStoppingStrategy, TrialObserver, EpochMetrics,
/// };
///
/// let config = EarlyStoppingConfig {
///     patience_epochs: 10,
///     min_delta: 1e-4,
///     strategy: EarlyStoppingStrategy::Plateau,
///     ..Default::default()
/// };
///
/// let mut observer = EarlyStoppingObserver::new(config);
///
/// // In adapter's train_with_params()
/// observer.on_trial_start(0, "lr=0.001");
///
/// for epoch in 0..100 {
///     // ... training ...
///     let metrics = EpochMetrics {
///         epoch,
///         train_loss: 0.4,
///         val_loss: 0.5,
///         timestamp: epoch as f64,
///     };
///     
///     match observer.on_epoch_complete(0, epoch, &metrics) {
///         ml::hyperopt::early_stopping::ObserverDecision::StopTrial => break,
///         _ => continue,
///     }
/// }
///
/// observer.on_trial_complete(0, 0.5);
/// ```
#[derive(Debug)]
pub struct EarlyStoppingObserver {
    config: EarlyStoppingConfig,
    state: HashMap<usize, EarlyStoppingState>,
    baseline_val_loss: Option<f64>,
    trial_best_losses: Vec<f64>,
}

impl EarlyStoppingObserver {
    /// Create new early stopping observer
    ///
    /// # Arguments
    ///
    /// * `config` - Early stopping configuration
    pub fn new(config: EarlyStoppingConfig) -> Self {
        Self {
            config,
            state: HashMap::new(),
            baseline_val_loss: None,
            trial_best_losses: Vec::new(),
        }
    }

    /// Get number of trials tracked
    pub fn trial_count(&self) -> usize {
        self.state.len()
    }

    /// Get or create state for trial
    fn get_or_create_state(&mut self, trial_num: usize) -> &mut EarlyStoppingState {
        self.state
            .entry(trial_num)
            .or_default()
    }

    /// Check if should stop based on strategy
    fn should_stop_trial(&self, epoch: usize, val_loss: f64, patience_counter: usize) -> bool {
        // Always respect min_epochs
        if epoch < self.config.min_epochs {
            return false;
        }

        // Apply strategy-specific logic
        match &self.config.strategy {
            EarlyStoppingStrategy::Plateau => patience_counter >= self.config.patience_epochs,
            EarlyStoppingStrategy::MedianPruner { warmup_steps } => {
                let strategy = MedianPrunerStrategy::new(*warmup_steps);
                strategy.should_stop(epoch, val_loss, &self.trial_best_losses)
            },
            EarlyStoppingStrategy::PercentilePruner {
                percentile,
                warmup_steps,
            } => {
                let strategy = PercentilePrunerStrategy::new(*percentile, *warmup_steps);
                strategy.should_stop(epoch, val_loss, &self.trial_best_losses)
            },
            EarlyStoppingStrategy::SuccessiveHalving { reduction_factor } => {
                // Need at least reduction_factor completed trials for meaningful comparison
                if self.trial_best_losses.len() < *reduction_factor {
                    return false;
                }

                // Sort completed trial losses ascending (NaN-safe)
                let mut sorted = self.trial_best_losses.clone();
                sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

                // Keep top 1/eta fraction, prune the rest
                let keep_count = (sorted.len() / reduction_factor).max(1);
                let threshold = sorted[keep_count - 1];

                // Prune if current trial is worse than the threshold
                val_loss > threshold
            },
            EarlyStoppingStrategy::Hyperband {
                max_resource,
                reduction_factor,
            } => {
                // Hyperband: SHA pruning but ONLY at rung epochs
                // Rung epochs: max_resource / η^k for k = 1, 2, ...
                let eta = *reduction_factor;
                let max_r = *max_resource;

                // Check if this epoch is a rung epoch
                let mut is_rung = false;
                let mut rung_r = max_r / eta;
                while rung_r >= 1 {
                    if epoch == rung_r {
                        is_rung = true;
                        break;
                    }
                    rung_r /= eta;
                }

                if !is_rung {
                    return false; // Only prune at rung epochs
                }

                // At rung epoch, apply SHA pruning logic
                if self.trial_best_losses.len() < eta {
                    return false;
                }

                let mut sorted_losses = self.trial_best_losses.clone();
                sorted_losses.sort_by(|a, b| {
                    a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)
                });

                let keep_count = (sorted_losses.len() / eta).max(1);
                let threshold = sorted_losses[keep_count - 1];

                val_loss > threshold
            },
        }
    }
}

impl TrialObserver for EarlyStoppingObserver {
    fn on_trial_start(&mut self, trial_num: usize, params: &str) {
        info!(
            "Early Stopping: Trial {} started with params: {}",
            trial_num, params
        );
        self.get_or_create_state(trial_num);
    }

    fn on_epoch_complete(
        &mut self,
        trial_num: usize,
        epoch: usize,
        metrics: &EpochMetrics,
    ) -> ObserverDecision {
        let val_loss = metrics.val_loss;
        let min_delta = self.config.min_delta;
        let patience_epochs = self.config.patience_epochs;

        // Get state
        let state = self.get_or_create_state(trial_num);

        // Record metrics
        state.record_epoch_metrics(metrics.clone());

        // Update best loss
        let improved = state.update(val_loss, min_delta);

        if improved {
            info!(
                "Trial {}, Epoch {}: val_loss improved to {:.6} (patience reset)",
                trial_num, epoch, val_loss
            );
        }

        // Check stopping criteria
        // Extract patience_counter before checking (to avoid borrow issues)
        let patience_counter = state.patience_counter;

        // Drop mutable borrow of state before calling should_stop_trial
        let should_stop = self.should_stop_trial(epoch, val_loss, patience_counter);

        if should_stop {
            // Get state again to mark as stopped
            let state = self.get_or_create_state(trial_num);
            state.mark_stopped(epoch);
            info!(
                "Trial {}, Epoch {}: Early stopping triggered (patience: {}/{})",
                trial_num, epoch, patience_counter, patience_epochs
            );
            return ObserverDecision::StopTrial;
        }

        ObserverDecision::Continue
    }

    fn on_trial_complete(&mut self, trial_num: usize, final_loss: f64) {
        info!(
            "Early Stopping: Trial {} completed with final loss: {:.6}",
            trial_num, final_loss
        );

        // Update global statistics
        self.trial_best_losses.push(final_loss);

        if self.baseline_val_loss.is_none() || final_loss < self.baseline_val_loss.unwrap_or(f64::INFINITY) {
            self.baseline_val_loss = Some(final_loss);
            info!("New baseline loss: {:.6}", final_loss);
        }
    }

    fn on_trial_failed(&mut self, trial_num: usize, error: &str) {
        warn!("Early Stopping: Trial {} failed: {}", trial_num, error);
    }
}

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

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

        // First update always improves
        assert!(state.update(1.0, 0.01));
        assert_eq!(state.best_val_loss, 1.0);
        assert_eq!(state.patience_counter, 0);

        // Significant improvement (> 0.01)
        assert!(state.update(0.5, 0.01));
        assert_eq!(state.best_val_loss, 0.5);
        assert_eq!(state.patience_counter, 0);

        // Marginal improvement (< 0.01)
        assert!(!state.update(0.49, 0.01));
        assert_eq!(state.best_val_loss, 0.5);
        assert_eq!(state.patience_counter, 1);

        // No improvement
        assert!(!state.update(0.6, 0.01));
        assert_eq!(state.best_val_loss, 0.5);
        assert_eq!(state.patience_counter, 2);
    }

    #[test]
    fn test_plateau_detection_basic() {
        let strategy = PlateauDetectionStrategy::new(3, 0.01);
        let mut state = EarlyStoppingState::new();

        // Improving - should not stop
        assert!(!strategy.should_stop(0, 1.0, &mut state, &[]));
        assert!(!strategy.should_stop(1, 0.5, &mut state, &[]));

        // Plateau - should stop after patience
        assert!(!strategy.should_stop(2, 0.5, &mut state, &[]));
        assert!(!strategy.should_stop(3, 0.5, &mut state, &[]));
        assert!(strategy.should_stop(4, 0.5, &mut state, &[]));
    }

    #[test]
    fn test_median_pruner_basic() {
        let strategy = MedianPrunerStrategy::new(0);
        let history = vec![0.3, 0.4, 0.5, 0.6, 0.7]; // median = 0.5

        // Worse than median
        assert!(strategy.should_stop(5, 0.8, &history));

        // Better than median
        assert!(!strategy.should_stop(5, 0.2, &history));

        // Equal to median
        assert!(!strategy.should_stop(5, 0.5, &history));
    }

    #[test]
    fn test_percentile_pruner_basic() {
        let strategy = PercentilePrunerStrategy::new(25.0, 0);
        let history = vec![0.2, 0.4, 0.6, 0.8]; // 25th = 0.4

        // Bottom 25%
        assert!(strategy.should_stop(5, 0.9, &history));

        // Above 25th percentile (0.3 < 0.4, so not pruned)
        assert!(!strategy.should_stop(5, 0.3, &history));
    }

    #[test]
    fn test_successive_halving_prunes_bottom_trials() {
        // 9 completed trials with losses 0.1..=0.9, reduction_factor=3
        // keep_count = max(1, 9/3) = 3
        // sorted: [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
        // threshold = sorted[2] = 0.3
        // New trial with val_loss=0.5 > 0.3 => should be pruned (StopTrial)
        let config = EarlyStoppingConfig {
            patience_epochs: 10,
            min_delta: 1e-4,
            min_epochs: 0,
            strategy: EarlyStoppingStrategy::SuccessiveHalving {
                reduction_factor: 3,
            },
            ..Default::default()
        };
        let mut observer = EarlyStoppingObserver::new(config);

        // Register 9 completed trials
        for i in 0..9 {
            observer.on_trial_start(i, &format!("trial_{}", i));
            let loss = (i as f64 + 1.0) / 10.0; // 0.1, 0.2, ..., 0.9
            observer.on_trial_complete(i, loss);
        }

        // Start a new trial (trial 9) with a bad loss
        observer.on_trial_start(9, "trial_9");
        let metrics = EpochMetrics {
            epoch: 1,
            train_loss: 0.5,
            val_loss: 0.5, // worse than threshold 0.3
            timestamp: 1.0,
        };
        let decision = observer.on_epoch_complete(9, 1, &metrics);
        assert_eq!(
            decision,
            ObserverDecision::StopTrial,
            "Trial with val_loss=0.5 should be pruned (threshold=0.3)"
        );
    }

    #[test]
    fn test_successive_halving_keeps_top_trials() {
        // 9 completed trials with losses 0.1..=0.9, reduction_factor=3
        // keep_count = max(1, 9/3) = 3
        // sorted: [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
        // threshold = sorted[2] = 0.3
        // New trial with val_loss=0.2 <= 0.3 => should continue
        let config = EarlyStoppingConfig {
            patience_epochs: 10,
            min_delta: 1e-4,
            min_epochs: 0,
            strategy: EarlyStoppingStrategy::SuccessiveHalving {
                reduction_factor: 3,
            },
            ..Default::default()
        };
        let mut observer = EarlyStoppingObserver::new(config);

        // Register 9 completed trials
        for i in 0..9 {
            observer.on_trial_start(i, &format!("trial_{}", i));
            let loss = (i as f64 + 1.0) / 10.0; // 0.1, 0.2, ..., 0.9
            observer.on_trial_complete(i, loss);
        }

        // Start a new trial (trial 9) with a good loss
        observer.on_trial_start(9, "trial_9");
        let metrics = EpochMetrics {
            epoch: 1,
            train_loss: 0.2,
            val_loss: 0.2, // within top 1/3 (threshold=0.3)
            timestamp: 1.0,
        };
        let decision = observer.on_epoch_complete(9, 1, &metrics);
        assert_eq!(
            decision,
            ObserverDecision::Continue,
            "Trial with val_loss=0.2 should continue (threshold=0.3)"
        );
    }

    #[test]
    fn test_hyperband_prunes_at_rung_epochs() {
        // max_resource=81, reduction_factor=3
        // Rung epochs: 81/3=27, 81/9=9, 81/27=3, 81/81=1
        let config = EarlyStoppingConfig {
            patience_epochs: 100,
            min_delta: 1e-4,
            min_epochs: 0,
            strategy: EarlyStoppingStrategy::Hyperband {
                max_resource: 81,
                reduction_factor: 3,
            },
            ..Default::default()
        };
        let mut observer = EarlyStoppingObserver::new(config);

        // Seed with 9 completed trials
        for trial in 0..9 {
            observer.on_trial_start(trial, &format!("trial_{}", trial));
            observer.on_trial_complete(trial, (trial as f64 + 1.0) * 0.1);
        }

        // Trial 9 with bad loss at rung epoch 27 (81/3=27)
        observer.on_trial_start(9, "trial_9");
        let metrics = EpochMetrics {
            epoch: 27,
            train_loss: 0.9,
            val_loss: 0.95,
            timestamp: 27.0,
        };
        let decision = observer.on_epoch_complete(9, 27, &metrics);
        assert_eq!(
            decision,
            ObserverDecision::StopTrial,
            "Hyperband should prune at rung epoch 27 (val_loss=0.95 > threshold)"
        );
    }

    #[test]
    fn test_hyperband_does_not_prune_between_rungs() {
        // max_resource=81, reduction_factor=3
        // Rung epochs: 27, 9, 3, 1 -- epoch 15 is NOT a rung
        let config = EarlyStoppingConfig {
            patience_epochs: 100,
            min_delta: 1e-4,
            min_epochs: 0,
            strategy: EarlyStoppingStrategy::Hyperband {
                max_resource: 81,
                reduction_factor: 3,
            },
            ..Default::default()
        };
        let mut observer = EarlyStoppingObserver::new(config);

        // Seed with 9 completed trials
        for trial in 0..9 {
            observer.on_trial_start(trial, &format!("trial_{}", trial));
            observer.on_trial_complete(trial, (trial as f64 + 1.0) * 0.1);
        }

        // Trial 9 with bad loss but NOT at a rung epoch (epoch 15)
        observer.on_trial_start(9, "trial_9");
        let metrics = EpochMetrics {
            epoch: 15,
            train_loss: 0.9,
            val_loss: 0.95,
            timestamp: 15.0,
        };
        let decision = observer.on_epoch_complete(9, 15, &metrics);
        assert_eq!(
            decision,
            ObserverDecision::Continue,
            "Hyperband should NOT prune between rung epochs (epoch 15 is not a rung)"
        );
    }
}
