quantrs2_core/

variational_optimization.rs

//! Enhanced variational parameter optimization using SciRS2
//!
//! This module provides advanced optimization techniques for variational quantum algorithms
//! leveraging SciRS2's optimization capabilities including:
//! - Gradient-based methods (BFGS, L-BFGS, Conjugate Gradient)
//! - Gradient-free methods (Nelder-Mead, Powell, COBYLA)
//! - Stochastic optimization (SPSA, Adam, RMSprop)
//! - Natural gradient descent for quantum circuits

use crate::{
    error::{QuantRS2Error, QuantRS2Result},
    variational::VariationalCircuit,
};
use scirs2_core::ndarray::{Array1, Array2};
// use scirs2_core::parallel_ops::*;
use crate::optimization_stubs::{minimize, Method, OptimizeResult, Options};
use crate::parallel_ops_stubs::*;
// use scirs2_core::optimization::{minimize, Method, OptimizeResult, Options};
use rustc_hash::FxHashMap;
use std::sync::{Arc, Mutex};

// Import SciRS2 optimization
// extern crate scirs2_optimize;
// use scirs2_optimize::unconstrained::{minimize, Method, Options};

// Import SciRS2 linear algebra for natural gradient
// extern crate scirs2_linalg;

/// Advanced optimizer for variational quantum circuits
pub struct VariationalQuantumOptimizer {
    /// Optimization method
    method: OptimizationMethod,
    /// Configuration
    config: OptimizationConfig,
    /// History of optimization
    history: OptimizationHistory,
    /// Fisher information matrix cache
    fisher_cache: Option<FisherCache>,
}

/// Optimization methods available
#[derive(Debug, Clone)]
pub enum OptimizationMethod {
    /// Standard gradient descent
    GradientDescent { learning_rate: f64 },
    /// Momentum-based gradient descent
    Momentum { learning_rate: f64, momentum: f64 },
    /// Adam optimizer
    Adam {
        learning_rate: f64,
        beta1: f64,
        beta2: f64,
        epsilon: f64,
    },
    /// RMSprop optimizer
    RMSprop {
        learning_rate: f64,
        decay_rate: f64,
        epsilon: f64,
    },
    /// Natural gradient descent
    NaturalGradient {
        learning_rate: f64,
        regularization: f64,
    },
    /// SciRS2 BFGS method
    BFGS,
    /// SciRS2 L-BFGS method
    LBFGS { memory_size: usize },
    /// SciRS2 Conjugate Gradient
    ConjugateGradient,
    /// SciRS2 Nelder-Mead simplex
    NelderMead,
    /// SciRS2 Powell's method
    Powell,
    /// Simultaneous Perturbation Stochastic Approximation
    SPSA {
        a: f64,
        c: f64,
        alpha: f64,
        gamma: f64,
    },
    /// Quantum Natural SPSA
    QNSPSA {
        learning_rate: f64,
        regularization: f64,
        spsa_epsilon: f64,
    },
}

/// Configuration for optimization
#[derive(Clone)]
pub struct OptimizationConfig {
    /// Maximum iterations
    pub max_iterations: usize,
    /// Function tolerance
    pub f_tol: f64,
    /// Gradient tolerance
    pub g_tol: f64,
    /// Parameter tolerance
    pub x_tol: f64,
    /// Enable parallel gradient computation
    pub parallel_gradients: bool,
    /// Batch size for stochastic methods
    pub batch_size: Option<usize>,
    /// Random seed
    pub seed: Option<u64>,
    /// Callback function after each iteration
    pub callback: Option<Arc<dyn Fn(&[f64], f64) + Send + Sync>>,
    /// Early stopping patience
    pub patience: Option<usize>,
    /// Gradient clipping value
    pub grad_clip: Option<f64>,
}

impl Default for OptimizationConfig {
    fn default() -> Self {
        Self {
            max_iterations: 100,
            f_tol: 1e-8,
            g_tol: 1e-8,
            x_tol: 1e-8,
            parallel_gradients: true,
            batch_size: None,
            seed: None,
            callback: None,
            patience: None,
            grad_clip: None,
        }
    }
}

/// Optimization history tracking
#[derive(Debug, Clone)]
pub struct OptimizationHistory {
    /// Parameter values at each iteration
    pub parameters: Vec<Vec<f64>>,
    /// Loss values
    pub loss_values: Vec<f64>,
    /// Gradient norms
    pub gradient_norms: Vec<f64>,
    /// Iteration times (ms)
    pub iteration_times: Vec<f64>,
    /// Total iterations
    pub total_iterations: usize,
    /// Converged flag
    pub converged: bool,
}

impl OptimizationHistory {
    fn new() -> Self {
        Self {
            parameters: Vec::new(),
            loss_values: Vec::new(),
            gradient_norms: Vec::new(),
            iteration_times: Vec::new(),
            total_iterations: 0,
            converged: false,
        }
    }
}

