torsh-optim 0.1.3

//! Bayesian Optimization for Hyperparameter Tuning
//!
//! This module implements Bayesian optimization algorithms for efficient optimization
//! of expensive black-box functions. It's particularly useful for hyperparameter
//! optimization where function evaluations are costly.
//!
//! Key features:
//! - Gaussian Process surrogate models
//! - Various acquisition functions (EI, PI, UCB, etc.)
//! - Sequential model-based optimization
//! - Multi-objective optimization support
//! - Constraint handling
//! - Parallel evaluation support

use crate::{OptimizerError, OptimizerResult};
use scirs2_core::RngExt;
use std::collections::HashMap;
use std::sync::Arc;
use torsh_core::device::CpuDevice;

/// Trait for objective functions in Bayesian optimization
pub trait ObjectiveFunction: Send + Sync {
    /// Evaluate the objective function (to be minimized)
    fn evaluate(&self, parameters: &[f32]) -> OptimizerResult<f32>;

    /// Get the dimension of the parameter space
    fn dimension(&self) -> usize;

    /// Get parameter bounds (required for Bayesian optimization)
    fn bounds(&self) -> (Vec<f32>, Vec<f32>);

    /// Get the name of the objective function
    fn name(&self) -> &str {
        "Unknown"
    }

    /// Whether the function is noisy (affects GP model)
    fn is_noisy(&self) -> bool {
        false
    }
}

/// Configuration for Bayesian optimization
#[derive(Debug, Clone)]
pub struct BayesianOptConfig {
    /// Maximum number of function evaluations
    pub max_evaluations: usize,
    /// Number of initial random evaluations
    pub initial_random_evaluations: usize,
    /// Acquisition function to use
    pub acquisition_function: AcquisitionFunction,
    /// Number of optimization restarts for acquisition function
    pub acquisition_restarts: usize,
    /// Gaussian Process hyperparameters
    pub gp_config: GaussianProcessConfig,
    /// Random seed for reproducibility
    pub seed: Option<u64>,
    /// Device for computations
    pub device: Arc<CpuDevice>,
    /// Verbose output
    pub verbose: bool,
}

impl Default for BayesianOptConfig {
    fn default() -> Self {
        Self {
            max_evaluations: 100,
            initial_random_evaluations: 10,
            acquisition_function: AcquisitionFunction::ExpectedImprovement { xi: 0.01 },
            acquisition_restarts: 10,
            gp_config: GaussianProcessConfig::default(),
            seed: None,
            device: Arc::new(CpuDevice::new()),
            verbose: false,
        }
    }
}

/// Available acquisition functions
#[derive(Debug, Clone)]
pub enum AcquisitionFunction {
    /// Expected Improvement
    ExpectedImprovement { xi: f32 },
    /// Probability of Improvement
    ProbabilityOfImprovement { xi: f32 },
    /// Upper Confidence Bound
    UpperConfidenceBound { kappa: f32 },
    /// Thompson Sampling
    ThompsonSampling,
    /// Entropy Search
    EntropySearch,
}

/// Gaussian Process configuration
#[derive(Debug, Clone)]
pub struct GaussianProcessConfig {
    /// Kernel type
    pub kernel: KernelType,
    /// Noise level
    pub noise_level: f32,
    /// Length scale bounds for optimization
    pub length_scale_bounds: (f32, f32),
    /// Signal variance bounds
    pub signal_variance_bounds: (f32, f32),
    /// Whether to optimize hyperparameters
    pub optimize_hyperparameters: bool,
}

impl Default for GaussianProcessConfig {
    fn default() -> Self {
        Self {
            kernel: KernelType::RBF { length_scale: 1.0 },
            noise_level: 1e-6,
            length_scale_bounds: (1e-3, 1e3),
            signal_variance_bounds: (1e-3, 1e3),
            optimize_hyperparameters: true,
        }
    }
}

