sklears_model_selection/

multi_fidelity_optimization.rs

//! Multi-Fidelity Bayesian Optimization
//!
//! This module provides multi-fidelity Bayesian optimization for efficient hyperparameter tuning
//! by leveraging multiple approximation levels (fidelities) of the objective function.
//! Lower fidelity evaluations are cheaper but less accurate, while higher fidelity evaluations
//! are more expensive but more accurate.

use scirs2_core::ndarray::Array1;
use scirs2_core::random::rngs::StdRng;
use scirs2_core::random::Rng;
use scirs2_core::random::SeedableRng;
use sklears_core::types::Float;
use std::collections::HashMap;

/// Fidelity levels for multi-fidelity optimization
#[derive(Debug, Clone)]
pub enum FidelityLevel {
    /// Low fidelity (fast, less accurate)
    Low {
        sample_fraction: Float,

        epochs_fraction: Float,

        cv_folds: usize,
    },
    /// Medium fidelity (moderate speed and accuracy)
    Medium {
        sample_fraction: Float,

        epochs_fraction: Float,
        cv_folds: usize,
    },
    /// High fidelity (slow, most accurate)
    High {
        sample_fraction: Float,
        epochs_fraction: Float,
        cv_folds: usize,
    },
    /// Custom fidelity with user-defined parameters
    Custom {
        parameters: HashMap<String, Float>,
        relative_cost: Float,
        accuracy_estimate: Float,
    },
}

/// Multi-fidelity optimization strategies
#[derive(Debug, Clone)]
pub enum MultiFidelityStrategy {
    /// Successive Halving with multiple fidelities
    SuccessiveHalving {
        eta: Float,

        min_fidelity: FidelityLevel,

        max_fidelity: FidelityLevel,
    },
    /// Multi-fidelity Bayesian Optimization (MFBO)
    BayesianOptimization {
        acquisition_function: AcquisitionFunction,

        fidelity_selection: FidelitySelectionMethod,
        correlation_model: CorrelationModel,
    },
    /// Hyperband with multi-fidelity
    Hyperband {
        max_budget: Float,
        eta: Float,
        fidelities: Vec<FidelityLevel>,
    },
    /// BOHB (Bayesian Optimization and Hyperband)
    BOHB {
        min_budget: Float,
        max_budget: Float,
        eta: Float,
        bandwidth_factor: Float,
    },
    /// Fabolas (Fast Bayesian Optimization on Large Datasets)
    Fabolas {
        min_dataset_fraction: Float,
        max_dataset_fraction: Float,
        cost_model: CostModel,
    },
    /// Multi-Task Gaussian Process
    MultiTaskGP {
        task_similarity: Float,
        shared_hyperparameters: Vec<String>,
    },
}

/// Acquisition functions for multi-fidelity optimization
#[derive(Debug, Clone)]
pub enum AcquisitionFunction {
    /// Expected Improvement with fidelity consideration
    ExpectedImprovement,
    /// Upper Confidence Bound with fidelity adjustment
    UpperConfidenceBound { beta: Float },
    /// Probability of Improvement
    ProbabilityOfImprovement,
    /// Knowledge Gradient
    KnowledgeGradient,
    /// Entropy Search
    EntropySearch,
    /// Multi-fidelity Expected Improvement
    MultiFidelityEI { fidelity_weight: Float },
}

/// Methods for selecting fidelity levels
#[derive(Debug, Clone)]
pub enum FidelitySelectionMethod {
    /// Always start with lowest fidelity
    LowestFirst,
    /// Dynamic selection based on uncertainty
    UncertaintyBased { threshold: Float },
    /// Cost-aware selection
    CostAware { budget_fraction: Float },
    /// Performance-based selection
    PerformanceBased { improvement_threshold: Float },
    /// Information-theoretic selection
    InformationTheoretic,
}

/// Models for correlation between fidelities
#[derive(Debug, Clone)]
pub enum CorrelationModel {
    /// Linear correlation between fidelities
    Linear { correlation_strength: Float },
    /// Exponential correlation
    Exponential { decay_rate: Float },
    /// Learned correlation using Gaussian Process
    GaussianProcess { kernel_type: String },
    /// Rank correlation
    RankCorrelation,
}

/// Cost models for different fidelity levels
#[derive(Debug, Clone)]
pub enum CostModel {
    /// Polynomial cost model
    Polynomial {
        degree: usize,

        coefficients: Vec<Float>,
    },
    /// Exponential cost model
    Exponential { base: Float, scale: Float },
    /// Linear cost model
    Linear { slope: Float, intercept: Float },
    /// Custom cost function
    Custom { cost_function: String },
}