/// Fisher information matrix cache
struct FisherCache {
    /// Cached Fisher matrix
    matrix: Arc<Mutex<Option<Array2<f64>>>>,
    /// Parameters for cached matrix
    params: Arc<Mutex<Option<Vec<f64>>>>,
    /// Cache validity threshold
    threshold: f64,
}

/// Optimizer state for stateful methods
struct OptimizerState {
    /// Momentum vectors
    momentum: FxHashMap<String, f64>,
    /// Adam first moment
    adam_m: FxHashMap<String, f64>,
    /// Adam second moment
    adam_v: FxHashMap<String, f64>,
    /// RMSprop moving average
    rms_avg: FxHashMap<String, f64>,
    /// Iteration counter
    iteration: usize,
}

impl VariationalQuantumOptimizer {
    /// Create a new optimizer
    pub fn new(method: OptimizationMethod, config: OptimizationConfig) -> Self {
        let fisher_cache = match &method {
            OptimizationMethod::NaturalGradient { .. } | OptimizationMethod::QNSPSA { .. } => {
                Some(FisherCache {
                    matrix: Arc::new(Mutex::new(None)),
                    params: Arc::new(Mutex::new(None)),
                    threshold: 1e-3,
                })
            }
            _ => None,
        };

        Self {
            method,
            config,
            history: OptimizationHistory::new(),
            fisher_cache,
        }
    }

    /// Optimize a variational circuit
    pub fn optimize(
        &mut self,
        circuit: &mut VariationalCircuit,
        cost_fn: impl Fn(&VariationalCircuit) -> QuantRS2Result<f64> + Send + Sync + 'static,
    ) -> QuantRS2Result<OptimizationResult> {
        let cost_fn = Arc::new(cost_fn);

        match &self.method {
            OptimizationMethod::BFGS
            | OptimizationMethod::LBFGS { .. }
            | OptimizationMethod::ConjugateGradient
            | OptimizationMethod::NelderMead
            | OptimizationMethod::Powell => self.optimize_with_scirs2(circuit, cost_fn),
            _ => self.optimize_custom(circuit, cost_fn),
        }
    }

    /// Optimize using SciRS2 methods
    fn optimize_with_scirs2(
        &mut self,
        circuit: &mut VariationalCircuit,
        cost_fn: Arc<dyn Fn(&VariationalCircuit) -> QuantRS2Result<f64> + Send + Sync>,
    ) -> QuantRS2Result<OptimizationResult> {
        let param_names = circuit.parameter_names();
        let initial_params: Vec<f64> = param_names
            .iter()
            .map(|name| circuit.get_parameters().get(name).copied().unwrap_or(0.0))
            .collect();

        let circuit_clone = Arc::new(Mutex::new(circuit.clone()));
        let param_names_clone = param_names.clone();

        // Create objective function for SciRS2
        let objective = move |params: &scirs2_core::ndarray::ArrayView1<f64>| -> f64 {
            let params_slice = params.as_slice().unwrap();
            let mut param_map = FxHashMap::default();
            for (name, &value) in param_names_clone.iter().zip(params_slice) {
                param_map.insert(name.clone(), value);
            }

            let mut circuit = circuit_clone.lock().unwrap();
            if circuit.set_parameters(&param_map).is_err() {
                return f64::INFINITY;
            }

            match cost_fn(&*circuit) {
                Ok(loss) => loss,
                Err(_) => f64::INFINITY,
            }
        };

        // Set up SciRS2 method
        let method = match &self.method {
            OptimizationMethod::BFGS => Method::BFGS,
            OptimizationMethod::LBFGS { memory_size: _ } => Method::LBFGS,
            OptimizationMethod::ConjugateGradient => Method::BFGS, // Use BFGS as fallback
            OptimizationMethod::NelderMead => Method::NelderMead,
            OptimizationMethod::Powell => Method::Powell,
            _ => unreachable!(),
        };

        // Configure options
        let options = Options {
            max_iter: self.config.max_iterations,
            ftol: self.config.f_tol,
            gtol: self.config.g_tol,
            xtol: self.config.x_tol,
            ..Default::default()
        };

        // Run optimization
        let start_time = std::time::Instant::now();
        let initial_array = scirs2_core::ndarray::Array1::from_vec(initial_params.clone());
        let result = minimize(objective, &initial_array, method, Some(options))
            .map_err(|e| QuantRS2Error::InvalidInput(format!("Optimization failed: {:?}", e)))?;

        // Update circuit with optimal parameters
        let mut final_params = FxHashMap::default();
        for (name, &value) in param_names.iter().zip(result.x.as_slice().unwrap()) {
            final_params.insert(name.clone(), value);
        }
        circuit.set_parameters(&final_params)?;

        // Update history
        self.history.parameters.push(result.x.to_vec());
        self.history.loss_values.push(result.fun);
        self.history.total_iterations = result.iterations;
        self.history.converged = result.success;

        Ok(OptimizationResult {
            optimal_parameters: final_params,
            final_loss: result.fun,
            iterations: result.iterations,
            converged: result.success,
            optimization_time: start_time.elapsed().as_secs_f64(),
            history: self.history.clone(),
        })
    }