/// Available kernel types
#[derive(Debug, Clone)]
pub enum KernelType {
    /// Radial Basis Function (RBF) / Squared Exponential
    RBF { length_scale: f32 },
    /// Matern kernel with nu = 1.5
    Matern15 { length_scale: f32 },
    /// Matern kernel with nu = 2.5
    Matern25 { length_scale: f32 },
    /// Linear kernel
    Linear { variance: f32 },
    /// Rational Quadratic kernel
    RationalQuadratic { length_scale: f32, alpha: f32 },
}

/// Data point in the optimization history
#[derive(Debug, Clone)]
pub struct DataPoint {
    /// Parameter values
    pub parameters: Vec<f32>,
    /// Objective function value
    pub value: f32,
    /// Evaluation time (optional)
    pub evaluation_time: Option<std::time::Duration>,
}

/// Result from Bayesian optimization
#[derive(Debug, Clone)]
pub struct BayesianOptResult {
    /// Best point found
    pub best_point: DataPoint,
    /// All evaluated points
    pub history: Vec<DataPoint>,
    /// Number of function evaluations
    pub evaluations: usize,
    /// Whether optimization converged
    pub converged: bool,
    /// Convergence reason
    pub convergence_reason: String,
    /// Acquisition function values at each iteration
    pub acquisition_history: Vec<f32>,
}

/// Simplified Gaussian Process implementation
#[derive(Debug)]
pub struct GaussianProcess {
    config: GaussianProcessConfig,
    /// Training inputs (X)
    x_train: Vec<Vec<f32>>,
    /// Training outputs (y)
    y_train: Vec<f32>,
    /// Gram matrix (K + noise*I)
    gram_matrix: Vec<Vec<f32>>,
    /// Cholesky decomposition of gram matrix
    chol_gram: Vec<Vec<f32>>,
    /// Alpha vector for predictions
    alpha: Vec<f32>,
    /// Mean of training data
    y_mean: f32,
    /// Standard deviation of training data
    y_std: f32,
    /// Whether the model has been fitted
    fitted: bool,
}

impl GaussianProcess {
    pub fn new(config: GaussianProcessConfig) -> Self {
        Self {
            config,
            x_train: Vec::new(),
            y_train: Vec::new(),
            gram_matrix: Vec::new(),
            chol_gram: Vec::new(),
            alpha: Vec::new(),
            y_mean: 0.0,
            y_std: 1.0,
            fitted: false,
        }
    }

    pub fn fit(&mut self, x_train: Vec<Vec<f32>>, y_train: Vec<f32>) -> OptimizerResult<()> {
        if x_train.is_empty() || y_train.is_empty() || x_train.len() != y_train.len() {
            return Err(OptimizerError::InvalidInput(
                "Invalid training data".to_string(),
            ));
        }

        self.x_train = x_train;
        self.y_train = y_train;

        // Normalize targets
        self.y_mean = self.y_train.iter().sum::<f32>() / self.y_train.len() as f32;
        let var = self
            .y_train
            .iter()
            .map(|&y| (y - self.y_mean).powi(2))
            .sum::<f32>()
            / self.y_train.len() as f32;
        self.y_std = var.sqrt().max(1e-6);

        let y_normalized: Vec<f32> = self
            .y_train
            .iter()
            .map(|&y| (y - self.y_mean) / self.y_std)
            .collect();

        // Build gram matrix
        let n = self.x_train.len();
        self.gram_matrix = vec![vec![0.0; n]; n];

        for i in 0..n {
            for j in 0..n {
                self.gram_matrix[i][j] = self.kernel(&self.x_train[i], &self.x_train[j]);
                if i == j {
                    self.gram_matrix[i][j] += self.config.noise_level;
                }
            }
        }

        // Cholesky decomposition (simplified)
        self.chol_gram = self.cholesky_decomposition(&self.gram_matrix)?;

        // Solve for alpha
        self.alpha = self.solve_triangular(&self.chol_gram, &y_normalized)?;

        self.fitted = true;
        Ok(())
    }