/// Multi-fidelity optimization configuration
#[derive(Debug, Clone)]
pub struct MultiFidelityConfig {
    pub strategy: MultiFidelityStrategy,
    pub max_evaluations: usize,
    pub max_budget: Float,
    pub early_stopping_patience: usize,
    pub fidelity_progression: FidelityProgression,
    pub random_state: Option<u64>,
    pub parallel_evaluations: usize,
}

/// Fidelity progression strategies
#[derive(Debug, Clone)]
pub enum FidelityProgression {
    /// Linear progression from low to high fidelity
    Linear,
    /// Exponential progression
    Exponential { growth_rate: Float },
    /// Adaptive progression based on performance
    Adaptive { adaptation_rate: Float },
    /// Conservative progression (slow increase)
    Conservative,
    /// Aggressive progression (fast increase)
    Aggressive,
}

/// Evaluation result at a specific fidelity
#[derive(Debug, Clone)]
pub struct FidelityEvaluation {
    pub hyperparameters: HashMap<String, Float>,
    pub fidelity: FidelityLevel,
    pub score: Float,
    pub cost: Float,
    pub evaluation_time: Float,
    pub uncertainty: Option<Float>,
    pub additional_metrics: HashMap<String, Float>,
}

/// Multi-fidelity optimization result
#[derive(Debug, Clone)]
pub struct MultiFidelityResult {
    pub best_hyperparameters: HashMap<String, Float>,
    pub best_score: Float,
    pub best_fidelity: FidelityLevel,
    pub optimization_history: Vec<FidelityEvaluation>,
    pub total_cost: Float,
    pub total_time: Float,
    pub convergence_curve: Vec<Float>,
    pub fidelity_usage: HashMap<String, usize>,
    pub cost_efficiency: Float,
}

/// Multi-fidelity Bayesian optimizer
#[derive(Debug)]
pub struct MultiFidelityOptimizer {
    config: MultiFidelityConfig,
    gaussian_process: MultiFidelityGP,
    evaluation_history: Vec<FidelityEvaluation>,
    current_best: Option<FidelityEvaluation>,
    rng: StdRng,
}

/// Multi-fidelity Gaussian Process
#[derive(Debug, Clone)]
pub struct MultiFidelityGP {
    observations: Vec<(Array1<Float>, Float, Float)>, // (hyperparams, fidelity, score)
    hyperparameters: GPHyperparameters,
    trained: bool,
}

/// Gaussian Process hyperparameters
#[derive(Debug, Clone)]
pub struct GPHyperparameters {
    pub length_scales: Array1<Float>,
    pub signal_variance: Float,
    pub noise_variance: Float,
    pub fidelity_correlation: Float,
}

impl Default for MultiFidelityConfig {
    fn default() -> Self {
        Self {
            strategy: MultiFidelityStrategy::BayesianOptimization {
                acquisition_function: AcquisitionFunction::ExpectedImprovement,
                fidelity_selection: FidelitySelectionMethod::UncertaintyBased { threshold: 0.1 },
                correlation_model: CorrelationModel::Linear {
                    correlation_strength: 0.8,
                },
            },
            max_evaluations: 100,
            max_budget: 1000.0,
            early_stopping_patience: 10,
            fidelity_progression: FidelityProgression::Adaptive {
                adaptation_rate: 0.1,
            },
            random_state: None,
            parallel_evaluations: 1,
        }
    }
}

impl MultiFidelityOptimizer {
    /// Create a new multi-fidelity optimizer
    pub fn new(config: MultiFidelityConfig) -> Self {
        let rng = match config.random_state {
            Some(seed) => StdRng::seed_from_u64(seed),
            None => {
                use scirs2_core::random::thread_rng;
                StdRng::from_rng(&mut thread_rng())
            }
        };

        let gaussian_process = MultiFidelityGP::new();

        Self {
            config,
            gaussian_process,
            evaluation_history: Vec::new(),
            current_best: None,
            rng,
        }
    }