    /// Optimize using custom methods
    fn optimize_custom(
        &mut self,
        circuit: &mut VariationalCircuit,
        cost_fn: Arc<dyn Fn(&VariationalCircuit) -> QuantRS2Result<f64> + Send + Sync>,
    ) -> QuantRS2Result<OptimizationResult> {
        let mut state = OptimizerState {
            momentum: FxHashMap::default(),
            adam_m: FxHashMap::default(),
            adam_v: FxHashMap::default(),
            rms_avg: FxHashMap::default(),
            iteration: 0,
        };

        let param_names = circuit.parameter_names();
        let start_time = std::time::Instant::now();
        let mut best_loss = f64::INFINITY;
        let mut patience_counter = 0;

        for iter in 0..self.config.max_iterations {
            let iter_start = std::time::Instant::now();

            // Compute loss
            let loss = cost_fn(circuit)?;

            // Check for improvement
            if loss < best_loss - self.config.f_tol {
                best_loss = loss;
                patience_counter = 0;
            } else if let Some(patience) = self.config.patience {
                patience_counter += 1;
                if patience_counter >= patience {
                    self.history.converged = true;
                    break;
                }
            }

            // Compute gradients
            let gradients = self.compute_gradients(circuit, &cost_fn)?;

            // Clip gradients if requested
            let gradients = if let Some(max_norm) = self.config.grad_clip {
                self.clip_gradients(gradients, max_norm)
            } else {
                gradients
            };

            // Update parameters based on method
            self.update_parameters(circuit, &gradients, &mut state)?;

            // Update history
            let current_params: Vec<f64> = param_names
                .iter()
                .map(|name| circuit.get_parameters().get(name).copied().unwrap_or(0.0))
                .collect();

            let grad_norm = gradients.values().map(|g| g * g).sum::<f64>().sqrt();

            self.history.parameters.push(current_params);
            self.history.loss_values.push(loss);
            self.history.gradient_norms.push(grad_norm);
            self.history
                .iteration_times
                .push(iter_start.elapsed().as_secs_f64() * 1000.0);
            self.history.total_iterations = iter + 1;

            // Callback
            if let Some(callback) = &self.config.callback {
                let params: Vec<f64> = param_names
                    .iter()
                    .map(|name| circuit.get_parameters().get(name).copied().unwrap_or(0.0))
                    .collect();
                callback(&params, loss);
            }

            // Check convergence
            if grad_norm < self.config.g_tol {
                self.history.converged = true;
                break;
            }

            state.iteration += 1;
        }

        let final_params = circuit.get_parameters();
        let final_loss = cost_fn(circuit)?;

        Ok(OptimizationResult {
            optimal_parameters: final_params,
            final_loss,
            iterations: self.history.total_iterations,
            converged: self.history.converged,
            optimization_time: start_time.elapsed().as_secs_f64(),
            history: self.history.clone(),
        })
    }

    /// Compute gradients for all parameters
    fn compute_gradients(
        &self,
        circuit: &VariationalCircuit,
        cost_fn: &Arc<dyn Fn(&VariationalCircuit) -> QuantRS2Result<f64> + Send + Sync>,
    ) -> QuantRS2Result<FxHashMap<String, f64>> {
        let param_names = circuit.parameter_names();

        if self.config.parallel_gradients {
            // Parallel gradient computation
            let gradients: Vec<(String, f64)> = param_names
                .par_iter()
                .map(|param_name| {
                    let grad = self
                        .compute_single_gradient(circuit, param_name, cost_fn)
                        .unwrap_or(0.0);
                    (param_name.clone(), grad)
                })
                .collect();

            Ok(gradients.into_iter().collect())
        } else {
            // Sequential gradient computation
            let mut gradients = FxHashMap::default();
            for param_name in &param_names {
                let grad = self.compute_single_gradient(circuit, param_name, cost_fn)?;
                gradients.insert(param_name.clone(), grad);
            }
            Ok(gradients)
        }
    }

    /// Compute gradient for a single parameter
    fn compute_single_gradient(
        &self,
        circuit: &VariationalCircuit,
        param_name: &str,
        cost_fn: &Arc<dyn Fn(&VariationalCircuit) -> QuantRS2Result<f64> + Send + Sync>,
    ) -> QuantRS2Result<f64> {
        match &self.method {
            OptimizationMethod::SPSA { c, .. } => {
                // SPSA gradient approximation
                self.spsa_gradient(circuit, param_name, cost_fn, *c)
            }
            _ => {
                // Parameter shift rule
                self.parameter_shift_gradient(circuit, param_name, cost_fn)
            }
        }
    }

