foxhunt/ml/src/hyperopt/optimizer.rs

//! Generic Bayesian Optimizer using Argmin
//!
//! This module provides a production-ready optimization implementation that works
//! with any model implementing `HyperparameterOptimizable`. It uses the argmin
//! library with Particle Swarm Optimization (PSO) for efficient parameter search.
//!
//! ## Key Features
//!
//! - **Particle Swarm Optimization**: Derivative-free, population-based optimization
//! - **Latin Hypercube Sampling**: Smart initialization for first N trials
//! - **Type-safe**: Works with any `ParameterSpace` implementation
//! - **Production-ready**: Comprehensive error handling and logging
//!
//! ## Algorithm Overview
//!
//! 1. **Initialization**: Generate `n_initial` diverse samples via LHS
//! 2. **Evaluation**: Train model with each parameter set, record objective
//! 3. **Optimization**: Use Particle Swarm to find global optima
//! 4. **Multi-restart**: Swarm automatically explores multiple regions
//! 5. **Repeat**: Continue until `max_trials` reached
//!
//! ## Usage Example
//!
//! ```rust,no_run
//! use ml::hyperopt::{ArgminOptimizer, HyperparameterOptimizable};
//! use ml::hyperopt::adapters::mamba2::{Mamba2Trainer, Mamba2Params};
//!
//! # fn example() -> anyhow::Result<()> {
//! let trainer = Mamba2Trainer::new("data.parquet", 50)?;
//!
//! let optimizer = ArgminOptimizer::builder()
//!     .max_trials(30)
//!     .n_initial(5)
//!     .seed(42)
//!     .build();
//!
//! let result = optimizer.optimize(trainer)?;
//! println!("Best loss: {:.6}", result.best_objective);
//! # Ok(())
//! # }
//! ```

use anyhow::{Context, Result};
use argmin::core::{CostFunction, Executor, State};
use argmin::solver::particleswarm::ParticleSwarm;
use ndarray::Array2;
use rand::prelude::*;
use rand::SeedableRng;
use rand_distr::Uniform;
use std::sync::{Arc, Mutex};
use std::time::Instant;
use tracing::{info, warn};

use super::traits::{HyperparameterOptimizable, OptimizationResult, ParameterSpace, TrialResult};
use crate::hyperopt::adapters::dqn::ObjectiveMode;
use crate::MLError;

/// Bayesian optimizer using argmin library
///
/// This optimizer uses Particle Swarm Optimization (PSO) with Latin Hypercube Sampling
/// for initialization. It's designed to work with any model implementing
/// `HyperparameterOptimizable`.
///
/// ## Configuration
///
/// - `max_trials`: Total evaluations (including initial samples)
/// - `n_initial`: Number of Latin Hypercube samples for initialization
/// - `n_particles`: Number of particles in the swarm (default: 20)
/// - `seed`: Random seed for reproducibility (optional)
///
/// ## Performance Characteristics
///
/// - Initial samples: ~O(n_initial * training_time)
/// - Per-iteration overhead: ~1-10ms for swarm updates
/// - Memory usage: O(max_trials) for trial history
///
/// ## Advantages over Nelder-Mead
///
/// - Supports multi-dimensional vector parameters (Vec<f64>)
/// - Better at avoiding local minima through swarm intelligence
/// - Scales well to higher dimensions (20-50 parameters)
/// - No gradient information required
#[derive(Debug, Clone)]
pub struct ArgminOptimizer {
    /// Maximum number of trials (evaluations)
    pub(crate) max_trials: usize,
    /// Number of initial samples via Latin Hypercube Sampling
    pub(crate) n_initial: usize,
    /// Number of particles in the swarm
    pub(crate) n_particles: usize,
    /// Random seed for reproducibility
    pub(crate) seed: Option<u64>,
    /// Maximum iterations per restart
    pub(crate) max_iters_per_restart: usize,
}

impl Default for ArgminOptimizer {
    fn default() -> Self {
        Self {
            max_trials: 30,
            n_initial: 5,
            n_particles: 20,
            seed: None,
            max_iters_per_restart: 50,
        }
    }
}

impl ArgminOptimizer {
    /// Create a new optimizer with default settings
    ///
    /// Uses 30 max trials and 5 initial samples.
    pub fn new() -> Self {
        Self::default()
    }