    pub fn predict(&self, x_test: &[f32]) -> OptimizerResult<(f32, f32)> {
        if !self.fitted {
            return Err(OptimizerError::InvalidInput("Model not fitted".to_string()));
        }

        // Compute kernel vector
        let k_star: Vec<f32> = self
            .x_train
            .iter()
            .map(|x_train| self.kernel(x_test, x_train))
            .collect();

        // Predictive mean
        let mut mean = 0.0;
        for i in 0..self.alpha.len() {
            mean += self.alpha[i] * k_star[i];
        }

        // Denormalize
        mean = mean * self.y_std + self.y_mean;

        // Predictive variance (simplified)
        let k_star_star = self.kernel(x_test, x_test);
        let v = self.solve_triangular(&self.chol_gram, &k_star)?;
        let var = k_star_star - v.iter().map(|x| x * x).sum::<f32>();
        let std = (var.max(0.0).sqrt() * self.y_std).max(1e-6);

        Ok((mean, std))
    }

    fn kernel(&self, x1: &[f32], x2: &[f32]) -> f32 {
        match &self.config.kernel {
            KernelType::RBF { length_scale } => {
                let dist_sq = x1
                    .iter()
                    .zip(x2.iter())
                    .map(|(a, b)| (a - b).powi(2))
                    .sum::<f32>();
                (-0.5 * dist_sq / length_scale.powi(2)).exp()
            }
            KernelType::Matern15 { length_scale } => {
                let dist = x1
                    .iter()
                    .zip(x2.iter())
                    .map(|(a, b)| (a - b).powi(2))
                    .sum::<f32>()
                    .sqrt();
                let scaled_dist = dist * 3.0_f32.sqrt() / length_scale;
                (1.0 + scaled_dist) * (-scaled_dist).exp()
            }
            KernelType::Matern25 { length_scale } => {
                let dist = x1
                    .iter()
                    .zip(x2.iter())
                    .map(|(a, b)| (a - b).powi(2))
                    .sum::<f32>()
                    .sqrt();
                let scaled_dist = dist * 5.0_f32.sqrt() / length_scale;
                (1.0 + scaled_dist + scaled_dist.powi(2) / 3.0) * (-scaled_dist).exp()
            }
            KernelType::Linear { variance } => {
                variance * x1.iter().zip(x2.iter()).map(|(a, b)| a * b).sum::<f32>()
            }
            KernelType::RationalQuadratic {
                length_scale,
                alpha,
            } => {
                let dist_sq = x1
                    .iter()
                    .zip(x2.iter())
                    .map(|(a, b)| (a - b).powi(2))
                    .sum::<f32>();
                (1.0 + dist_sq / (2.0 * alpha * length_scale.powi(2))).powf(-alpha)
            }
        }
    }

    fn cholesky_decomposition(&self, matrix: &[Vec<f32>]) -> OptimizerResult<Vec<Vec<f32>>> {
        let n = matrix.len();
        let mut l = vec![vec![0.0; n]; n];

        for i in 0..n {
            for j in 0..=i {
                if i == j {
                    let sum: f32 = (0..j).map(|k| (l[j][k] as f32).powi(2)).sum();
                    let val = matrix[j][j] - sum;
                    if val <= 0.0 {
                        return Err(OptimizerError::NumericalError(
                            "Matrix not positive definite".to_string(),
                        ));
                    }
                    l[j][j] = val.sqrt();
                } else {
                    let sum: f32 = (0..j).map(|k| l[i][k] * l[j][k]).sum();
                    l[i][j] = (matrix[i][j] - sum) / l[j][j];
                }
            }
        }

        Ok(l)
    }

    fn solve_triangular(&self, l: &[Vec<f32>], b: &[f32]) -> OptimizerResult<Vec<f32>> {
        let n = l.len();
        let mut x = vec![0.0; n];

        // Forward substitution
        for i in 0..n {
            let sum: f32 = (0..i).map(|j| l[i][j] * x[j]).sum();
            x[i] = (b[i] - sum) / l[i][i];
        }

        // Backward substitution
        for i in (0..n).rev() {
            let sum: f32 = ((i + 1)..n).map(|j| l[j][i] * x[j]).sum();
            x[i] = (x[i] - sum) / l[i][i];
        }

        Ok(x)
    }
}