    /// Parameter shift rule gradient
    fn parameter_shift_gradient(
        &self,
        circuit: &VariationalCircuit,
        param_name: &str,
        cost_fn: &Arc<dyn Fn(&VariationalCircuit) -> QuantRS2Result<f64> + Send + Sync>,
    ) -> QuantRS2Result<f64> {
        let current_params = circuit.get_parameters();
        let current_value = *current_params.get(param_name).ok_or_else(|| {
            QuantRS2Error::InvalidInput(format!("Parameter {} not found", param_name))
        })?;

        // Shift parameter by +π/2
        let mut circuit_plus = circuit.clone();
        let mut params_plus = current_params.clone();
        params_plus.insert(
            param_name.to_string(),
            current_value + std::f64::consts::PI / 2.0,
        );
        circuit_plus.set_parameters(&params_plus)?;
        let loss_plus = cost_fn(&circuit_plus)?;

        // Shift parameter by -π/2
        let mut circuit_minus = circuit.clone();
        let mut params_minus = current_params.clone();
        params_minus.insert(
            param_name.to_string(),
            current_value - std::f64::consts::PI / 2.0,
        );
        circuit_minus.set_parameters(&params_minus)?;
        let loss_minus = cost_fn(&circuit_minus)?;

        Ok((loss_plus - loss_minus) / 2.0)
    }

    /// SPSA gradient approximation
    fn spsa_gradient(
        &self,
        circuit: &VariationalCircuit,
        param_name: &str,
        cost_fn: &Arc<dyn Fn(&VariationalCircuit) -> QuantRS2Result<f64> + Send + Sync>,
        epsilon: f64,
    ) -> QuantRS2Result<f64> {
        use scirs2_core::random::prelude::*;

        let mut rng = if let Some(seed) = self.config.seed {
            StdRng::seed_from_u64(seed)
        } else {
            StdRng::from_seed(thread_rng().gen())
        };

        let current_params = circuit.get_parameters();
        let perturbation = if rng.gen::<bool>() { epsilon } else { -epsilon };

        // Positive perturbation
        let mut circuit_plus = circuit.clone();
        let mut params_plus = current_params.clone();
        for (name, value) in params_plus.iter_mut() {
            if name == param_name {
                *value += perturbation;
            }
        }
        circuit_plus.set_parameters(&params_plus)?;
        let loss_plus = cost_fn(&circuit_plus)?;

        // Negative perturbation
        let mut circuit_minus = circuit.clone();
        let mut params_minus = current_params.clone();
        for (name, value) in params_minus.iter_mut() {
            if name == param_name {
                *value -= perturbation;
            }
        }
        circuit_minus.set_parameters(&params_minus)?;
        let loss_minus = cost_fn(&circuit_minus)?;

        Ok((loss_plus - loss_minus) / (2.0 * perturbation))
    }

    /// Clip gradients by norm
    fn clip_gradients(
        &self,
        mut gradients: FxHashMap<String, f64>,
        max_norm: f64,
    ) -> FxHashMap<String, f64> {
        let norm = gradients.values().map(|g| g * g).sum::<f64>().sqrt();

        if norm > max_norm {
            let scale = max_norm / norm;
            for grad in gradients.values_mut() {
                *grad *= scale;
            }
        }

        gradients
    }