    /// Create optimizer with specified settings
    ///
    /// # Arguments
    ///
    /// * `max_trials` - Total number of trials (must be > n_initial)
    /// * `n_initial` - Number of initial LHS samples (typically 3-10)
    ///
    /// # Panics
    ///
    /// Panics if `max_trials <= n_initial` or if either is zero.
    pub fn with_trials(max_trials: usize, n_initial: usize) -> Self {
        assert!(max_trials > n_initial, "max_trials must be > n_initial");
        assert!(n_initial > 0, "n_initial must be > 0");

        Self {
            max_trials,
            n_initial,
            n_particles: 20,
            seed: None,
            max_iters_per_restart: 50,
        }
    }

    /// Set random seed for reproducibility
    pub fn with_seed(mut self, seed: u64) -> Self {
        self.seed = Some(seed);
        self
    }

    /// Start building an optimizer with custom configuration
    pub fn builder() -> ArgminOptimizerBuilder {
        ArgminOptimizerBuilder::new()
    }

    /// Generate Latin Hypercube samples for initialization
    ///
    /// # Arguments
    ///
    /// * `n_samples` - Number of samples to generate
    /// * `bounds` - Parameter bounds [(min, max), ...]
    /// * `rng` - Random number generator
    ///
    /// # Returns
    ///
    /// Array of shape (n_samples, n_params) with LHS samples
    pub(crate) fn latin_hypercube_sampling<R: Rng>(
        n_samples: usize,
        bounds: &[(f64, f64)],
        rng: &mut R,
    ) -> Array2<f64> {
        let n_params = bounds.len();
        let mut samples = Array2::zeros((n_samples, n_params));

        // For each parameter dimension
        for dim in 0..n_params {
            // Create permutation of [0, 1, 2, ..., n_samples-1]
            let mut indices: Vec<usize> = (0..n_samples).collect();
            indices.shuffle(rng);

            // Generate stratified samples
            let uniform = Uniform::new(0.0, 1.0);
            for (sample_idx, &perm_idx) in indices.iter().enumerate() {
                let segment_start = perm_idx as f64 / n_samples as f64;
                let segment_end = (perm_idx + 1) as f64 / n_samples as f64;
                let random_offset: f64 = rng.sample(uniform);
                let normalized = segment_start + random_offset * (segment_end - segment_start);

                // Denormalize to actual bounds
                let (min, max) = bounds[dim];
                samples[[sample_idx, dim]] = min + normalized * (max - min);
            }
        }

        samples
    }