/// Bayesian optimization algorithm
#[derive(Debug)]
pub struct BayesianOptimizer {
    config: BayesianOptConfig,
    gp: GaussianProcess,
    history: Vec<DataPoint>,
    best_point: Option<DataPoint>,
}

impl BayesianOptimizer {
    pub fn new(config: BayesianOptConfig) -> Self {
        let gp = GaussianProcess::new(config.gp_config.clone());

        Self {
            config,
            gp,
            history: Vec::new(),
            best_point: None,
        }
    }

    pub fn optimize<F: ObjectiveFunction>(
        &mut self,
        objective: &F,
    ) -> OptimizerResult<BayesianOptResult> {
        let dimension = objective.dimension();
        let (lower_bounds, upper_bounds) = objective.bounds();

        if lower_bounds.len() != dimension || upper_bounds.len() != dimension {
            return Err(OptimizerError::InvalidInput(
                "Bounds dimension mismatch".to_string(),
            ));
        }

        // ✅ SciRS2 Policy Compliant - Using scirs2_core::random instead of direct rand
        use scirs2_core::random::{Random, Rng};

        let mut rng = if let Some(seed) = self.config.seed {
            Random::seed(seed)
        } else {
            Random::seed(0)
        };

        let mut acquisition_history = Vec::new();

        // Initial random evaluations
        for i in 0..self
            .config
            .initial_random_evaluations
            .min(self.config.max_evaluations)
        {
            let params = self.sample_random_point(&mut rng, &lower_bounds, &upper_bounds);
            let start_time = std::time::Instant::now();
            let value = objective.evaluate(&params)?;
            let evaluation_time = start_time.elapsed();

            let point = DataPoint {
                parameters: params,
                value,
                evaluation_time: Some(evaluation_time),
            };

            self.add_observation(point.clone());

            if self.config.verbose {
                println!(
                    "Initial evaluation {}: f({:?}) = {:.6e}",
                    i + 1,
                    point.parameters,
                    point.value
                );
            }
        }

        // Main optimization loop
        let mut evaluations = self.config.initial_random_evaluations;

        while evaluations < self.config.max_evaluations {
            // Fit GP model
            let x_train: Vec<Vec<f32>> =
                self.history.iter().map(|p| p.parameters.clone()).collect();
            let y_train: Vec<f32> = self.history.iter().map(|p| p.value).collect();

            self.gp.fit(x_train, y_train)?;

            // Optimize acquisition function
            let next_point = self.optimize_acquisition(&mut rng, &lower_bounds, &upper_bounds)?;

            // Evaluate objective at next point
            let start_time = std::time::Instant::now();
            let value = objective.evaluate(&next_point)?;
            let evaluation_time = start_time.elapsed();

            let point = DataPoint {
                parameters: next_point,
                value,
                evaluation_time: Some(evaluation_time),
            };

            // Compute acquisition value for history
            let acq_value = self.compute_acquisition(&point.parameters)?;
            acquisition_history.push(acq_value);

            self.add_observation(point.clone());
            evaluations += 1;

            if self.config.verbose {
                println!(
                    "Iteration {}: f({:?}) = {:.6e}, Best = {:.6e}",
                    evaluations,
                    point.parameters,
                    point.value,
                    self.best_point
                        .as_ref()
                        .expect("best_point should exist after add_observation")
                        .value
                );
            }
        }

        Ok(BayesianOptResult {
            best_point: self
                .best_point
                .clone()
                .expect("best_point should exist after optimization"),
            history: self.history.clone(),
            evaluations,
            converged: false, // Could implement convergence criteria
            convergence_reason: "Maximum evaluations reached".to_string(),
            acquisition_history,
        })
    }

    fn add_observation(&mut self, point: DataPoint) {
        if self.best_point.is_none()
            || point.value
                < self
                    .best_point
                    .as_ref()
                    .expect("best_point should exist after is_none check")
                    .value
        {
            self.best_point = Some(point.clone());
        }
        self.history.push(point);
    }

    fn sample_random_point<R: scirs2_core::random::Rng>(
        &self,
        rng: &mut R,
        lower: &[f32],
        upper: &[f32],
    ) -> Vec<f32> {
        (0..lower.len())
            .map(|i| rng.random::<f32>() * (upper[i] - lower[i]) + lower[i])
            .collect()
    }