    /// Update parameters based on optimization method
    fn update_parameters(
        &self,
        circuit: &mut VariationalCircuit,
        gradients: &FxHashMap<String, f64>,
        state: &mut OptimizerState,
    ) -> QuantRS2Result<()> {
        let mut new_params = circuit.get_parameters();

        match &self.method {
            OptimizationMethod::GradientDescent { learning_rate } => {
                // Simple gradient descent
                for (param_name, &grad) in gradients {
                    if let Some(value) = new_params.get_mut(param_name) {
                        *value -= learning_rate * grad;
                    }
                }
            }
            OptimizationMethod::Momentum {
                learning_rate,
                momentum,
            } => {
                // Momentum-based gradient descent
                for (param_name, &grad) in gradients {
                    let velocity = state.momentum.entry(param_name.clone()).or_insert(0.0);
                    *velocity = momentum * *velocity - learning_rate * grad;

                    if let Some(value) = new_params.get_mut(param_name) {
                        *value += *velocity;
                    }
                }
            }
            OptimizationMethod::Adam {
                learning_rate,
                beta1,
                beta2,
                epsilon,
            } => {
                // Adam optimizer
                let t = state.iteration as f64 + 1.0;
                let lr_t = learning_rate * (1.0 - beta2.powf(t)).sqrt() / (1.0 - beta1.powf(t));

                for (param_name, &grad) in gradients {
                    let m = state.adam_m.entry(param_name.clone()).or_insert(0.0);
                    let v = state.adam_v.entry(param_name.clone()).or_insert(0.0);

                    *m = beta1 * *m + (1.0 - beta1) * grad;
                    *v = beta2 * *v + (1.0 - beta2) * grad * grad;

                    if let Some(value) = new_params.get_mut(param_name) {
                        *value -= lr_t * *m / (v.sqrt() + epsilon);
                    }
                }
            }
            OptimizationMethod::RMSprop {
                learning_rate,
                decay_rate,
                epsilon,
            } => {
                // RMSprop optimizer
                for (param_name, &grad) in gradients {
                    let avg = state.rms_avg.entry(param_name.clone()).or_insert(0.0);
                    *avg = decay_rate * *avg + (1.0 - decay_rate) * grad * grad;

                    if let Some(value) = new_params.get_mut(param_name) {
                        *value -= learning_rate * grad / (avg.sqrt() + epsilon);
                    }
                }
            }
            OptimizationMethod::NaturalGradient {
                learning_rate,
                regularization,
            } => {
                // Natural gradient descent
                let fisher_inv =
                    self.compute_fisher_inverse(circuit, gradients, *regularization)?;
                let natural_grad = self.apply_fisher_inverse(&fisher_inv, gradients);

                for (param_name, &nat_grad) in &natural_grad {
                    if let Some(value) = new_params.get_mut(param_name) {
                        *value -= learning_rate * nat_grad;
                    }
                }
            }
            OptimizationMethod::SPSA { a, alpha, .. } => {
                // SPSA parameter update
                let ak = a / (state.iteration as f64 + 1.0).powf(*alpha);

                for (param_name, &grad) in gradients {
                    if let Some(value) = new_params.get_mut(param_name) {
                        *value -= ak * grad;
                    }
                }
            }
            OptimizationMethod::QNSPSA {
                learning_rate,
                regularization,
                ..
            } => {
                // Quantum Natural SPSA
                let fisher_inv =
                    self.compute_fisher_inverse(circuit, gradients, *regularization)?;
                let natural_grad = self.apply_fisher_inverse(&fisher_inv, gradients);

                for (param_name, &nat_grad) in &natural_grad {
                    if let Some(value) = new_params.get_mut(param_name) {
                        *value -= learning_rate * nat_grad;
                    }
                }
            }
            _ => {
                // Should not reach here for SciRS2 methods
                return Err(QuantRS2Error::InvalidInput(
                    "Invalid optimization method".to_string(),
                ));
            }
        }

        circuit.set_parameters(&new_params)
    }

    /// Compute Fisher information matrix inverse
    fn compute_fisher_inverse(
        &self,
        circuit: &VariationalCircuit,
        gradients: &FxHashMap<String, f64>,
        regularization: f64,
    ) -> QuantRS2Result<Array2<f64>> {
        let param_names: Vec<_> = gradients.keys().cloned().collect();
        let n_params = param_names.len();

        // Check cache
        if let Some(cache) = &self.fisher_cache {
            if let Some(cached_matrix) = cache.matrix.lock().unwrap().as_ref() {
                if let Some(cached_params) = cache.params.lock().unwrap().as_ref() {
                    let current_params: Vec<f64> = param_names
                        .iter()
                        .map(|name| circuit.get_parameters().get(name).copied().unwrap_or(0.0))
                        .collect();

                    let diff_norm: f64 = current_params
                        .iter()
                        .zip(cached_params.iter())
                        .map(|(a, b)| (a - b).powi(2))
                        .sum::<f64>()
                        .sqrt();

                    if diff_norm < cache.threshold {
                        return Ok(cached_matrix.clone());
                    }
                }
            }
        }

        // Compute Fisher information matrix
        let mut fisher = Array2::zeros((n_params, n_params));

        // Simplified Fisher matrix computation
        // In practice, this would involve quantum state overlaps
        for i in 0..n_params {
            for j in i..n_params {
                // Approximation: use gradient outer product
                let value = gradients[&param_names[i]] * gradients[&param_names[j]];
                fisher[[i, j]] = value;
                fisher[[j, i]] = value;
            }
        }

        // Add regularization
        for i in 0..n_params {
            fisher[[i, i]] += regularization;
        }

        // Compute inverse using simple matrix inversion
        // For now, use a simple inversion approach
        // TODO: Use ndarray-linalg when trait import issues are resolved
        let n = fisher.nrows();
        let mut fisher_inv = Array2::eye(n);

        // Simple inversion using Gaussian elimination (placeholder)
        // In practice, should use proper numerical methods
        if n == 1 {
            fisher_inv[[0, 0]] = 1.0 / fisher[[0, 0]];
        } else if n == 2 {
            let det = fisher[[0, 0]] * fisher[[1, 1]] - fisher[[0, 1]] * fisher[[1, 0]];
            if det.abs() < 1e-10 {
                return Err(QuantRS2Error::InvalidInput(
                    "Fisher matrix is singular".to_string(),
                ));
            }
            fisher_inv[[0, 0]] = fisher[[1, 1]] / det;
            fisher_inv[[0, 1]] = -fisher[[0, 1]] / det;
            fisher_inv[[1, 0]] = -fisher[[1, 0]] / det;
            fisher_inv[[1, 1]] = fisher[[0, 0]] / det;
        } else {
            // For larger matrices, return identity as placeholder
            // TODO: Implement proper inversion
        }

        // Update cache
        if let Some(cache) = &self.fisher_cache {
            let current_params: Vec<f64> = param_names
                .iter()
                .map(|name| circuit.get_parameters().get(name).copied().unwrap_or(0.0))
                .collect();

            *cache.matrix.lock().unwrap() = Some(fisher_inv.clone());
            *cache.params.lock().unwrap() = Some(current_params);
        }

        Ok(fisher_inv)
    }