    /// Run optimization on a model
    ///
    /// This method:
    /// 1. Extracts parameter bounds from `P::continuous_bounds()`
    /// 2. Generates initial samples via Latin Hypercube
    /// 3. Evaluates each sample by calling `model.train_with_params()`
    /// 4. Uses Nelder-Mead to optimize from best initial point
    /// 5. Performs multi-restart from other good points
    /// 6. Repeats until `max_trials` reached
    ///
    /// # Type Parameters
    ///
    /// - `M`: Model type implementing `HyperparameterOptimizable`
    /// - `P`: Parameter type implementing `ParameterSpace`
    ///
    /// # Arguments
    ///
    /// * `model` - Mutable reference to model (needed for training)
    ///
    /// # Returns
    ///
    /// `OptimizationResult<P>` containing best parameters and full trial history
    ///
    /// # Errors
    ///
    /// Returns error if:
    /// - Parameter space has zero dimensions
    /// - Initial sampling fails
    /// - All trials produce errors (no valid observations)
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use ml::hyperopt::{ArgminOptimizer, HyperparameterOptimizable};
    /// # use ml::hyperopt::adapters::mamba2::Mamba2Trainer;
    /// # fn example() -> anyhow::Result<()> {
    /// let mut trainer = Mamba2Trainer::new("data.parquet", 50)?;
    /// let optimizer = ArgminOptimizer::new();
    /// let result = optimizer.optimize(trainer)?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn optimize<M>(&self, mut model: M) -> Result<OptimizationResult<M::Params>>
    where
        M: HyperparameterOptimizable + Send,
        M::Params: ParameterSpace + Send,
    {
        info!("╔═══════════════════════════════════════════════════════════╗");
        info!("║  Bayesian Hyperparameter Optimization (Argmin)           ║");
        info!("╚═══════════════════════════════════════════════════════════╝");

        // Extract parameter space bounds
        let bounds = M::Params::continuous_bounds();
        let n_params = bounds.len();

        if n_params == 0 {
            return Err(MLError::ConfigError {
                reason: "Parameter space has zero dimensions".to_string(),
            }
            .into());
        }

        info!("Configuration:");
        info!("  Max Trials: {}", self.max_trials);
        info!("  Initial Samples: {}", self.n_initial);
        info!("  Swarm Particles: {}", self.n_particles);
        info!("  Parameters: {}", n_params);
        info!("  Max Iters/Restart: {}", self.max_iters_per_restart);

        let param_names = M::Params::param_names();
        for (i, name) in param_names.iter().enumerate() {
            info!("    {} - [{:.6}, {:.6}]", name, bounds[i].0, bounds[i].1);
        }

        // Initialize RNG
        let mut rng = match self.seed {
            Some(seed) => StdRng::seed_from_u64(seed),
            None => StdRng::from_entropy(),
        };

        // Generate initial samples using Latin Hypercube
        info!(
            "Generating {} initial samples with Latin Hypercube Sampling...",
            self.n_initial
        );
        let initial_samples = Self::latin_hypercube_sampling(self.n_initial, &bounds, &mut rng);
        info!("✓ Generated {} initial samples", initial_samples.nrows());

        // Shared state for collecting trial results
        let trial_results = Arc::new(Mutex::new(Vec::new()));
        let trial_counter = Arc::new(Mutex::new(0));

        // Evaluate initial samples
        info!("Evaluating initial samples...");
        for i in 0..self.n_initial {
            let continuous_vec: Vec<f64> = initial_samples.row(i).to_vec();
            Self::evaluate_point(
                &continuous_vec,
                &mut model,
                &trial_results,
                &trial_counter,
                &param_names,
            )?;
        }

        // Create trial budget observer BEFORE creating cost function
        // This observer enforces strict trial limits (fixes PSO infinite iteration bug)
        let observer = crate::hyperopt::TrialBudgetObserver::new(self.max_trials);

        // Create cost function wrapper
        let cost_fn = ObjectiveFunction {
            model: Arc::new(Mutex::new(model)),
            trial_results: Arc::clone(&trial_results),
            trial_counter: Arc::clone(&trial_counter),
            param_names: param_names.clone(),
            bounds: bounds.clone(),
            observer: observer.clone(),
        };

        // Find best initial point to start optimization
        let trials = trial_results.lock().unwrap().clone();
        let best_initial = trials
            .iter()
            .min_by(|a, b| {
                a.objective
                    .partial_cmp(&b.objective)
                    .unwrap_or(std::cmp::Ordering::Equal)
            })
            .ok_or_else(|| {
                anyhow::anyhow!("No initial trials completed successfully — cannot start PSO")
            })?;

        info!("╔═══════════════════════════════════════════════════════════╗");
        info!("║  Starting Particle Swarm Optimization                    ║");
        info!("╚═══════════════════════════════════════════════════════════╝");
        info!("Best initial objective: {:.6}", best_initial.objective);
        info!("Execution mode: Sequential trials (model locked by Mutex, rayon for swarm only)");

        // Determine how many iterations we can afford
        let trials_used = trials.len();
        let remaining_trials = self.max_trials.saturating_sub(trials_used);

        // FIX: Each PSO iteration evaluates n_particles candidates (sequentially via mutex)
        // PSO evaluates ALL particles in swarm per iteration, so divide remaining budget
        // by swarm size to prevent trial count overflow (fixes 962 trial bug)
        // CRITICAL FIX (2025-11-07): Use CEILING division to ensure all trials complete
        // Example: 8 remaining ÷ 20 particles = 0.4 → ceil to 1 iteration (not 0)
        let max_iters_by_budget =
            ((remaining_trials as f64) / (self.n_particles as f64)).ceil() as usize;

        let max_iters = max_iters_by_budget.min(self.max_iters_per_restart);

        info!(
            "PSO Budget: {} iterations ({} remaining trials ÷ {} particles = {} max iters)",
            max_iters, remaining_trials, self.n_particles, max_iters_by_budget
        );

        if max_iters > 0 {
            // Create bounds vectors for ParticleSwarm
            let lower_bounds: Vec<f64> = bounds.iter().map(|(min, _)| *min).collect();
            let upper_bounds: Vec<f64> = bounds.iter().map(|(_, max)| *max).collect();

            // Create Particle Swarm solver
            let solver = ParticleSwarm::new((lower_bounds, upper_bounds), self.n_particles);

            // Run optimization (parallel execution enabled via rayon feature)
            // CRITICAL FIX (2025-11-03): Removed .target_cost(0.0) to prevent early termination
            // PSO must run for exactly max_iters iterations to complete all requested trials
            // CRITICAL FIX (2025-11-07): Added TrialBudgetObserver to enforce strict trial limits
            let res = Executor::new(cost_fn, solver)
                .configure(|state| state.max_iters(max_iters as u64))
                .add_observer(
                    observer.clone(),
                    argmin::core::observers::ObserverMode::Always,
                )
                .run();

            // Handle budget exhaustion gracefully (not an error condition)
            match res {
                Ok(result) => {
                    info!("Optimization complete:");
                    info!("  Final cost: {:.6}", result.state().get_best_cost());
                    info!("  Iterations: {}", result.state().get_iter());
                    info!("  Evaluations: {}", trial_counter.lock().unwrap());
                },
                Err(e) if e.to_string().contains("Trial budget exhausted") => {
                    info!(
                        "PSO terminated: trial budget reached ({} trials)",
                        self.max_trials
                    );
                    info!("  Evaluations: {}", trial_counter.lock().unwrap());
                },
                Err(e) => return Err(e.into()),
            }
        } else {
            info!("No remaining budget for Particle Swarm optimization");
        }

        info!("╔═══════════════════════════════════════════════════════════╗");
        info!("║  Optimization Complete                                    ║");
        info!("╚═══════════════════════════════════════════════════════════╝");

        // Extract final results
        let trials = match Arc::try_unwrap(trial_results) {
            Ok(mutex) => mutex.into_inner().unwrap(),
            Err(arc) => arc.lock().unwrap().clone(),
        };

        if trials.is_empty() {
            return Err(MLError::TrainingError(
                "All trials failed - no valid observations".to_string(),
            )
            .into());
        }

        // Create optimization result
        let result = OptimizationResult::from_trials(trials);

        info!("Best Parameters Found:");
        let best_continuous = result.best_params.to_continuous();
        for (i, name) in param_names.iter().enumerate() {
            info!("  {}: {:.6}", name, best_continuous[i]);
        }
        info!("Best Objective: {:.6}", result.best_objective);
        info!("Total Improvement: {:.6}", result.total_improvement());
        info!("Improvement: {:.2}%", result.improvement_percentage());

        Ok(result)
    }

