//! 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, ObjectiveMode, OptimizationResult, ParameterSpace, TrialResult,
};
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 max_trials: usize,
    /// Number of initial samples via Latin Hypercube Sampling
    pub n_initial: usize,
    /// Number of particles in the swarm
    pub n_particles: usize,
    /// Random seed for reproducibility
    pub seed: Option<u64>,
    /// Maximum iterations per restart
    pub max_iters_per_restart: usize,
    /// Model name for Prometheus metrics progress reporting (e.g. "dqn", "ppo")
    pub model_name: Option<String>,
}

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

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,
            model_name: None,
        }
    }

    /// 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 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(())
    /// # }
    /// ```
    #[allow(clippy::unwrap_in_result)]
    #[allow(clippy::cognitive_complexity)]
    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!("╚═══════════════════════════════════════════════════════════╝");

        // Detect hardware budget and extract VRAM-aware parameter space bounds
        let budget = crate::traits::HardwareBudget::detect();
        info!("Hardware budget: {}MB GPU memory", budget.gpu_memory_mb);
        let bounds = M::Params::continuous_bounds_for(&budget);
        let n_params = bounds.len();

        if n_params == 0 {
            return Err(MLError::ConfigError("Parameter space has zero dimensions".to_owned())
            .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(std::sync::atomic::AtomicUsize::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,
            )?;
            if let Some(ref name) = self.model_name {
                common::metrics::training_metrics::set_hyperopt_trial(
                    name,
                    (i + 1) as f64,
                    self.max_trials as f64,
                );
            }
        }

        // Create trial budget observer BEFORE creating cost function.
        // Pass the shared trial_counter so LHS evaluations already counted above
        // are included in the budget — fixes the "0 forever" bug where PSO ran
        // unconstrained because the observer had its own separate counter.
        let observer = crate::TrialBudgetObserver::new(self.max_trials, Arc::clone(&trial_counter));

        // Create cost function wrapper
        let optimization_start = Instant::now();
        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(),
            model_name: self.model_name.clone(),
            max_trials: self.max_trials,
            optimization_start,
        };

        // Find best initial point to start optimization
        let trials = trial_results.lock().map_err(|e| {
            anyhow::anyhow!("Failed to lock trial results: {}", e)
        })?.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);

        // Algorithm selection by budget:
        // - Budget < 2× swarm → LHS (space-filling, optimal for small budgets)
        // - Budget ≥ 2× swarm → PSO (swarm optimization, needs full iterations)
        // PSO with a partial swarm is worse than random — 75% of particles
        // have no cost value, so gbest is computed from incomplete data.
        if remaining_trials < self.n_particles * 2 {
            info!(
                "Using LHS for remaining {} trials (budget < {} particles × 2 = need PSO budget ≥ {})",
                remaining_trials, self.n_particles, self.n_particles * 2
            );
            let extra_samples = Self::latin_hypercube_sampling(remaining_trials, &bounds, &mut rng);
            let mut model_guard = cost_fn.model.lock().map_err(|e| anyhow::anyhow!("lock: {e}"))?;
            for i in 0..remaining_trials {
                let continuous_vec: Vec<f64> = extra_samples.row(i).to_vec();
                Self::evaluate_point(
                    &continuous_vec,
                    &mut *model_guard,
                    &trial_results,
                    &trial_counter,
                    &param_names,
                )?;
            }
            drop(model_guard);
        }

        let max_iters_by_budget =
            ((remaining_trials as f64) / (self.n_particles as f64)).ceil() as usize;
        let max_iters = if remaining_trials < self.n_particles * 2 {
            0 // LHS already handled above
        } else {
            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 with tuned parameters.
            // Default argmin: inertia=0.72, cognitive=1.19, social=1.19
            // Problem: social >> inertia → premature convergence to global best.
            // Fix: higher inertia (0.9) for exploration, lower social (0.8) to
            // reduce collapse toward the single LHS best point.
            let solver = ParticleSwarm::new((lower_bounds, upper_bounds), self.n_particles)
                .with_inertia_factor(0.9)
                .map_err(|e| anyhow::anyhow!("PSO inertia: {e}"))?
                .with_cognitive_factor(1.5)
                .map_err(|e| anyhow::anyhow!("PSO cognitive: {e}"))?
                .with_social_factor(0.8)
                .map_err(|e| anyhow::anyhow!("PSO social: {e}"))?;

            // 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,
                    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.load(std::sync::atomic::Ordering::Relaxed));
                },
                Err(e) if e.to_string().contains("Trial budget exhausted") => {
                    info!(
                        "PSO terminated: trial budget reached ({} trials)",
                        self.max_trials
                    );
                    info!("  Evaluations: {}", trial_counter.load(std::sync::atomic::Ordering::Relaxed));
                },
                Err(e) => return Err(e),
            }
        } 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().map_err(|e| {
                anyhow::anyhow!("Failed to unwrap trial results mutex: {}", e)
            })?,
            Err(arc) => arc.lock().map_err(|e| {
                anyhow::anyhow!("Failed to lock trial results: {}", e)
            })?.clone(),
        };

        if trials.is_empty() {
            return Err(MLError::TrainingError(
                "All trials failed - no valid observations".to_owned(),
            )
            .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 in the parameter space.
    ///
    /// Converts `continuous_vec` to typed parameters, trains the model,
    /// records the trial result, and returns the objective value.
    #[allow(clippy::unwrap_in_result)]
    #[allow(clippy::cognitive_complexity)]
    pub(crate) fn evaluate_point<M>(
        continuous_vec: &[f64],
        model: &mut M,
        trial_results: &Arc<Mutex<Vec<TrialResult<M::Params>>>>,
        trial_counter: &Arc<std::sync::atomic::AtomicUsize>,
        _param_names: &[&'static str],
    ) -> Result<f64>
    where
        M: HyperparameterOptimizable,
        M::Params: ParameterSpace,
    {
        let start_time = Instant::now();

        // Increment trial counter
        let trial_num = trial_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;

        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 — NaN/divergence is a valid PSO outcome,
        // not a fatal error. Score failed trials with maximum penalty so PSO
        // learns to avoid that region of the search space.
        let (objective, trial_metrics) = match model.train_with_params(params.clone()) {
            Ok(metrics) => {
                let raw_objective = M::extract_objective(&metrics);
                let obj = if raw_objective.is_finite() {
                    raw_objective
                } else {
                    tracing::error!("Trial {} produced NON-FINITE objective ({:?}), using penalty 1e6", trial_num, raw_objective);
                    1e6
                };
                (obj, M::extract_metrics(&metrics))
            }
            Err(e) => {
                tracing::error!("Trial {} training FAILED: {:#}", trial_num, e);
                eprintln!("!!! FULL ERROR CHAIN: {:?}", e);
                (1e6, None)
            }
        };
        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().map_err(|e| {
                anyhow::anyhow!("Failed to lock trial results: {}", e)
            })?;
            results.push(TrialResult {
                trial_num,
                params,
                objective,
                duration_secs,
                metrics: trial_metrics,
            });
        }

        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`.
    #[allow(clippy::cognitive_complexity, clippy::similar_names)]
    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()
        );

        // Phase B (Sharpe-based refinement, seeded from the top Phase A
        // configs) needs `M: Clone` so the model consumed by Phase A can be
        // reconstructed for the second pass. `DQNTrainer` does not currently
        // implement `Clone`, so this routine returns the Phase A result as
        // the two-phase result. Callers that need Sharpe-mode refinement
        // should either use `optimize_parallel` (which requires `M: Clone`)
        // or run a second `optimize` pass on a freshly constructed model.

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

        Ok(phase_a_result)
    }

    /// Run optimization with parallel trial evaluation.
    ///
    /// Requires `M: Clone` so each evaluation thread gets its own model instance.
    /// Parallelism is controlled by `RAYON_NUM_THREADS` (for PSO phase) and
    /// internal thread scoping (for LHS phase, capped at `n_initial` threads).
    ///
    /// DQN/PPO trials are CPU-bound (~1 core each, 278MB VRAM). On an 8-vCPU L4
    /// with 24GB VRAM, 4-6 concurrent trials safely fit. This gives ~4-5x speedup
    /// over sequential execution.
    #[allow(clippy::unwrap_in_result)]
    #[allow(clippy::cognitive_complexity, clippy::too_many_lines)]
    pub fn optimize_parallel<M>(&self, model: M) -> Result<OptimizationResult<M::Params>>
    where
        M: HyperparameterOptimizable + Clone + Send,
        M::Params: ParameterSpace + Send,
    {
        info!("╔═══════════════════════════════════════════════════════════╗");
        info!("║  Parallel Hyperparameter Optimization (Argmin + Rayon)   ║");
        info!("╚═══════════════════════════════════════════════════════════╝");

        let rayon_threads = rayon::current_num_threads();
        info!("Parallel execution: {} rayon threads", rayon_threads);

        // Detect hardware budget and extract VRAM-aware parameter space bounds
        let budget = crate::traits::HardwareBudget::detect();
        info!("Hardware budget: {}MB GPU memory", budget.gpu_memory_mb);

        // Compute VRAM-aware concurrency strategy
        let strategy = budget.plan_hyperopt(
            M::Params::model_overhead_mb(),
            M::Params::per_sample_mb(),
            64.0,
            4096.0,
        );
        info!("HyperoptStrategy: {} concurrent trials, batch bounds [{}, {}], {:.0}MB/trial, {:.0}MB total",
            strategy.max_concurrent_trials,
            strategy.batch_size_bounds.0, strategy.batch_size_bounds.1,
            strategy.per_trial_memory_mb, strategy.total_reserved_mb);

        // Auto-scale PSO particles to match GPU concurrency for maximum
        // hardware utilization. Each PSO iteration evaluates all particles,
        // so matching particle count to VRAM slots fills every GPU slot per
        // iteration. Cap at max_trials — no benefit from more particles than
        // the total trial budget. On CPU (max_concurrent = 1), fall back to
        // configured n_particles for sequential PSO exploration.
        let effective_particles = if strategy.max_concurrent_trials > 1 {
            let scaled = strategy.max_concurrent_trials.min(self.max_trials);
            if scaled != self.n_particles {
                info!("Auto-scaled PSO particles: {} -> {} (GPU fits {} concurrent, budget {})",
                    self.n_particles, scaled, strategy.max_concurrent_trials, self.max_trials);
            }
            scaled
        } else {
            self.n_particles.min(self.max_trials)
        };

        // Trial budget is user-configured — never auto-inflate.
        let effective_max_trials = self.max_trials;

        let bounds = M::Params::continuous_bounds_for(&budget);
        let n_params = bounds.len();

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

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

        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(std::sync::atomic::AtomicUsize::new(0));

        // OOM Layer 2: Start at ceil(max/2), ramp up after first success
        let initial_concurrency = strategy.max_concurrent_trials.div_ceil(2).max(1);
        let concurrency = std::sync::atomic::AtomicUsize::new(initial_concurrency);

        info!(
            "Evaluating {} initial samples with ramp-up concurrency (start: {}, max: {})...",
            self.n_initial, initial_concurrency, strategy.max_concurrent_trials
        );

        let initial_vecs: Vec<Vec<f64>> = (0..self.n_initial)
            .map(|i| initial_samples.row(i).to_vec())
            .collect();

        for chunk in initial_vecs.chunks(initial_concurrency) {
            // OOM Layer 4: VRAM watchdog — check free memory before spawning
            if let Ok((_total_mb, free_mb, _name)) = ml_core::memory_optimization::auto_batch_size::detect_gpu_memory() {
                let total_mb = budget.gpu_memory_mb as f64;
                if total_mb > 0.0 && free_mb / total_mb < 0.10 {
                    warn!("VRAM watchdog: only {:.0}MB free ({:.0}%), reducing concurrency",
                        free_mb, (free_mb / total_mb) * 100.0);
                    let current = concurrency.load(std::sync::atomic::Ordering::Relaxed);
                    concurrency.store((current / 2).max(1), std::sync::atomic::Ordering::Relaxed);
                }
            }

            let lhs_errors: Vec<Option<anyhow::Error>> = std::thread::scope(|scope| {
                let handles: Vec<_> = chunk.iter().map(|continuous_vec| {
                    let mut model_clone = model.clone();
                    let results = Arc::clone(&trial_results);
                    let counter = Arc::clone(&trial_counter);
                    let names = param_names.clone();
                    let vec = continuous_vec.clone();
                    scope.spawn(move || {
                        Self::evaluate_point(&vec, &mut model_clone, &results, &counter, &names).err()
                    })
                }).collect();
                handles.into_iter().map(|h| h.join().unwrap_or(None)).collect()
            });

            for err in lhs_errors.into_iter().flatten() {
                tracing::error!("LHS evaluation FAILED (continuing): {:#}", err);
            }

            // OOM Layer 2: Ramp up after first chunk succeeds
            let current = concurrency.load(std::sync::atomic::Ordering::Relaxed);
            if current < strategy.max_concurrent_trials {
                info!("Ramp-up: {} -> {} concurrent trials (first chunk succeeded)",
                    current, strategy.max_concurrent_trials);
                concurrency.store(strategy.max_concurrent_trials, std::sync::atomic::Ordering::Relaxed);
            }

            if let Some(ref name) = self.model_name {
                let done = trial_counter.load(std::sync::atomic::Ordering::Relaxed);
                common::metrics::training_metrics::set_hyperopt_trial(
                    name,
                    done as f64,
                    self.max_trials as f64,
                );
            }
        }

        // Create trial budget observer sharing the same counter as LHS/evaluate_point.
        let observer = crate::TrialBudgetObserver::new(effective_max_trials, Arc::clone(&trial_counter));

        // Create parallel cost function — clones model per evaluation instead of holding mutex
        let cost_fn = ParallelObjectiveFunction {
            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(),
            model_name: self.model_name.clone(),
            max_trials: self.max_trials,
            optimization_start: Instant::now(),
        };

        // Find best initial point
        let trials = trial_results.lock().map_err(|e| {
            anyhow::anyhow!("Failed to lock trial results: {}", e)
        })?.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 Parallel Particle Swarm Optimization           ║");
        info!("╚═══════════════════════════════════════════════════════════╝");
        info!("Best initial objective: {:.6}", best_initial.objective);
        let final_concurrency = concurrency.load(std::sync::atomic::Ordering::Relaxed);
        info!("Execution mode: Parallel trials (clone-per-particle, {} VRAM-concurrent, {} rayon threads)",
            final_concurrency, rayon_threads);
        if rayon_threads > strategy.max_concurrent_trials {
            warn!("Rayon threads ({}) exceed VRAM concurrency ({}). Consider capping rayon pool at startup.",
                rayon_threads, strategy.max_concurrent_trials);
        }

        // Determine PSO iterations
        let trials_used = trials.len();
        let remaining_trials = effective_max_trials.saturating_sub(trials_used);
        let max_iters_by_budget =
            ((remaining_trials as f64) / (effective_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, effective_particles, max_iters_by_budget
        );

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

            let solver = ParticleSwarm::new((lower_bounds, upper_bounds), effective_particles)
                .with_inertia_factor(0.9)
                .map_err(|e| anyhow::anyhow!("PSO inertia: {e}"))?
                .with_cognitive_factor(1.5)
                .map_err(|e| anyhow::anyhow!("PSO cognitive: {e}"))?
                .with_social_factor(0.8)
                .map_err(|e| anyhow::anyhow!("PSO social: {e}"))?;

            let res = Executor::new(cost_fn, solver)
                .configure(|state| state.max_iters(max_iters as u64))
                .add_observer(
                    observer,
                    argmin::core::observers::ObserverMode::Always,
                )
                .run();

            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.load(std::sync::atomic::Ordering::Relaxed));
                },
                Err(e) if e.to_string().contains("Trial budget exhausted") => {
                    info!(
                        "PSO terminated: trial budget reached ({} trials, configured {})",
                        effective_max_trials, self.max_trials
                    );
                    info!("  Evaluations: {}", trial_counter.load(std::sync::atomic::Ordering::Relaxed));
                },
                Err(e) => return Err(e),
            }
        } 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().map_err(|e| {
                anyhow::anyhow!("Failed to unwrap trial results mutex: {}", e)
            })?,
            Err(arc) => arc.lock().map_err(|e| {
                anyhow::anyhow!("Failed to lock trial results: {}", e)
            })?.clone(),
        };

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

        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)
    }
}

/// Run optimization using Tree-Parzen Estimator (TPE) instead of PSO.
///
/// TPE builds separate density models for "good" and "bad" trials,
/// then suggests new points by maximizing Expected Improvement.
/// Better sample efficiency than PSO in 20-50D spaces.
///
/// # Arguments
///
/// * `model` - Model implementing `HyperparameterOptimizable` (consumed)
/// * `max_trials` - Total number of trials to evaluate
/// * `n_initial` - Number of initial LHS samples before TPE kicks in
/// * `seed` - Optional random seed for reproducibility
///
/// # Returns
///
/// `OptimizationResult<M::Params>` containing best parameters and full trial history
#[allow(clippy::cognitive_complexity)]
pub fn optimize_with_tpe<M>(
    mut model: M,
    max_trials: usize,
    n_initial: usize,
    seed: Option<u64>,
    model_name: Option<&str>,
) -> Result<OptimizationResult<M::Params>>
where
    M: HyperparameterOptimizable + Send,
    M::Params: ParameterSpace + Send,
{
    use crate::tpe::{TpeConfig, TpeOptimizer};

    info!("╔═══════════════════════════════════════════════════════════╗");
    info!("║  Bayesian Hyperparameter Optimization (TPE)              ║");
    info!("╚═══════════════════════════════════════════════════════════╝");

    let budget = crate::traits::HardwareBudget::detect();
    let bounds = M::Params::continuous_bounds_for(&budget);
    let n_params = bounds.len();

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

    // Clamp n_initial to valid range: at least 1, at most max_trials - 1
    let n_initial = n_initial.max(1).min(max_trials.saturating_sub(1).max(1));

    info!("Configuration:");
    info!("  Optimizer: TPE (Tree-Parzen Estimator)");
    info!("  Max Trials: {}", max_trials);
    info!("  Initial LHS Samples: {}", n_initial);
    info!("  Parameters: {}", n_params);
    let n_candidates = 256_usize.max(8 * n_params);
    let gamma = if max_trials < 50 { 0.15 } else { 0.25 };
    info!("  Gamma (good quantile): {:.2}", gamma);
    info!("  EI Candidates: {} (scaled for {}D)", n_candidates, n_params);

    let param_names = M::Params::param_names();
    for (i, name) in param_names.iter().enumerate() {
        if let Some(&(lo, hi)) = bounds.get(i) {
            info!("    {} - [{:.6}, {:.6}]", name, lo, hi);
        }
    }

    let tpe_config = TpeConfig {
        n_dims: n_params,
        max_trials,
        n_initial,
        gamma,
        n_candidates,
        seed,
    };
    let mut tpe = TpeOptimizer::new(tpe_config);

    let trial_results = Arc::new(Mutex::new(Vec::new()));
    let trial_counter = Arc::new(std::sync::atomic::AtomicUsize::new(0));
    let tpe_start = Instant::now();

    for trial_idx in 0..max_trials {
        let continuous_vec = tpe.suggest(&bounds);

        info!(
            "TPE Trial {}/{}: evaluating suggested point",
            trial_idx + 1,
            max_trials
        );

        ArgminOptimizer::evaluate_point(
            &continuous_vec,
            &mut model,
            &trial_results,
            &trial_counter,
            &param_names,
        )?;

        // Feed the objective back to TPE so it can update its density models
        let results = trial_results.lock().map_err(|e| {
            MLError::ConcurrencyError {
                operation: format!("lock trial results: {}", e),
            }
        })?;
        if let Some(last) = results.last() {
            tpe.add_trial(continuous_vec, last.objective);
            info!(
                "TPE Trial {}/{}: objective = {:.6}",
                trial_idx + 1,
                max_trials,
                last.objective
            );
        }

        if let Some(name) = model_name {
            common::metrics::training_metrics::set_hyperopt_trial(
                name,
                (trial_idx + 1) as f64,
                max_trials as f64,
            );
            common::metrics::training_metrics::set_hyperopt_elapsed(
                name,
                tpe_start.elapsed().as_secs_f64(),
            );
            if let Some(last) = results.last() {
                common::metrics::training_metrics::record_hyperopt_trial_duration(
                    name,
                    last.duration_secs,
                );
            }
            if let Some(best) = results.iter().map(|r| r.objective).reduce(f64::min) {
                common::metrics::training_metrics::set_hyperopt_best_objective(name, best);
                common::metrics::training_metrics::set_hyperopt_trial_best_loss(name, best);
            }
        }
    }

    // Extract results
    let trials = match Arc::try_unwrap(trial_results) {
        Ok(mutex) => mutex.into_inner().map_err(|e| {
            anyhow::anyhow!("Failed to unwrap trial results mutex: {}", e)
        })?,
        Err(arc) => arc.lock().map_err(|e| {
            anyhow::anyhow!("Failed to lock trial results: {}", e)
        })?.clone(),
    };

    if trials.is_empty() {
        return Err(
            MLError::ModelError("No valid trials completed".to_owned()).into(),
        );
    }

    let result = OptimizationResult::from_trials(trials);

    info!("═══ TPE Optimization Complete ═══");
    info!(
        "  Best trial: objective {:.6}",
        result.best_objective
    );
    info!("  Total trials: {}", result.all_trials.len());

    Ok(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<std::sync::atomic::AtomicUsize>,
    param_names: Vec<&'static str>,
    bounds: Vec<(f64, f64)>,
    observer: crate::TrialBudgetObserver,
    model_name: Option<String>,
    max_trials: usize,
    optimization_start: Instant,
}

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

    #[allow(clippy::cognitive_complexity)]
    fn cost(&self, param: &Self::Param) -> Result<Self::Output, argmin::core::Error> {
        // Increment the shared trial counter FIRST — this is the canonical counter
        // used by both the observer (for budget enforcement) and the results log.
        let trial_num = self.trial_counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1;

        // Check budget immediately after incrementing — abort if we've exceeded the limit.
        if trial_num > self.max_trials {
            warn!(
                "Trial budget exhausted: {}/{} trials. Stopping PSO.",
                trial_num, self.max_trials
            );
            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();

        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().map_err(|e| {
            argmin::core::Error::msg(format!("Failed to lock model: {}", e))
        })?;
        let metrics = match model.train_with_params(params.clone()) {
            Ok(m) => m,
            Err(e) => {
                tracing::error!("Trial {} training FAILED: {:#}", 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 trial_metrics = M::extract_metrics(&metrics);
        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().map_err(|e| {
                argmin::core::Error::msg(format!("Failed to lock trial results: {}", e))
            })?;
            results.push(TrialResult {
                trial_num,
                params,
                objective,
                duration_secs,
                metrics: trial_metrics,
            });
        }

        if let Some(ref name) = self.model_name {
            common::metrics::training_metrics::set_hyperopt_trial(
                name,
                trial_num as f64,
                self.max_trials as f64,
            );
            common::metrics::training_metrics::record_hyperopt_trial_duration(name, duration_secs);
            common::metrics::training_metrics::set_hyperopt_elapsed(
                name,
                self.optimization_start.elapsed().as_secs_f64(),
            );
            // Update running best objective
            if let Ok(results) = self.trial_results.lock() {
                if let Some(best) = results.iter().map(|r| r.objective).reduce(f64::min) {
                    common::metrics::training_metrics::set_hyperopt_best_objective(name, best);
                    common::metrics::training_metrics::set_hyperopt_trial_best_loss(name, best);
                }
            }
        }

        Ok(objective)
    }
}

/// Parallel cost function for argmin — clones model per evaluation.
///
/// Unlike [`ObjectiveFunction`] which holds an `Arc<Mutex<M>>` lock for the
/// entire training duration, this variant clones the model and releases the
/// lock immediately. This allows argmin's rayon-based PSO to evaluate multiple
/// particles truly in parallel.
///
/// Each clone gets its own `InternalDQNTrainer` (created inside `train_with_params`),
/// so GPU/CPU resources are the only shared concern. DQN uses ~1 core + 278MB VRAM
/// per trial, allowing 4-6 concurrent trials on an 8-vCPU L4 with 24GB VRAM.
struct ParallelObjectiveFunction<M>
where
    M: HyperparameterOptimizable + Clone,
    M::Params: ParameterSpace,
{
    model: Arc<Mutex<M>>,
    trial_results: Arc<Mutex<Vec<TrialResult<M::Params>>>>,
    trial_counter: Arc<std::sync::atomic::AtomicUsize>,
    param_names: Vec<&'static str>,
    bounds: Vec<(f64, f64)>,
    observer: crate::TrialBudgetObserver,
    model_name: Option<String>,
    max_trials: usize,
    optimization_start: Instant,
}

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

    #[allow(clippy::cognitive_complexity)]
    fn cost(&self, param: &Self::Param) -> Result<Self::Output, argmin::core::Error> {
        // Increment the shared trial counter FIRST — shared with LHS phase and observer.
        let trial_num = self.trial_counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1;

        // Check budget immediately after incrementing.
        if trial_num > self.max_trials {
            warn!(
                "Trial budget exhausted: {}/{} trials. Stopping PSO.",
                trial_num, self.max_trials
            );
            return Err(argmin::core::Error::msg("Trial budget exhausted"));
        }

        // Clone model and release mutex immediately — enables true parallel evaluation
        let mut model = {
            let guard = self.model.lock().map_err(|e| {
                argmin::core::Error::msg(format!("Failed to lock model: {}", e))
            })?;
            guard.clone()
        }; // Mutex released here — other particles can proceed concurrently

        // 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();

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

        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);
            },
        };

        info!("  Parameters (converted): {:?}", params);

        let metrics = match model.train_with_params(params.clone()) {
            Ok(m) => m,
            Err(e) => {
                tracing::error!("Trial {} training FAILED: {:#}", trial_num, e);
                return Ok(1e6);
            },
        };

        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 trial_metrics = M::extract_metrics(&metrics);
        let duration_secs = start_time.elapsed().as_secs_f64();

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

        {
            let mut results = self.trial_results.lock().map_err(|e| {
                argmin::core::Error::msg(format!("Failed to lock trial results: {}", e))
            })?;
            results.push(TrialResult {
                trial_num,
                params,
                objective,
                duration_secs,
                metrics: trial_metrics,
            });
        }

        if let Some(ref name) = self.model_name {
            common::metrics::training_metrics::set_hyperopt_trial(
                name,
                trial_num as f64,
                self.max_trials as f64,
            );
            common::metrics::training_metrics::record_hyperopt_trial_duration(name, duration_secs);
            common::metrics::training_metrics::set_hyperopt_elapsed(
                name,
                self.optimization_start.elapsed().as_secs_f64(),
            );
            if let Ok(results) = self.trial_results.lock() {
                if let Some(best) = results.iter().map(|r| r.objective).reduce(f64::min) {
                    common::metrics::training_metrics::set_hyperopt_best_objective(name, best);
                    common::metrics::training_metrics::set_hyperopt_trial_best_loss(name, best);
                }
            }
        }

        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,
    model_name: Option<String>,
}

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,
            model_name: None,
        }
    }

    /// 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
    }

    /// Set model name for per-trial Prometheus progress metrics
    pub fn model_name<S: Into<String>>(mut self, name: S) -> Self {
        self.model_name = Some(name.into());
        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,
            model_name: self.model_name,
        }
    }
}

#[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);
    }
}