    /// Apply Fisher information matrix inverse to gradients
    fn apply_fisher_inverse(
        &self,
        fisher_inv: &Array2<f64>,
        gradients: &FxHashMap<String, f64>,
    ) -> FxHashMap<String, f64> {
        let param_names: Vec<_> = gradients.keys().cloned().collect();
        let grad_vec: Vec<f64> = param_names.iter().map(|name| gradients[name]).collect();

        let grad_array = Array1::from_vec(grad_vec);
        let natural_grad = fisher_inv.dot(&grad_array);

        let mut result = FxHashMap::default();
        for (i, name) in param_names.iter().enumerate() {
            result.insert(name.clone(), natural_grad[i]);
        }

        result
    }
}

/// Optimization result
#[derive(Debug, Clone)]
pub struct OptimizationResult {
    /// Optimal parameters
    pub optimal_parameters: FxHashMap<String, f64>,
    /// Final loss value
    pub final_loss: f64,
    /// Number of iterations
    pub iterations: usize,
    /// Whether optimization converged
    pub converged: bool,
    /// Total optimization time (seconds)
    pub optimization_time: f64,
    /// Full optimization history
    pub history: OptimizationHistory,
}

/// Create optimized VQE optimizer
pub fn create_vqe_optimizer() -> VariationalQuantumOptimizer {
    let config = OptimizationConfig {
        max_iterations: 200,
        f_tol: 1e-10,
        g_tol: 1e-10,
        parallel_gradients: true,
        grad_clip: Some(1.0),
        ..Default::default()
    };

    VariationalQuantumOptimizer::new(OptimizationMethod::LBFGS { memory_size: 10 }, config)
}

/// Create optimized QAOA optimizer
pub fn create_qaoa_optimizer() -> VariationalQuantumOptimizer {
    let config = OptimizationConfig {
        max_iterations: 100,
        parallel_gradients: true,
        ..Default::default()
    };

    VariationalQuantumOptimizer::new(OptimizationMethod::BFGS, config)
}

/// Create natural gradient optimizer
pub fn create_natural_gradient_optimizer(learning_rate: f64) -> VariationalQuantumOptimizer {
    let config = OptimizationConfig {
        max_iterations: 100,
        parallel_gradients: true,
        ..Default::default()
    };

    VariationalQuantumOptimizer::new(
        OptimizationMethod::NaturalGradient {
            learning_rate,
            regularization: 1e-4,
        },
        config,
    )
}

/// Create SPSA optimizer for noisy quantum devices
pub fn create_spsa_optimizer() -> VariationalQuantumOptimizer {
    let config = OptimizationConfig {
        max_iterations: 500,
        seed: Some(42),
        ..Default::default()
    };

    VariationalQuantumOptimizer::new(
        OptimizationMethod::SPSA {
            a: 0.1,
            c: 0.1,
            alpha: 0.602,
            gamma: 0.101,
        },
        config,
    )
}

/// Constrained optimization for variational circuits
pub struct ConstrainedVariationalOptimizer {
    /// Base optimizer
    base_optimizer: VariationalQuantumOptimizer,
    /// Constraints
    constraints: Vec<Constraint>,
}

/// Constraint for optimization
#[derive(Clone)]
pub struct Constraint {
    /// Constraint function
    pub function: Arc<dyn Fn(&FxHashMap<String, f64>) -> f64 + Send + Sync>,
    /// Constraint type
    pub constraint_type: ConstraintType,
    /// Constraint value
    pub value: f64,
}