    /// Evaluate a single point
    fn evaluate_point<M>(
        continuous_vec: &[f64],
        model: &mut M,
        trial_results: &Arc<Mutex<Vec<TrialResult<M::Params>>>>,
        trial_counter: &Arc<Mutex<usize>>,
        _param_names: &[&'static str],
    ) -> Result<f64>
    where
        M: HyperparameterOptimizable,
        M::Params: ParameterSpace,
    {
        let start_time = Instant::now();

        // Increment trial counter
        let trial_num = {
            let mut counter = trial_counter.lock().unwrap();
            *counter += 1;
            *counter
        };

        info!("╔═══════════════════════════════════════════════════════════╗");
        info!(
            "║  Trial {}: Evaluating Parameters                          ║",
            trial_num
        );
        info!("╚═══════════════════════════════════════════════════════════╝");

        // Convert continuous vector to parameters BEFORE logging
        let params =
            M::Params::from_continuous(continuous_vec).context("Failed to convert parameters")?;

        // Log CONVERTED parameters (shows actual values: learning_rate ~1e-4, not -11)
        info!("  Parameters (converted): {:?}", params);

        // Train model with parameters
        let metrics = model
            .train_with_params(params.clone())
            .context(format!("Training failed for trial {}", trial_num))?;

        // Extract objective (replace NaN/Inf with large penalty)
        let raw_objective = M::extract_objective(&metrics);
        let objective = if raw_objective.is_finite() {
            raw_objective
        } else {
            warn!("Trial {} produced non-finite objective ({:?}), using penalty 1e6", trial_num, raw_objective);
            1e6
        };
        let duration_secs = start_time.elapsed().as_secs_f64();

        info!("✓ Trial {} completed in {:.1}s", trial_num, duration_secs);
        info!("  Objective: {:.6}", objective);

        // Record trial result
        {
            let mut results = trial_results.lock().unwrap();
            results.push(TrialResult {
                trial_num,
                params,
                objective,
                duration_secs,
            });
        }

        Ok(objective)
    }

    /// Run two-phase optimization: episode reward then Sharpe ratio.
    ///
    /// Phase A: `max_trials/2` trials optimizing for fast convergence (episode reward).
    /// Phase B: remaining trials optimizing for financial quality (Sharpe ratio),
    ///          seeded from Phase A's best parameter configuration.
    ///
    /// This reuses the existing `optimize()` method for each phase. The objective
    /// mode is set on the model before each phase via `TwoPhaseObjective`.
    ///
    /// # Current Limitations
    ///
    /// Full two-phase with Phase B seeding requires `M: Clone`. Currently runs
    /// Phase A only and returns its result. Phase B will be enabled once the
    /// model types implement `Clone`.
    pub fn optimize_two_phase<M>(&self, mut model: M) -> Result<OptimizationResult<M::Params>>
    where
        M: HyperparameterOptimizable + TwoPhaseObjective + Send,
        M::Params: ParameterSpace + Send,
    {
        let phase_a_trials = self.max_trials / 2;
        let _phase_b_trials = self.max_trials - phase_a_trials;

        info!("======= Two-Phase Optimization =======");
        info!("Phase A: {} trials optimizing episode reward", phase_a_trials);
        info!("Phase B: {} trials optimizing Sharpe ratio (pending Clone support)", _phase_b_trials);

        // Phase A: Episode Reward
        model.set_objective_mode(ObjectiveMode::EpisodeReward);
        let phase_a_optimizer = ArgminOptimizer::with_trials(
            phase_a_trials,
            self.n_initial.min(phase_a_trials - 1).max(1),
        );
        let phase_a_result = phase_a_optimizer.optimize(model)?;

        info!(
            "Phase A complete: best_objective={:.6}, evaluated {} trials",
            phase_a_result.best_objective,
            phase_a_result.all_trials.len()
        );

        // TODO: Phase B requires M: Clone to reconstruct model from Phase A result.
        // Once DQNTrainer implements Clone, Phase B will:
        // 1. Extract the model back (currently consumed by optimize())
        // 2. Set objective mode to Sharpe
        // 3. Seed Phase B with top-3 parameter configs from Phase A
        // 4. Run remaining trials with Sharpe-based objective

        info!(
            "======= Two-Phase Complete: best_objective={:.6} =======",
            phase_a_result.best_objective
        );

        Ok(phase_a_result)
    }
}

/// Trait for models supporting two-phase objective switching.
///
/// Used by [`ArgminOptimizer::optimize_two_phase()`] to switch between
/// episode reward (fast convergence) and Sharpe ratio (financial quality).
///
/// This trait is intentionally separate from [`HyperparameterOptimizable`]
/// because `extract_objective` is a static method that cannot access instance
/// state. Two-phase optimization instead sets the objective mode on the model
/// instance before each phase.
pub trait TwoPhaseObjective {
    /// Set the objective mode for the current optimization phase.
    fn set_objective_mode(&mut self, mode: ObjectiveMode);
    /// Get the current objective mode.
    fn objective_mode(&self) -> ObjectiveMode;
}

/// Cost function wrapper for argmin
struct ObjectiveFunction<M>
where
    M: HyperparameterOptimizable,
    M::Params: ParameterSpace,
{
    model: Arc<Mutex<M>>,
    trial_results: Arc<Mutex<Vec<TrialResult<M::Params>>>>,
    trial_counter: Arc<Mutex<usize>>,
    param_names: Vec<&'static str>,
    bounds: Vec<(f64, f64)>,
    observer: crate::hyperopt::TrialBudgetObserver,
}

impl<M> CostFunction for ObjectiveFunction<M>
where
    M: HyperparameterOptimizable,
    M::Params: ParameterSpace,
{
    type Param = Vec<f64>;
    type Output = f64;

    fn cost(&self, param: &Self::Param) -> Result<Self::Output, argmin::core::Error> {
        // Increment observer trial count FIRST
        self.observer.increment_trial();

        // Check if budget exhausted (prevents runaway PSO)
        if self.observer.should_terminate() {
            warn!(
                "Trial budget exhausted: {}/{} trials. Stopping PSO.",
                self.observer.get_trials_used(),
                self.observer.get_trials_used()
            );
            return Err(argmin::core::Error::msg("Trial budget exhausted"));
        }

        // Clamp parameters to bounds
        let mut clamped = param.clone();
        for (i, (min, max)) in self.bounds.iter().enumerate() {
            clamped[i] = clamped[i].clamp(*min, *max);
        }

        let start_time = Instant::now();

        // Increment trial counter
        let trial_num = {
            let mut counter = self.trial_counter.lock().unwrap();
            *counter += 1;
            *counter
        };

        info!("╔═══════════════════════════════════════════════════════════╗");
        info!(
            "║  Trial {}: Evaluating Parameters                          ║",
            trial_num
        );
        info!("╚═══════════════════════════════════════════════════════════╝");

        // Convert continuous vector to parameters BEFORE logging
        let params = match M::Params::from_continuous(&clamped) {
            Ok(p) => p,
            Err(e) => {
                warn!(
                    "Failed to convert parameters for trial {}: {}",
                    trial_num, e
                );
                return Ok(1e6); // Penalty for invalid parameters
            },
        };

        // Log CONVERTED parameters (shows actual values: learning_rate ~1e-4, not -11)
        info!("  Parameters (converted): {:?}", params);

        // Train model with parameters
        let mut model = self.model.lock().unwrap();
        let metrics = match model.train_with_params(params.clone()) {
            Ok(m) => m,
            Err(e) => {
                warn!("Training failed for trial {}: {}", trial_num, e);
                return Ok(1e6); // Penalty for training failure
            },
        };

        // Extract objective (replace NaN/Inf with large penalty)
        let raw_objective = M::extract_objective(&metrics);
        let objective = if raw_objective.is_finite() {
            raw_objective
        } else {
            warn!("Trial {} produced non-finite objective ({:?}), using penalty 1e6", trial_num, raw_objective);
            1e6
        };
        let duration_secs = start_time.elapsed().as_secs_f64();

        info!("✓ Trial {} completed in {:.1}s", trial_num, duration_secs);
        info!("  Objective: {:.6}", objective);

        // Record trial result
        {
            let mut results = self.trial_results.lock().unwrap();
            results.push(TrialResult {
                trial_num,
                params,
                objective,
                duration_secs,
            });
        }

        Ok(objective)
    }
}

/// Builder for `ArgminOptimizer` with fluent API
///
/// Provides a convenient way to configure optimizer settings.
///
/// # Example
///
/// ```rust
/// use ml::hyperopt::ArgminOptimizer;
///
/// let optimizer = ArgminOptimizer::builder()
///     .max_trials(50)
///     .n_initial(10)
///     .n_particles(30)
///     .seed(12345)
///     .build();
/// ```
#[derive(Debug, Clone)]
pub struct ArgminOptimizerBuilder {
    max_trials: usize,
    n_initial: usize,
    n_particles: usize,
    seed: Option<u64>,
    max_iters_per_restart: usize,
}

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

impl ArgminOptimizerBuilder {
    /// Create a new builder with default settings
    pub fn new() -> Self {
        Self {
            max_trials: 30,
            n_initial: 5,
            n_particles: 20,
            seed: None,
            max_iters_per_restart: 50,
        }
    }