    fn optimize_acquisition<R: scirs2_core::random::Rng>(
        &self,
        rng: &mut R,
        lower: &[f32],
        upper: &[f32],
    ) -> OptimizerResult<Vec<f32>> {
        let mut best_point = Vec::new();
        let mut best_acquisition = f32::NEG_INFINITY;

        // Multi-restart optimization of acquisition function
        for _ in 0..self.config.acquisition_restarts {
            let start_point = self.sample_random_point(rng, lower, upper);
            let optimized_point = self.local_optimize_acquisition(start_point, lower, upper)?;
            let acq_value = self.compute_acquisition(&optimized_point)?;

            if acq_value > best_acquisition {
                best_acquisition = acq_value;
                best_point = optimized_point;
            }
        }

        if best_point.is_empty() {
            // Fallback to random sampling
            best_point = self.sample_random_point(rng, lower, upper);
        }

        Ok(best_point)
    }

    fn local_optimize_acquisition(
        &self,
        start_point: Vec<f32>,
        lower: &[f32],
        upper: &[f32],
    ) -> OptimizerResult<Vec<f32>> {
        // Simple gradient-free optimization using random search
        let mut best_point = start_point;
        let mut best_value = self.compute_acquisition(&best_point)?;

        // ✅ SciRS2 Policy Compliant - Using scirs2_core::random instead of direct rand
        use scirs2_core::random::{Random, Rng};
        let mut rng = Random::seed(0);
        let step_size = 0.1;

        for _ in 0..100 {
            let mut candidate = best_point.clone();

            // Random perturbation
            for i in 0..candidate.len() {
                let perturbation =
                    (rng.random::<f32>() * 2.0 - 1.0) * step_size * (upper[i] - lower[i]);
                candidate[i] = (candidate[i] + perturbation).max(lower[i]).min(upper[i]);
            }

            let value = self.compute_acquisition(&candidate)?;
            if value > best_value {
                best_value = value;
                best_point = candidate;
            }
        }

        Ok(best_point)
    }

    fn compute_acquisition(&self, point: &[f32]) -> OptimizerResult<f32> {
        let (mean, std) = self.gp.predict(point)?;
        let best_value = self
            .best_point
            .as_ref()
            .expect("best_point should exist for acquisition computation")
            .value;

        match &self.config.acquisition_function {
            AcquisitionFunction::ExpectedImprovement { xi } => {
                let improvement = best_value - mean - xi;
                if std <= 0.0 {
                    return Ok(0.0);
                }
                let z = improvement / std;
                Ok(improvement * self.normal_cdf(z) + std * self.normal_pdf(z))
            }
            AcquisitionFunction::ProbabilityOfImprovement { xi } => {
                let improvement = best_value - mean - xi;
                if std <= 0.0 {
                    return Ok(0.0);
                }
                let z = improvement / std;
                Ok(self.normal_cdf(z))
            }
            AcquisitionFunction::UpperConfidenceBound { kappa } => {
                Ok(-(mean - kappa * std)) // Negative because we're maximizing acquisition but minimizing objective
            }
            _ => {
                // Simplified implementations for other acquisition functions
                Ok(-(mean - std)) // Default to UCB-like behavior
            }
        }
    }

    fn normal_cdf(&self, x: f32) -> f32 {
        // Approximation of the standard normal CDF
        0.5 * (1.0 + self.erf(x / 2.0_f32.sqrt()))
    }

    fn normal_pdf(&self, x: f32) -> f32 {
        // Standard normal PDF
        (2.0 * std::f32::consts::PI).sqrt().recip() * (-0.5 * x * x).exp()
    }

    fn erf(&self, x: f32) -> f32 {
        // Approximation of the error function
        let a1 = 0.254829592;
        let a2 = -0.284496736;
        let a3 = 1.421413741;
        let a4 = -1.453152027;
        let a5 = 1.061405429;
        let p = 0.3275911;

        let sign = if x < 0.0 { -1.0 } else { 1.0 };
        let x = x.abs();

        let t = 1.0 / (1.0 + p * x);
        let y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * (-x * x).exp();

        sign * y
    }
}