    /// Optimize hyperparameters using multi-fidelity approach
    pub fn optimize<F>(
        &mut self,
        evaluation_fn: F,
        parameter_bounds: &[(Float, Float)],
    ) -> Result<MultiFidelityResult, Box<dyn std::error::Error>>
    where
        F: Fn(
            &HashMap<String, Float>,
            &FidelityLevel,
        ) -> Result<FidelityEvaluation, Box<dyn std::error::Error>>,
    {
        let start_time = std::time::Instant::now();
        let mut total_cost = 0.0;
        let mut convergence_curve = Vec::new();
        let mut fidelity_usage = HashMap::new();

        match &self.config.strategy {
            MultiFidelityStrategy::SuccessiveHalving { .. } => {
                self.successive_halving_optimize(
                    &evaluation_fn,
                    parameter_bounds,
                    &mut total_cost,
                    &mut convergence_curve,
                    &mut fidelity_usage,
                )?;
            }
            MultiFidelityStrategy::BayesianOptimization { .. } => {
                self.bayesian_optimize(
                    &evaluation_fn,
                    parameter_bounds,
                    &mut total_cost,
                    &mut convergence_curve,
                    &mut fidelity_usage,
                )?;
            }
            MultiFidelityStrategy::Hyperband { .. } => {
                self.hyperband_optimize(
                    &evaluation_fn,
                    parameter_bounds,
                    &mut total_cost,
                    &mut convergence_curve,
                    &mut fidelity_usage,
                )?;
            }
            MultiFidelityStrategy::BOHB { .. } => {
                self.bohb_optimize(
                    &evaluation_fn,
                    parameter_bounds,
                    &mut total_cost,
                    &mut convergence_curve,
                    &mut fidelity_usage,
                )?;
            }
            MultiFidelityStrategy::Fabolas { .. } => {
                self.fabolas_optimize(
                    &evaluation_fn,
                    parameter_bounds,
                    &mut total_cost,
                    &mut convergence_curve,
                    &mut fidelity_usage,
                )?;
            }
            MultiFidelityStrategy::MultiTaskGP { .. } => {
                self.multi_task_gp_optimize(
                    &evaluation_fn,
                    parameter_bounds,
                    &mut total_cost,
                    &mut convergence_curve,
                    &mut fidelity_usage,
                )?;
            }
        }

        let total_time = start_time.elapsed().as_secs_f64() as Float;
        let cost_efficiency = if total_cost > 0.0 {
            self.current_best.as_ref().map_or(0.0, |best| best.score) / total_cost
        } else {
            0.0
        };

        Ok(MultiFidelityResult {
            best_hyperparameters: self
                .current_best
                .as_ref()
                .map(|best| best.hyperparameters.clone())
                .unwrap_or_default(),
            best_score: self.current_best.as_ref().map_or(0.0, |best| best.score),
            best_fidelity: self
                .current_best
                .as_ref()
                .map(|best| best.fidelity.clone())
                .unwrap_or(self.get_default_fidelity()),
            optimization_history: self.evaluation_history.clone(),
            total_cost,
            total_time,
            convergence_curve,
            fidelity_usage,
            cost_efficiency,
        })
    }

    /// Successive halving with multi-fidelity
    fn successive_halving_optimize<F>(
        &mut self,
        evaluation_fn: &F,
        parameter_bounds: &[(Float, Float)],
        total_cost: &mut Float,
        convergence_curve: &mut Vec<Float>,
        fidelity_usage: &mut HashMap<String, usize>,
    ) -> Result<(), Box<dyn std::error::Error>>
    where
        F: Fn(
            &HashMap<String, Float>,
            &FidelityLevel,
        ) -> Result<FidelityEvaluation, Box<dyn std::error::Error>>,
    {
        let (eta, min_fidelity, max_fidelity) = match &self.config.strategy {
            MultiFidelityStrategy::SuccessiveHalving {
                eta,
                min_fidelity,
                max_fidelity,
            } => (*eta, min_fidelity.clone(), max_fidelity.clone()),
            _ => unreachable!(),
        };

        let mut configurations = self.generate_initial_configurations(parameter_bounds, 50)?;
        let mut current_fidelity = min_fidelity;

        while configurations.len() > 1 && !self.should_stop() {
            let mut evaluations = Vec::new();

            // Evaluate all configurations at current fidelity
            for config in &configurations {
                let evaluation = evaluation_fn(config, &current_fidelity)?;
                *total_cost += evaluation.cost;
                *fidelity_usage
                    .entry(self.fidelity_to_string(&current_fidelity))
                    .or_insert(0) += 1;

                self.evaluation_history.push(evaluation.clone());
                evaluations.push(evaluation.clone());

                if self.update_best(&evaluation) {
                    convergence_curve.push(self.current_best.as_ref().unwrap().score);
                } else if let Some(best) = &self.current_best {
                    convergence_curve.push(best.score);
                }
            }

            // Keep top 1/eta configurations
            evaluations.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap());
            let keep_count = (configurations.len() as Float / eta).max(1.0) as usize;

            configurations = evaluations
                .iter()
                .take(keep_count)
                .map(|eval| eval.hyperparameters.clone())
                .collect();

            // Increase fidelity
            current_fidelity = self.increase_fidelity(&current_fidelity, &max_fidelity);
        }