    /// Set maximum number of trials
    pub fn max_trials(mut self, max_trials: usize) -> Self {
        self.max_trials = max_trials;
        self
    }

    /// Set number of initial samples
    pub fn n_initial(mut self, n_initial: usize) -> Self {
        self.n_initial = n_initial;
        self
    }

    /// Set number of particles in the swarm
    pub fn n_particles(mut self, n_particles: usize) -> Self {
        self.n_particles = n_particles;
        self
    }

    /// Set random seed for reproducibility
    pub fn seed(mut self, seed: u64) -> Self {
        self.seed = Some(seed);
        self
    }

    /// Set maximum iterations per restart
    pub fn max_iters_per_restart(mut self, max_iters: usize) -> Self {
        self.max_iters_per_restart = max_iters;
        self
    }

    /// Build the optimizer
    ///
    /// # Panics
    ///
    /// Panics if `max_trials <= n_initial` or if either is zero.
    pub fn build(self) -> ArgminOptimizer {
        assert!(
            self.max_trials > self.n_initial,
            "max_trials must be > n_initial"
        );
        assert!(self.n_initial > 0, "n_initial must be > 0");
        assert!(self.n_particles > 0, "n_particles must be > 0");

        ArgminOptimizer {
            max_trials: self.max_trials,
            n_initial: self.n_initial,
            n_particles: self.n_particles,
            seed: self.seed,
            max_iters_per_restart: self.max_iters_per_restart,
        }
    }
}

// Re-export for backward compatibility
pub type EgoboxOptimizer = ArgminOptimizer;
pub type EgoboxOptimizerBuilder = ArgminOptimizerBuilder;

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