/// Constraint type
#[derive(Debug, Clone, Copy)]
pub enum ConstraintType {
    /// Equality constraint
    Eq,
    /// Inequality constraint
    Ineq,
}

impl ConstrainedVariationalOptimizer {
    /// Create a new constrained optimizer
    pub fn new(base_optimizer: VariationalQuantumOptimizer) -> Self {
        Self {
            base_optimizer,
            constraints: Vec::new(),
        }
    }

    /// Add an equality constraint
    pub fn add_equality_constraint(
        &mut self,
        constraint_fn: impl Fn(&FxHashMap<String, f64>) -> f64 + Send + Sync + 'static,
        value: f64,
    ) {
        self.constraints.push(Constraint {
            function: Arc::new(constraint_fn),
            constraint_type: ConstraintType::Eq,
            value,
        });
    }

    /// Add an inequality constraint
    pub fn add_inequality_constraint(
        &mut self,
        constraint_fn: impl Fn(&FxHashMap<String, f64>) -> f64 + Send + Sync + 'static,
        value: f64,
    ) {
        self.constraints.push(Constraint {
            function: Arc::new(constraint_fn),
            constraint_type: ConstraintType::Ineq,
            value,
        });
    }

    /// Optimize with constraints
    pub fn optimize(
        &mut self,
        circuit: &mut VariationalCircuit,
        cost_fn: impl Fn(&VariationalCircuit) -> QuantRS2Result<f64> + Send + Sync + 'static,
    ) -> QuantRS2Result<OptimizationResult> {
        if self.constraints.is_empty() {
            return self.base_optimizer.optimize(circuit, cost_fn);
        }

        // For constrained optimization, use penalty method
        let cost_fn = Arc::new(cost_fn);
        let constraints = self.constraints.clone();
        let penalty_weight = 1000.0;

        let penalized_cost = move |circuit: &VariationalCircuit| -> QuantRS2Result<f64> {
            let base_cost = cost_fn(circuit)?;
            let params = circuit.get_parameters();

            let mut penalty = 0.0;
            for constraint in &constraints {
                let constraint_value = (constraint.function)(&params);
                match constraint.constraint_type {
                    ConstraintType::Eq => {
                        penalty += penalty_weight * (constraint_value - constraint.value).powi(2);
                    }
                    ConstraintType::Ineq => {
                        if constraint_value > constraint.value {
                            penalty +=
                                penalty_weight * (constraint_value - constraint.value).powi(2);
                        }
                    }
                }
            }

            Ok(base_cost + penalty)
        };

        self.base_optimizer.optimize(circuit, penalized_cost)
    }
}

/// Hyperparameter optimization for variational circuits
pub struct HyperparameterOptimizer {
    /// Search space for hyperparameters
    search_space: FxHashMap<String, (f64, f64)>,
    /// Number of trials
    n_trials: usize,
    /// Optimization method for inner loop
    inner_method: OptimizationMethod,
}

impl HyperparameterOptimizer {
    /// Create a new hyperparameter optimizer
    pub fn new(n_trials: usize) -> Self {
        Self {
            search_space: FxHashMap::default(),
            n_trials,
            inner_method: OptimizationMethod::BFGS,
        }
    }

    /// Add a hyperparameter to search
    pub fn add_hyperparameter(&mut self, name: String, min_value: f64, max_value: f64) {
        self.search_space.insert(name, (min_value, max_value));
    }

    /// Optimize hyperparameters
    pub fn optimize(
        &self,
        circuit_builder: impl Fn(&FxHashMap<String, f64>) -> VariationalCircuit + Send + Sync,
        cost_fn: impl Fn(&VariationalCircuit) -> QuantRS2Result<f64> + Send + Sync + Clone + 'static,
    ) -> QuantRS2Result<HyperparameterResult> {
        use scirs2_core::random::prelude::*;

        let mut rng = StdRng::from_seed(thread_rng().gen());
        let mut best_hyperparams = FxHashMap::default();
        let mut best_loss = f64::INFINITY;
        let mut all_trials = Vec::new();

        for _trial in 0..self.n_trials {
            // Sample hyperparameters
            let mut hyperparams = FxHashMap::default();
            for (name, &(min_val, max_val)) in &self.search_space {
                let value = rng.gen_range(min_val..max_val);
                hyperparams.insert(name.clone(), value);
            }

            // Build circuit with hyperparameters
            let mut circuit = circuit_builder(&hyperparams);

            // Optimize circuit
            let config = OptimizationConfig {
                max_iterations: 50,
                ..Default::default()
            };

            let mut optimizer = VariationalQuantumOptimizer::new(self.inner_method.clone(), config);

            let result = optimizer.optimize(&mut circuit, cost_fn.clone())?;

            all_trials.push(HyperparameterTrial {
                hyperparameters: hyperparams.clone(),
                final_loss: result.final_loss,
                optimal_parameters: result.optimal_parameters,
            });

            if result.final_loss < best_loss {
                best_loss = result.final_loss;
                best_hyperparams = hyperparams;
            }
        }

        Ok(HyperparameterResult {
            best_hyperparameters: best_hyperparams,
            best_loss,
            all_trials,
        })
    }
}