        Ok(())
    }

    /// Bayesian optimization with multi-fidelity
    fn bayesian_optimize<F>(
        &mut self,
        evaluation_fn: &F,
        parameter_bounds: &[(Float, Float)],
        total_cost: &mut Float,
        convergence_curve: &mut Vec<Float>,
        fidelity_usage: &mut HashMap<String, usize>,
    ) -> Result<(), Box<dyn std::error::Error>>
    where
        F: Fn(
            &HashMap<String, Float>,
            &FidelityLevel,
        ) -> Result<FidelityEvaluation, Box<dyn std::error::Error>>,
    {
        let (acquisition_function, fidelity_selection, _correlation_model) =
            match &self.config.strategy {
                MultiFidelityStrategy::BayesianOptimization {
                    acquisition_function,
                    fidelity_selection,
                    correlation_model,
                } => (
                    acquisition_function.clone(),
                    fidelity_selection.clone(),
                    correlation_model.clone(),
                ),
                _ => unreachable!(),
            };

        // Initialize with random evaluations
        let init_evaluations = 5;
        for _ in 0..init_evaluations {
            let config = self.sample_random_configuration(parameter_bounds)?;
            let fidelity = self.select_fidelity(&fidelity_selection, None)?;

            let evaluation = evaluation_fn(&config, &fidelity)?;
            *total_cost += evaluation.cost;
            *fidelity_usage
                .entry(self.fidelity_to_string(&fidelity))
                .or_insert(0) += 1;

            self.evaluation_history.push(evaluation.clone());
            if self.update_best(&evaluation) {
                convergence_curve.push(self.current_best.as_ref().unwrap().score);
            } else if let Some(best) = &self.current_best {
                convergence_curve.push(best.score);
            }
        }

        // Update Gaussian Process
        self.gaussian_process.update(&self.evaluation_history)?;

        // Bayesian optimization loop
        while self.evaluation_history.len() < self.config.max_evaluations && !self.should_stop() {
            // Select next configuration and fidelity
            let next_config = self.optimize_acquisition(&acquisition_function, parameter_bounds)?;
            let next_fidelity = self.select_fidelity(&fidelity_selection, Some(&next_config))?;

            let evaluation = evaluation_fn(&next_config, &next_fidelity)?;
            *total_cost += evaluation.cost;
            *fidelity_usage
                .entry(self.fidelity_to_string(&next_fidelity))
                .or_insert(0) += 1;

            self.evaluation_history.push(evaluation.clone());
            if self.update_best(&evaluation) {
                convergence_curve.push(self.current_best.as_ref().unwrap().score);
            } else if let Some(best) = &self.current_best {
                convergence_curve.push(best.score);
            }

            // Update Gaussian Process periodically
            if self.evaluation_history.len() % 5 == 0 {
                self.gaussian_process.update(&self.evaluation_history)?;
            }
        }

        Ok(())
    }

    /// Hyperband optimization
    fn hyperband_optimize<F>(
        &mut self,
        evaluation_fn: &F,
        parameter_bounds: &[(Float, Float)],
        total_cost: &mut Float,
        convergence_curve: &mut Vec<Float>,
        fidelity_usage: &mut HashMap<String, usize>,
    ) -> Result<(), Box<dyn std::error::Error>>
    where
        F: Fn(
            &HashMap<String, Float>,
            &FidelityLevel,
        ) -> Result<FidelityEvaluation, Box<dyn std::error::Error>>,
    {
        let (max_budget, eta, fidelities) = match &self.config.strategy {
            MultiFidelityStrategy::Hyperband {
                max_budget,
                eta,
                fidelities,
            } => (*max_budget, *eta, fidelities.clone()),
            _ => unreachable!(),
        };

        let log_eta = eta.ln();
        let s_max = (max_budget.ln() / log_eta).floor() as usize;

        for s in 0..=s_max {
            let n = ((s_max + 1) as Float * eta.powi(s as i32) / (s + 1) as Float).ceil() as usize;
            let r = max_budget * eta.powi(-(s as i32));

            let mut configurations = self.generate_initial_configurations(parameter_bounds, n)?;
            let current_budget = r;

            for i in 0..=s {
                let n_i = (n as Float * eta.powi(-(i as i32))).floor() as usize;
                let r_i = current_budget * eta.powi(i as i32);

                if configurations.len() > n_i {
                    configurations.truncate(n_i);
                }

                let fidelity = self.budget_to_fidelity(r_i, &fidelities);
                let mut evaluations = Vec::new();

                for config in &configurations {
                    let evaluation = evaluation_fn(config, &fidelity)?;
                    *total_cost += evaluation.cost;
                    *fidelity_usage
                        .entry(self.fidelity_to_string(&fidelity))
                        .or_insert(0) += 1;

                    self.evaluation_history.push(evaluation.clone());
                    evaluations.push(evaluation.clone());

                    if self.update_best(&evaluation) {
                        convergence_curve.push(self.current_best.as_ref().unwrap().score);
                    } else if let Some(best) = &self.current_best {
                        convergence_curve.push(best.score);
                    }
                }

                // Keep top configurations
                evaluations.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap());
                configurations = evaluations
                    .iter()
                    .take(n_i)
                    .map(|eval| eval.hyperparameters.clone())
                    .collect();
            }
        }

        Ok(())
    }

    /// BOHB optimization (Bayesian Optimization and Hyperband)
    fn bohb_optimize<F>(
        &mut self,
        evaluation_fn: &F,
        parameter_bounds: &[(Float, Float)],
        total_cost: &mut Float,
        convergence_curve: &mut Vec<Float>,
        fidelity_usage: &mut HashMap<String, usize>,
    ) -> Result<(), Box<dyn std::error::Error>>
    where
        F: Fn(
            &HashMap<String, Float>,
            &FidelityLevel,
        ) -> Result<FidelityEvaluation, Box<dyn std::error::Error>>,
    {
        // Simplified BOHB implementation combining Hyperband with Bayesian optimization
        // Start with Hyperband for exploration
        self.hyperband_optimize(
            evaluation_fn,
            parameter_bounds,
            total_cost,
            convergence_curve,
            fidelity_usage,
        )?;

        // Continue with Bayesian optimization for exploitation
        let remaining_budget = self.config.max_budget - *total_cost;
        if remaining_budget > 0.0 {
            self.bayesian_optimize(
                evaluation_fn,
                parameter_bounds,
                total_cost,
                convergence_curve,
                fidelity_usage,
            )?;
        }

        Ok(())
    }

    /// Fabolas optimization
    fn fabolas_optimize<F>(
        &mut self,
        evaluation_fn: &F,
        parameter_bounds: &[(Float, Float)],
        total_cost: &mut Float,
        convergence_curve: &mut Vec<Float>,
        fidelity_usage: &mut HashMap<String, usize>,
    ) -> Result<(), Box<dyn std::error::Error>>
    where
        F: Fn(
            &HashMap<String, Float>,
            &FidelityLevel,
        ) -> Result<FidelityEvaluation, Box<dyn std::error::Error>>,
    {
        // Simplified Fabolas implementation focusing on dataset size as fidelity
        let (min_fraction, max_fraction, _cost_model) = match &self.config.strategy {
            MultiFidelityStrategy::Fabolas {
                min_dataset_fraction,
                max_dataset_fraction,
                cost_model,
            } => (*min_dataset_fraction, *max_dataset_fraction, cost_model),
            _ => unreachable!(),
        };

        let mut current_fraction = min_fraction;
        let fraction_step = (max_fraction - min_fraction) / 10.0;

        while current_fraction <= max_fraction && !self.should_stop() {
            let fidelity = FidelityLevel::Custom {
                parameters: {
                    let mut params = HashMap::new();
                    params.insert("dataset_fraction".to_string(), current_fraction);
                    params
                },
                relative_cost: current_fraction,
                accuracy_estimate: current_fraction.sqrt(),
            };

            let config = self.sample_random_configuration(parameter_bounds)?;
            let evaluation = evaluation_fn(&config, &fidelity)?;

            *total_cost += evaluation.cost;
            *fidelity_usage
                .entry(self.fidelity_to_string(&fidelity))
                .or_insert(0) += 1;

            self.evaluation_history.push(evaluation.clone());
            if self.update_best(&evaluation) {
                convergence_curve.push(self.current_best.as_ref().unwrap().score);
            } else if let Some(best) = &self.current_best {
                convergence_curve.push(best.score);
            }

            current_fraction += fraction_step;
        }

        Ok(())
    }

    /// Multi-task Gaussian Process optimization
    fn multi_task_gp_optimize<F>(
        &mut self,
        evaluation_fn: &F,
        parameter_bounds: &[(Float, Float)],
        total_cost: &mut Float,
        convergence_curve: &mut Vec<Float>,
        fidelity_usage: &mut HashMap<String, usize>,
    ) -> Result<(), Box<dyn std::error::Error>>
    where
        F: Fn(
            &HashMap<String, Float>,
            &FidelityLevel,
        ) -> Result<FidelityEvaluation, Box<dyn std::error::Error>>,
    {
        // Simplified multi-task GP implementation
        // Treat each fidelity as a separate task
        let fidelities = vec![
            FidelityLevel::Low {
                sample_fraction: 0.1,
                epochs_fraction: 0.1,
                cv_folds: 3,
            },
            FidelityLevel::Medium {
                sample_fraction: 0.5,
                epochs_fraction: 0.5,
                cv_folds: 5,
            },
            FidelityLevel::High {
                sample_fraction: 1.0,
                epochs_fraction: 1.0,
                cv_folds: 10,
            },
        ];

        while self.evaluation_history.len() < self.config.max_evaluations && !self.should_stop() {
            for fidelity in &fidelities {
                let config = self.sample_random_configuration(parameter_bounds)?;
                let evaluation = evaluation_fn(&config, fidelity)?;

                *total_cost += evaluation.cost;
                *fidelity_usage
                    .entry(self.fidelity_to_string(fidelity))
                    .or_insert(0) += 1;

                self.evaluation_history.push(evaluation.clone());
                if self.update_best(&evaluation) {
                    convergence_curve.push(self.current_best.as_ref().unwrap().score);
                } else if let Some(best) = &self.current_best {
                    convergence_curve.push(best.score);
                }

                if self.evaluation_history.len() >= self.config.max_evaluations {
                    break;
                }
            }
        }

        Ok(())
    }

    /// Generate initial random configurations
    fn generate_initial_configurations(
        &mut self,
        parameter_bounds: &[(Float, Float)],
        n: usize,
    ) -> Result<Vec<HashMap<String, Float>>, Box<dyn std::error::Error>> {
        let mut configurations = Vec::new();

        for _ in 0..n {
            configurations.push(self.sample_random_configuration(parameter_bounds)?);
        }

        Ok(configurations)
    }

    /// Sample a random configuration
    fn sample_random_configuration(
        &mut self,
        parameter_bounds: &[(Float, Float)],
    ) -> Result<HashMap<String, Float>, Box<dyn std::error::Error>> {
        let mut config = HashMap::new();

        for (i, &(low, high)) in parameter_bounds.iter().enumerate() {
            let value = self.rng.gen_range(low..high + 1.0);
            config.insert(format!("param_{}", i), value);
        }

        Ok(config)
    }

    /// Select fidelity level based on strategy
    fn select_fidelity(
        &mut self,
        method: &FidelitySelectionMethod,
        _config: Option<&HashMap<String, Float>>,
    ) -> Result<FidelityLevel, Box<dyn std::error::Error>> {
        match method {
            FidelitySelectionMethod::LowestFirst => Ok(FidelityLevel::Low {
                sample_fraction: 0.1,
                epochs_fraction: 0.1,
                cv_folds: 3,
            }),
            FidelitySelectionMethod::UncertaintyBased { threshold } => {
                // Use uncertainty to determine fidelity
                if self.evaluation_history.len() < 5 {
                    Ok(FidelityLevel::Low {
                        sample_fraction: 0.1,
                        epochs_fraction: 0.1,
                        cv_folds: 3,
                    })
                } else {
                    let avg_uncertainty = self
                        .evaluation_history
                        .iter()
                        .filter_map(|eval| eval.uncertainty)
                        .sum::<Float>()
                        / self.evaluation_history.len() as Float;

                    if avg_uncertainty > *threshold {
                        Ok(FidelityLevel::High {
                            sample_fraction: 1.0,
                            epochs_fraction: 1.0,
                            cv_folds: 10,
                        })
                    } else {
                        Ok(FidelityLevel::Medium {
                            sample_fraction: 0.5,
                            epochs_fraction: 0.5,
                            cv_folds: 5,
                        })
                    }
                }
            }
            FidelitySelectionMethod::CostAware { budget_fraction } => {
                let used_budget_fraction = self
                    .evaluation_history
                    .iter()
                    .map(|e| e.cost)
                    .sum::<Float>()
                    / self.config.max_budget;

                if used_budget_fraction < *budget_fraction {
                    Ok(FidelityLevel::Low {
                        sample_fraction: 0.1,
                        epochs_fraction: 0.1,
                        cv_folds: 3,
                    })
                } else {
                    Ok(FidelityLevel::High {
                        sample_fraction: 1.0,
                        epochs_fraction: 1.0,
                        cv_folds: 10,
                    })
                }
            }
            _ => Ok(FidelityLevel::Medium {
                sample_fraction: 0.5,
                epochs_fraction: 0.5,
                cv_folds: 5,
            }),
        }
    }

    /// Optimize acquisition function
    fn optimize_acquisition(
        &mut self,
        acquisition_function: &AcquisitionFunction,
        parameter_bounds: &[(Float, Float)],
    ) -> Result<HashMap<String, Float>, Box<dyn std::error::Error>> {
        // Simplified acquisition optimization - random search
        let n_candidates = 100;
        let mut best_config = self.sample_random_configuration(parameter_bounds)?;
        let mut best_acquisition_value = Float::NEG_INFINITY;

        for _ in 0..n_candidates {
            let candidate = self.sample_random_configuration(parameter_bounds)?;
            let acquisition_value = self.evaluate_acquisition(&candidate, acquisition_function)?;

            if acquisition_value > best_acquisition_value {
                best_acquisition_value = acquisition_value;
                best_config = candidate;
            }
        }

        Ok(best_config)
    }

    /// Evaluate acquisition function
    fn evaluate_acquisition(
        &mut self,
        config: &HashMap<String, Float>,
        acquisition_function: &AcquisitionFunction,
    ) -> Result<Float, Box<dyn std::error::Error>> {
        // Simplified acquisition function evaluation
        match acquisition_function {
            AcquisitionFunction::ExpectedImprovement => {
                // Mock EI calculation
                let config_vec: Vec<Float> = config.values().cloned().collect();
                let config_sum = config_vec.iter().sum::<Float>();
                Ok(config_sum + self.rng.random::<Float>() * 0.1)
            }
            AcquisitionFunction::UpperConfidenceBound { beta } => {
                // Mock UCB calculation
                let config_vec: Vec<Float> = config.values().cloned().collect();
                let config_sum = config_vec.iter().sum::<Float>();
                Ok(config_sum + beta * self.rng.random::<Float>())
            }
            _ => {
                // Default to random value
                Ok(self.rng.random::<Float>())
            }
        }
    }

    /// Increase fidelity level
    fn increase_fidelity(&self, current: &FidelityLevel, max: &FidelityLevel) -> FidelityLevel {
        match (current, max) {
            (FidelityLevel::Low { .. }, _) => FidelityLevel::Medium {
                sample_fraction: 0.5,
                epochs_fraction: 0.5,
                cv_folds: 5,
            },
            (FidelityLevel::Medium { .. }, _) => FidelityLevel::High {
                sample_fraction: 1.0,
                epochs_fraction: 1.0,
                cv_folds: 10,
            },
            _ => current.clone(),
        }
    }

    /// Convert budget to fidelity level
    fn budget_to_fidelity(&self, budget: Float, fidelities: &[FidelityLevel]) -> FidelityLevel {
        if budget < 0.3 {
            fidelities
                .first()
                .unwrap_or(&FidelityLevel::Low {
                    sample_fraction: 0.1,
                    epochs_fraction: 0.1,
                    cv_folds: 3,
                })
                .clone()
        } else if budget < 0.7 {
            fidelities
                .get(1)
                .unwrap_or(&FidelityLevel::Medium {
                    sample_fraction: 0.5,
                    epochs_fraction: 0.5,
                    cv_folds: 5,
                })
                .clone()
        } else {
            fidelities
                .get(2)
                .unwrap_or(&FidelityLevel::High {
                    sample_fraction: 1.0,
                    epochs_fraction: 1.0,
                    cv_folds: 10,
                })
                .clone()
        }
    }

    /// Convert fidelity to string for tracking
    fn fidelity_to_string(&self, fidelity: &FidelityLevel) -> String {
        match fidelity {
            FidelityLevel::Low { .. } => "Low".to_string(),
            FidelityLevel::Medium { .. } => "Medium".to_string(),
            FidelityLevel::High { .. } => "High".to_string(),
            FidelityLevel::Custom { .. } => "Custom".to_string(),
        }
    }

    /// Update best configuration
    fn update_best(&mut self, evaluation: &FidelityEvaluation) -> bool {
        match &self.current_best {
            Some(current) => {
                if evaluation.score > current.score {
                    self.current_best = Some(evaluation.clone());
                    true
                } else {
                    false
                }
            }
            None => {
                self.current_best = Some(evaluation.clone());
                true
            }
        }
    }

    /// Check if optimization should stop
    fn should_stop(&self) -> bool {
        self.evaluation_history.len() >= self.config.max_evaluations
    }

    /// Get default fidelity level
    fn get_default_fidelity(&self) -> FidelityLevel {
        FidelityLevel::Medium {
            sample_fraction: 0.5,
            epochs_fraction: 0.5,
            cv_folds: 5,
        }
    }
}