/// Example objective functions for testing
pub mod test_functions {
    use super::*;

    /// 1D quadratic function: f(x) = (x - 0.5)^2
    pub struct Quadratic1D;

    impl ObjectiveFunction for Quadratic1D {
        fn evaluate(&self, parameters: &[f32]) -> OptimizerResult<f32> {
            if parameters.len() != 1 {
                return Err(OptimizerError::InvalidInput(
                    "Expected 1D input".to_string(),
                ));
            }
            Ok((parameters[0] - 0.5).powi(2))
        }

        fn dimension(&self) -> usize {
            1
        }

        fn bounds(&self) -> (Vec<f32>, Vec<f32>) {
            (vec![0.0], vec![1.0])
        }

        fn name(&self) -> &str {
            "Quadratic 1D"
        }
    }

    /// 2D Branin function
    pub struct Branin;

    impl ObjectiveFunction for Branin {
        fn evaluate(&self, parameters: &[f32]) -> OptimizerResult<f32> {
            if parameters.len() != 2 {
                return Err(OptimizerError::InvalidInput(
                    "Expected 2D input".to_string(),
                ));
            }

            let x1 = parameters[0];
            let x2 = parameters[1];

            let a = 1.0;
            let b = 5.1 / (4.0 * std::f32::consts::PI.powi(2));
            let c = 5.0 / std::f32::consts::PI;
            let r = 6.0;
            let s = 10.0;
            let t = 1.0 / (8.0 * std::f32::consts::PI);

            let term1 = a * (x2 - b * x1.powi(2) + c * x1 - r).powi(2);
            let term2 = s * (1.0 - t) * x1.cos();
            let term3 = s;

            Ok(term1 + term2 + term3)
        }

        fn dimension(&self) -> usize {
            2
        }

        fn bounds(&self) -> (Vec<f32>, Vec<f32>) {
            (vec![-5.0, 0.0], vec![10.0, 15.0])
        }

        fn name(&self) -> &str {
            "Branin"
        }
    }
}

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

    #[test]
    fn test_gaussian_process_fit_predict() {
        let config = GaussianProcessConfig::default();
        let mut gp = GaussianProcess::new(config);

        let x_train = vec![vec![0.0], vec![0.5], vec![1.0]];
        let y_train = vec![0.0, 0.25, 1.0];

        gp.fit(x_train, y_train).unwrap();

        let (mean, std) = gp.predict(&[0.75]).unwrap();
        assert!(mean >= 0.0 && mean <= 1.0);
        assert!(std > 0.0);
    }

    #[test]
    fn test_bayesian_optimization_quadratic() {
        let config = BayesianOptConfig {
            max_evaluations: 20,
            initial_random_evaluations: 5,
            acquisition_function: AcquisitionFunction::ExpectedImprovement { xi: 0.01 },
            seed: Some(42),
            verbose: false,
            ..Default::default()
        };

        let mut optimizer = BayesianOptimizer::new(config);
        let objective = Quadratic1D;

        let result = optimizer.optimize(&objective).unwrap();

        // Should find minimum near x = 0.5
        assert!(result.best_point.parameters[0] > 0.3 && result.best_point.parameters[0] < 0.7);
        assert!(result.best_point.value < 0.1);
        assert_eq!(result.evaluations, 20);
    }

    #[test]
    fn test_bayesian_optimization_branin() {
        let config = BayesianOptConfig {
            max_evaluations: 30,
            initial_random_evaluations: 10,
            acquisition_function: AcquisitionFunction::UpperConfidenceBound { kappa: 2.576 },
            seed: Some(42),
            verbose: false,
            ..Default::default()
        };

        let mut optimizer = BayesianOptimizer::new(config);
        let objective = Branin;

        let result = optimizer.optimize(&objective).unwrap();

        // Branin function has global minimum around 0.398
        assert!(result.best_point.value < 5.0); // Should find a reasonably good solution
        assert_eq!(result.evaluations, 30);
    }
}