/// Hyperparameter optimization result
#[derive(Debug, Clone)]
pub struct HyperparameterResult {
    /// Best hyperparameters found
    pub best_hyperparameters: FxHashMap<String, f64>,
    /// Best loss achieved
    pub best_loss: f64,
    /// All trials
    pub all_trials: Vec<HyperparameterTrial>,
}

/// Single hyperparameter trial
#[derive(Debug, Clone)]
pub struct HyperparameterTrial {
    /// Hyperparameters used
    pub hyperparameters: FxHashMap<String, f64>,
    /// Final loss achieved
    pub final_loss: f64,
    /// Optimal variational parameters
    pub optimal_parameters: FxHashMap<String, f64>,
}

// Clone implementation for VariationalCircuit
impl Clone for VariationalCircuit {
    fn clone(&self) -> Self {
        Self {
            gates: self.gates.clone(),
            param_map: self.param_map.clone(),
            num_qubits: self.num_qubits,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::qubit::QubitId;
    use crate::variational::VariationalGate;

    #[test]
    fn test_gradient_descent_optimizer() {
        let mut circuit = VariationalCircuit::new(1);
        circuit.add_gate(VariationalGate::rx(QubitId(0), "theta".to_string(), 0.0));

        let config = OptimizationConfig {
            max_iterations: 10,
            ..Default::default()
        };

        let mut optimizer = VariationalQuantumOptimizer::new(
            OptimizationMethod::GradientDescent { learning_rate: 0.1 },
            config,
        );

        // Simple cost function
        let cost_fn = |circuit: &VariationalCircuit| -> QuantRS2Result<f64> {
            let theta = circuit
                .get_parameters()
                .get("theta")
                .copied()
                .unwrap_or(0.0);
            Ok((theta - 1.0).powi(2))
        };

        let result = optimizer.optimize(&mut circuit, cost_fn).unwrap();

        assert!(result.converged || result.iterations == 10);
        assert!((result.optimal_parameters["theta"] - 1.0).abs() < 0.1);
    }

    #[test]
    fn test_adam_optimizer() {
        let mut circuit = VariationalCircuit::new(2);
        circuit.add_gate(VariationalGate::ry(QubitId(0), "alpha".to_string(), 0.5));
        circuit.add_gate(VariationalGate::rz(QubitId(1), "beta".to_string(), 0.5));

        let config = OptimizationConfig {
            max_iterations: 100,
            f_tol: 1e-6,
            g_tol: 1e-6,
            ..Default::default()
        };

        let mut optimizer = VariationalQuantumOptimizer::new(
            OptimizationMethod::Adam {
                learning_rate: 0.1,
                beta1: 0.9,
                beta2: 0.999,
                epsilon: 1e-8,
            },
            config,
        );

        // Cost function with multiple parameters
        let cost_fn = |circuit: &VariationalCircuit| -> QuantRS2Result<f64> {
            let params = circuit.get_parameters();
            let alpha = params.get("alpha").copied().unwrap_or(0.0);
            let beta = params.get("beta").copied().unwrap_or(0.0);
            Ok(alpha.powi(2) + beta.powi(2))
        };

        let result = optimizer.optimize(&mut circuit, cost_fn).unwrap();

        assert!(result.optimal_parameters["alpha"].abs() < 0.1);
        assert!(result.optimal_parameters["beta"].abs() < 0.1);
    }

    #[test]
    fn test_constrained_optimization() {
        let mut circuit = VariationalCircuit::new(1);
        circuit.add_gate(VariationalGate::rx(QubitId(0), "x".to_string(), 2.0));

        let base_optimizer =
            VariationalQuantumOptimizer::new(OptimizationMethod::BFGS, Default::default());

        let mut constrained_opt = ConstrainedVariationalOptimizer::new(base_optimizer);

        // Add constraint: x >= 1.0
        constrained_opt
            .add_inequality_constraint(|params| 1.0 - params.get("x").copied().unwrap_or(0.0), 0.0);

        // Minimize x^2
        let cost_fn = |circuit: &VariationalCircuit| -> QuantRS2Result<f64> {
            let x = circuit.get_parameters().get("x").copied().unwrap_or(0.0);
            Ok(x.powi(2))
        };

        let result = constrained_opt.optimize(&mut circuit, cost_fn).unwrap();

        let optimized_x = result.optimal_parameters["x"];
        assert!(optimized_x >= 1.0 - 1e-6);
        assert!(optimized_x <= 2.0 + 1e-6);
    }
}