impl MultiFidelityGP {
    /// Create a new multi-fidelity Gaussian Process
    fn new() -> Self {
        Self {
            observations: Vec::new(),
            hyperparameters: GPHyperparameters {
                length_scales: Array1::from_elem(1, 1.0),
                signal_variance: 1.0,
                noise_variance: 0.1,
                fidelity_correlation: 0.8,
            },
            trained: false,
        }
    }

    /// Update the GP with new observations
    fn update(
        &mut self,
        evaluations: &[FidelityEvaluation],
    ) -> Result<(), Box<dyn std::error::Error>> {
        self.observations.clear();

        for eval in evaluations {
            let params: Vec<Float> = eval.hyperparameters.values().cloned().collect();
            let fidelity_value = self.fidelity_to_value(&eval.fidelity);
            self.observations
                .push((Array1::from_vec(params), fidelity_value, eval.score));
        }

        // Simplified GP update
        self.trained = true;
        Ok(())
    }

    /// Convert fidelity to numerical value
    fn fidelity_to_value(&self, fidelity: &FidelityLevel) -> Float {
        match fidelity {
            FidelityLevel::Low { .. } => 0.1,
            FidelityLevel::Medium { .. } => 0.5,
            FidelityLevel::High { .. } => 1.0,
            FidelityLevel::Custom { relative_cost, .. } => *relative_cost,
        }
    }
}