    #[derive(Debug, Clone)]
    struct TestParams {
        x: f64,
        y: f64,
    }

    impl ParameterSpace for TestParams {
        fn continuous_bounds() -> Vec<(f64, f64)> {
            vec![(-5.0, 5.0), (-5.0, 5.0)]
        }

        fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
            Ok(Self { x: x[0], y: x[1] })
        }

        fn to_continuous(&self) -> Vec<f64> {
            vec![self.x, self.y]
        }

        fn param_names() -> Vec<&'static str> {
            vec!["x", "y"]
        }
    }

    #[derive(Debug, Clone)]
    struct TestMetrics {
        loss: f64,
    }

    struct TestModel;

    impl HyperparameterOptimizable for TestModel {
        type Params = TestParams;
        type Metrics = TestMetrics;

        fn train_with_params(&mut self, params: Self::Params) -> Result<Self::Metrics, MLError> {
            // Rosenbrock function: (1-x)^2 + 100(y-x^2)^2
            // Minimum at (1, 1) with value 0
            let loss = (1.0 - params.x).powi(2) + 100.0 * (params.y - params.x.powi(2)).powi(2);
            Ok(TestMetrics { loss })
        }

        fn extract_objective(metrics: &Self::Metrics) -> f64 {
            metrics.loss
        }
    }

    #[test]
    fn test_optimizer_builder() {
        let optimizer = ArgminOptimizer::builder()
            .max_trials(20)
            .n_initial(3)
            .seed(42)
            .build();

        assert_eq!(optimizer.max_trials, 20);
        assert_eq!(optimizer.n_initial, 3);
        assert_eq!(optimizer.seed, Some(42));
    }

    #[test]
    fn test_latin_hypercube_sampling() {
        let mut rng = StdRng::seed_from_u64(42);
        let bounds = vec![(-5.0, 5.0), (-5.0, 5.0)];
        let samples = ArgminOptimizer::latin_hypercube_sampling(10, &bounds, &mut rng);

        assert_eq!(samples.nrows(), 10);
        assert_eq!(samples.ncols(), 2);

        // Check all samples are within bounds
        for i in 0..10 {
            assert!(samples[[i, 0]] >= -5.0 && samples[[i, 0]] <= 5.0);
            assert!(samples[[i, 1]] >= -5.0 && samples[[i, 1]] <= 5.0);
        }
    }

    #[test]
    #[ignore] // Expensive test - run manually
    fn test_optimizer_rosenbrock() {
        let model = TestModel;
        let optimizer = ArgminOptimizer::with_trials(15, 3);

        let result = optimizer.optimize(model).unwrap();

        // Should find minimum near (1, 1)
        assert!(
            result.best_objective < 1.0,
            "Expected to find near-optimal solution"
        );
        assert!(
            result.all_trials.len() <= 15,
            "Should have at most 15 trials"
        );
    }

    #[test]
    fn test_parameter_conversion_with_log_scale() {
        // Test that log-scale parameters are converted correctly
        #[derive(Debug, Clone, PartialEq)]
        struct LogScaleParams {
            learning_rate: f64,
            weight_decay: f64,
            batch_size: usize,
        }

        impl ParameterSpace for LogScaleParams {
            fn continuous_bounds() -> Vec<(f64, f64)> {
                vec![
                    (1e-5_f64.ln(), 1e-2_f64.ln()), // learning_rate (log scale)
                    (1e-6_f64.ln(), 1e-2_f64.ln()), // weight_decay (log scale)
                    (16.0, 256.0),                  // batch_size (linear)
                ]
            }

            fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
                Ok(Self {
                    learning_rate: x[0].exp(),
                    weight_decay: x[1].exp(),
                    batch_size: x[2].round() as usize,
                })
            }

            fn to_continuous(&self) -> Vec<f64> {
                vec![
                    self.learning_rate.ln(),
                    self.weight_decay.ln(),
                    self.batch_size as f64,
                ]
            }

            fn param_names() -> Vec<&'static str> {
                vec!["learning_rate", "weight_decay", "batch_size"]
            }
        }

        // Test continuous space values (negative for log-scale)
        let continuous = vec![-9.21, -11.51, 64.0]; // ln(1e-4), ln(1e-5), 64

        // Convert to actual parameters
        let params = LogScaleParams::from_continuous(&continuous).unwrap();

        // Verify conversions (use relative tolerance for floating point)
        // Allow 1% relative error due to log/exp floating point precision
        assert!(
            (params.learning_rate - 1e-4).abs() / 1e-4 < 0.01,
            "Learning rate should be ~1e-4 (±1%), got {}",
            params.learning_rate
        );
        assert!(
            (params.weight_decay - 1e-5).abs() / 1e-5 < 0.01,
            "Weight decay should be ~1e-5 (±1%), got {}",
            params.weight_decay
        );
        assert_eq!(params.batch_size, 64, "Batch size should be 64");

        // Verify round-trip
        let recovered = params.to_continuous();
        assert!((recovered[0] - continuous[0]).abs() < 1e-6);
        assert!((recovered[1] - continuous[1]).abs() < 1e-6);
        assert!((recovered[2] - continuous[2]).abs() < 1e-6);
    }

    #[test]
    fn test_parameter_logging_shows_converted_values() {
        // This test verifies that parameters are logged AFTER conversion,
        // so Debug output shows actual values (learning_rate ~1e-4, not -11)

        #[derive(Debug, Clone)]
        struct LogParams {
            value: f64,
        }

        impl ParameterSpace for LogParams {
            fn continuous_bounds() -> Vec<(f64, f64)> {
                vec![(1e-5_f64.ln(), 1e-2_f64.ln())]
            }

            fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
                Ok(Self { value: x[0].exp() })
            }

            fn to_continuous(&self) -> Vec<f64> {
                vec![self.value.ln()]
            }

            fn param_names() -> Vec<&'static str> {
                vec!["value"]
            }
        }

        // Raw continuous value (log scale)
        let continuous = vec![-11.51]; // ln(1e-5)

        // Convert to actual parameter
        let params = LogParams::from_continuous(&continuous).unwrap();

        // Debug output should show converted value (in scientific notation or decimal)
        let debug_str = format!("{:?}", params);

        // Should show value close to 1e-5 (not the log value -11.5)
        assert!(
            debug_str.contains("e-5") || debug_str.contains("0.00001"),
            "Debug output should show actual value ~1e-5, got: {}",
            debug_str
        );

        // Should NOT contain raw log value
        assert!(
            !debug_str.contains("-11."),
            "Debug output should NOT show raw log value -11.x"
        );
    }

    #[test]
    fn test_two_phase_optimizer_config() {
        let optimizer = ArgminOptimizer::with_trials(30, 5);
        assert_eq!(optimizer.max_trials, 30);
        // Two-phase would run 15 + 15 trials
        let phase_a = optimizer.max_trials / 2;
        let phase_b = optimizer.max_trials - phase_a;
        assert_eq!(phase_a, 15);
        assert_eq!(phase_b, 15);
    }

    #[test]
    fn test_two_phase_optimizer_odd_trials() {
        let optimizer = ArgminOptimizer::with_trials(31, 5);
        // Odd total: 15 + 16 trials
        let phase_a = optimizer.max_trials / 2;
        let phase_b = optimizer.max_trials - phase_a;
        assert_eq!(phase_a, 15);
        assert_eq!(phase_b, 16);
    }

    #[test]
    fn test_two_phase_objective_trait_object_safety() {
        // Verify TwoPhaseObjective can be used with the ObjectiveMode enum
        use super::TwoPhaseObjective;

        struct MockModel {
            mode: ObjectiveMode,
        }

        impl TwoPhaseObjective for MockModel {
            fn set_objective_mode(&mut self, mode: ObjectiveMode) {
                self.mode = mode;
            }
            fn objective_mode(&self) -> ObjectiveMode {
                self.mode
            }
        }

        let mut model = MockModel {
            mode: ObjectiveMode::Sharpe,
        };
        assert_eq!(model.objective_mode(), ObjectiveMode::Sharpe);

        model.set_objective_mode(ObjectiveMode::EpisodeReward);
        assert_eq!(model.objective_mode(), ObjectiveMode::EpisodeReward);
    }
}