/// Convenience function for multi-fidelity optimization
pub fn multi_fidelity_optimize<F>(
    evaluation_fn: F,
    parameter_bounds: &[(Float, Float)],
    config: Option<MultiFidelityConfig>,
) -> Result<MultiFidelityResult, Box<dyn std::error::Error>>
where
    F: Fn(
        &HashMap<String, Float>,
        &FidelityLevel,
    ) -> Result<FidelityEvaluation, Box<dyn std::error::Error>>,
{
    let config = config.unwrap_or_default();
    let mut optimizer = MultiFidelityOptimizer::new(config);
    optimizer.optimize(evaluation_fn, parameter_bounds)
}

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

    fn mock_evaluation_function(
        hyperparameters: &HashMap<String, Float>,
        fidelity: &FidelityLevel,
    ) -> Result<FidelityEvaluation, Box<dyn std::error::Error>> {
        let score = hyperparameters.values().sum::<Float>() * 0.1;
        let cost = match fidelity {
            FidelityLevel::Low { .. } => 1.0,
            FidelityLevel::Medium { .. } => 5.0,
            FidelityLevel::High { .. } => 10.0,
            FidelityLevel::Custom { relative_cost, .. } => *relative_cost * 10.0,
        };

        Ok(FidelityEvaluation {
            hyperparameters: hyperparameters.clone(),
            fidelity: fidelity.clone(),
            score,
            cost,
            evaluation_time: cost,
            uncertainty: Some(0.1),
            additional_metrics: HashMap::new(),
        })
    }

    #[test]
    fn test_multi_fidelity_optimizer_creation() {
        let config = MultiFidelityConfig::default();
        let optimizer = MultiFidelityOptimizer::new(config);
        assert_eq!(optimizer.evaluation_history.len(), 0);
    }

    #[test]
    fn test_multi_fidelity_optimization() {
        let config = MultiFidelityConfig {
            max_evaluations: 10,
            max_budget: 100.0,
            ..Default::default()
        };

        let parameter_bounds = vec![(0.0, 1.0), (0.0, 1.0)];

        let result =
            multi_fidelity_optimize(mock_evaluation_function, &parameter_bounds, Some(config))
                .unwrap();

        assert!(result.best_score >= 0.0);
        assert!(result.total_cost > 0.0);
        assert!(!result.optimization_history.is_empty());
    }

    #[test]
    fn test_fidelity_levels() {
        let low_fidelity = FidelityLevel::Low {
            sample_fraction: 0.1,
            epochs_fraction: 0.1,
            cv_folds: 3,
        };

        let evaluation = mock_evaluation_function(
            &HashMap::from([("param_0".to_string(), 0.5)]),
            &low_fidelity,
        )
        .unwrap();

        assert_eq!(evaluation.cost, 1.0);
    }

    #[test]
    fn test_successive_halving_strategy() {
        let config = MultiFidelityConfig {
            strategy: MultiFidelityStrategy::SuccessiveHalving {
                eta: 2.0,
                min_fidelity: FidelityLevel::Low {
                    sample_fraction: 0.1,
                    epochs_fraction: 0.1,
                    cv_folds: 3,
                },
                max_fidelity: FidelityLevel::High {
                    sample_fraction: 1.0,
                    epochs_fraction: 1.0,
                    cv_folds: 10,
                },
            },
            max_evaluations: 20,
            max_budget: 200.0,
            ..Default::default()
        };

        let parameter_bounds = vec![(0.0, 1.0), (0.0, 1.0)];

        let result =
            multi_fidelity_optimize(mock_evaluation_function, &parameter_bounds, Some(config))
                .unwrap();

        assert!(result.best_score >= 0.0);
        assert!(!result.fidelity_usage.is_empty());
    }
}