quantrs2-anneal 0.2.1

//! Variational Quantum Annealing for advanced optimization
//!
//! This module implements variational quantum annealing (VQA) algorithms that combine
//! classical optimization with quantum annealing to solve complex optimization problems.
//! VQA uses parameterized quantum circuits and classical optimization to find optimal
//! solutions through iterative refinement.

use scirs2_core::random::prelude::*;
use scirs2_core::random::ChaCha8Rng;
use scirs2_core::random::{Rng, SeedableRng};
use std::collections::HashMap;
use std::time::{Duration, Instant};
use thiserror::Error;

use crate::ising::{IsingError, IsingModel};
use crate::simulator::{AnnealingParams, AnnealingSolution, QuantumAnnealingSimulator};

/// Errors that can occur in variational quantum annealing
#[derive(Error, Debug)]
pub enum VqaError {
    /// Ising model error
    #[error("Ising error: {0}")]
    IsingError(#[from] IsingError),

    /// Invalid variational parameters
    #[error("Invalid parameters: {0}")]
    InvalidParameters(String),

    /// Optimization failed
    #[error("Optimization failed: {0}")]
    OptimizationFailed(String),

    /// Circuit construction error
    #[error("Circuit error: {0}")]
    CircuitError(String),

    /// Convergence error
    #[error("Convergence error: {0}")]
    ConvergenceError(String),
}

/// Result type for VQA operations
pub type VqaResult<T> = Result<T, VqaError>;

/// Types of variational ansatz circuits
#[derive(Debug, Clone, PartialEq)]
pub enum AnsatzType {
    /// Hardware-efficient ansatz with parameterized rotations
    HardwareEfficient {
        depth: usize,
        entangling_gates: EntanglingGateType,
    },

    /// QAOA-inspired ansatz with problem and mixer terms
    QaoaInspired {
        layers: usize,
        mixer_type: MixerType,
    },

    /// Adiabatic-inspired ansatz with time evolution
    AdiabaticInspired {
        time_steps: usize,
        evolution_time: f64,
    },

    /// Custom ansatz with user-defined structure
    Custom { structure: Vec<QuantumGate> },
}

/// Types of entangling gates for hardware-efficient ansatz
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum EntanglingGateType {
    /// CNOT gates in nearest-neighbor topology
    CNot,
    /// Controlled-Z gates
    CZ,
    /// Ising-style ZZ interactions
    ZZ,
    /// XY gates for spin exchange
    XY,
}

/// Types of mixer operations
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum MixerType {
    /// X-rotation mixer (transverse field)
    XRotation,
    /// XY mixer for hard constraints
    XY,
    /// Multi-angle mixer
    MultiAngle,
}

/// Quantum gate representation for variational circuits
#[derive(Debug, Clone, PartialEq)]
pub enum QuantumGate {
    /// Single-qubit rotation around X-axis
    RX { qubit: usize, angle: ParameterRef },
    /// Single-qubit rotation around Y-axis
    RY { qubit: usize, angle: ParameterRef },
    /// Single-qubit rotation around Z-axis
    RZ { qubit: usize, angle: ParameterRef },
    /// Two-qubit CNOT gate
    CNOT { control: usize, target: usize },
    /// Two-qubit controlled-Z gate
    CZ { control: usize, target: usize },
    /// Parameterized ZZ interaction
    ZZ {
        qubit1: usize,
        qubit2: usize,
        angle: ParameterRef,
    },
}

/// Reference to a variational parameter
#[derive(Debug, Clone, PartialEq)]
pub struct ParameterRef {
    /// Parameter index
    pub index: usize,
    /// Optional scaling factor
    pub scale: f64,
}

impl ParameterRef {
    /// Create a new parameter reference
    #[must_use]
    pub const fn new(index: usize) -> Self {
        Self { index, scale: 1.0 }
    }

    /// Create a scaled parameter reference
    #[must_use]
    pub const fn scaled(index: usize, scale: f64) -> Self {
        Self { index, scale }
    }
}

/// Variational quantum annealing configuration
#[derive(Debug, Clone)]
pub struct VqaConfig {
    /// Ansatz type and structure
    pub ansatz: AnsatzType,

    /// Classical optimizer for parameters
    pub optimizer: ClassicalOptimizer,

    /// Maximum iterations for variational optimization
    pub max_iterations: usize,

    /// Convergence tolerance
    pub convergence_tolerance: f64,

    /// Number of quantum annealing shots per evaluation
    pub num_shots: usize,

    /// Base annealing parameters
    pub annealing_params: AnnealingParams,

    /// Parameter initialization range
    pub parameter_init_range: (f64, f64),

    /// Use gradient-based optimization
    pub use_gradients: bool,

    /// Finite difference step for gradient estimation
    pub gradient_step: f64,

    /// Random seed
    pub seed: Option<u64>,

    /// Maximum runtime
    pub max_runtime: Option<Duration>,

    /// Logging frequency
    pub log_frequency: usize,
}

impl Default for VqaConfig {
    fn default() -> Self {
        Self {
            ansatz: AnsatzType::HardwareEfficient {
                depth: 3,
                entangling_gates: EntanglingGateType::CNot,
            },
            optimizer: ClassicalOptimizer::Adam {
                learning_rate: 0.01,
                beta1: 0.9,
                beta2: 0.999,
                epsilon: 1e-8,
            },
            max_iterations: 100,
            convergence_tolerance: 1e-6,
            num_shots: 100,
            annealing_params: AnnealingParams::default(),
            parameter_init_range: (-0.5, 0.5),
            use_gradients: true,
            gradient_step: 0.01,
            seed: None,
            max_runtime: Some(Duration::from_secs(3600)),
            log_frequency: 10,
        }
    }
}

/// Classical optimizers for variational parameters
#[derive(Debug, Clone)]
pub enum ClassicalOptimizer {
    /// Gradient descent
    GradientDescent { learning_rate: f64 },

    /// Adam optimizer
    Adam {
        learning_rate: f64,
        beta1: f64,
        beta2: f64,
        epsilon: f64,
    },

    /// `RMSprop` optimizer
    RMSprop {
        learning_rate: f64,
        decay_rate: f64,
        epsilon: f64,
    },

    /// Nelder-Mead simplex
    NelderMead {
        initial_simplex_size: f64,
        alpha: f64,
        gamma: f64,
        rho: f64,
        sigma: f64,
    },

    /// BFGS quasi-Newton method
    BFGS {
        line_search_tolerance: f64,
        max_line_search_iterations: usize,
    },
}

/// Variational quantum annealing results
#[derive(Debug, Clone)]
pub struct VqaResults {
    /// Best solution found
    pub best_solution: Vec<i8>,

    /// Best energy achieved
    pub best_energy: f64,

    /// Optimal variational parameters
    pub optimal_parameters: Vec<f64>,

    /// Energy history over iterations
    pub energy_history: Vec<f64>,

    /// Parameter history over iterations
    pub parameter_history: Vec<Vec<f64>>,

    /// Gradient norms over iterations
    pub gradient_norms: Vec<f64>,

    /// Number of iterations completed
    pub iterations_completed: usize,

    /// Convergence achieved
    pub converged: bool,

    /// Total optimization time
    pub total_time: Duration,

    /// Statistics about the optimization
    pub statistics: VqaStatistics,
}

/// Statistics for VQA optimization
#[derive(Debug, Clone)]
pub struct VqaStatistics {
    /// Total function evaluations
    pub function_evaluations: usize,

    /// Total gradient evaluations
    pub gradient_evaluations: usize,

    /// Total quantum annealing time
    pub total_annealing_time: Duration,

    /// Average energy per iteration
    pub average_energy: f64,

    /// Energy variance over iterations
    pub energy_variance: f64,

    /// Parameter update statistics
    pub parameter_stats: ParameterStatistics,

    /// Classical optimizer performance
    pub optimizer_stats: OptimizerStatistics,
}

/// Statistics about parameter updates
#[derive(Debug, Clone)]
pub struct ParameterStatistics {
    /// Average parameter magnitude
    pub average_magnitude: f64,

    /// Parameter variance
    pub parameter_variance: f64,

    /// Number of parameter updates
    pub num_updates: usize,

    /// Largest parameter change per iteration
    pub max_parameter_change: Vec<f64>,
}

/// Statistics about classical optimizer performance
#[derive(Debug, Clone)]
pub struct OptimizerStatistics {
    /// Step acceptance rate
    pub step_acceptance_rate: f64,

    /// Average step size
    pub average_step_size: f64,

    /// Number of line search iterations (for applicable optimizers)
    pub line_search_iterations: usize,

    /// Optimizer-specific metrics
    pub optimizer_metrics: HashMap<String, f64>,
}

/// Variational quantum annealing optimizer
pub struct VariationalQuantumAnnealer {
    /// Configuration
    config: VqaConfig,

    /// Current variational parameters
    parameters: Vec<f64>,

    /// Classical optimizer state
    optimizer_state: OptimizerState,

    /// Random number generator
    rng: ChaCha8Rng,

    /// Optimization history
    history: OptimizationHistory,
}

/// Internal state for classical optimizers
#[derive(Debug)]
enum OptimizerState {
    GradientDescent {
        momentum: Option<Vec<f64>>,
    },

    Adam {
        m: Vec<f64>, // First moment estimate
        v: Vec<f64>, // Second moment estimate
        t: usize,    // Time step
    },

    RMSprop {
        s: Vec<f64>, // Moving average of squared gradients
    },

    NelderMead {
        simplex: Vec<Vec<f64>>,
        function_values: Vec<f64>,
    },

    BFGS {
        hessian_inverse: Vec<Vec<f64>>,
        previous_gradient: Option<Vec<f64>>,
        previous_parameters: Option<Vec<f64>>,
    },
}

/// Optimization history tracking
#[derive(Debug)]
struct OptimizationHistory {
    energies: Vec<f64>,
    parameters: Vec<Vec<f64>>,
    gradients: Vec<Vec<f64>>,
    function_evals: usize,
    gradient_evals: usize,
    start_time: Instant,
}

impl VariationalQuantumAnnealer {
    /// Create a new variational quantum annealer
    pub fn new(config: VqaConfig) -> VqaResult<Self> {
        let num_parameters = Self::count_parameters(&config.ansatz)?;

        let rng = match config.seed {
            Some(seed) => ChaCha8Rng::seed_from_u64(seed),
            None => ChaCha8Rng::seed_from_u64(thread_rng().random()),
        };

        let mut vqa = Self {
            config: config.clone(),
            parameters: vec![0.0; num_parameters],
            optimizer_state: Self::initialize_optimizer_state(&config.optimizer, num_parameters)?,
            rng,
            history: OptimizationHistory {
                energies: Vec::new(),
                parameters: Vec::new(),
                gradients: Vec::new(),
                function_evals: 0,
                gradient_evals: 0,
                start_time: Instant::now(),
            },
        };

        vqa.initialize_parameters()?;
        Ok(vqa)
    }

    /// Count the number of parameters in an ansatz.
    ///
    /// The hardware-efficient ansatz parameter count depends on the number of
    /// qubits, which is not known until a problem is supplied. At construction
    /// time we therefore size it with a single-qubit estimate; [`optimize`]
    /// recomputes the exact count for the actual problem (via
    /// [`Self::exact_parameter_count`]) and resizes the parameter vector before
    /// optimization begins, so no gates are ever silently dropped.
    fn count_parameters(ansatz: &AnsatzType) -> VqaResult<usize> {
        Self::exact_parameter_count(ansatz, 1)
    }

    /// Exact number of variational parameters for `ansatz` on `num_qubits`.
    ///
    /// This mirrors the gate-emission logic of the circuit builders so the
    /// parameter vector is always exactly the size the circuit consumes:
    /// - hardware-efficient: `depth · (2·n + e)` where `e` is the per-layer
    ///   count of parameterised entangling gates (`n-1` for `ZZ`/`XY`, `0` for
    ///   `CNot`/`CZ`),
    /// - QAOA-inspired: `2·layers` (one γ and one β per layer),
    /// - adiabatic-inspired: one angle per time step,
    /// - custom: one past the largest referenced parameter index.
    fn exact_parameter_count(ansatz: &AnsatzType, num_qubits: usize) -> VqaResult<usize> {
        match ansatz {
            AnsatzType::HardwareEfficient {
                depth,
                entangling_gates,
            } => {
                let rotations_per_layer = 2 * num_qubits;
                let entangling_params_per_layer = match entangling_gates {
                    EntanglingGateType::ZZ | EntanglingGateType::XY => num_qubits.saturating_sub(1),
                    EntanglingGateType::CNot | EntanglingGateType::CZ => 0,
                };
                Ok(depth * (rotations_per_layer + entangling_params_per_layer))
            }

            AnsatzType::QaoaInspired { layers, .. } => {
                // For QAOA: 2 parameters per layer (gamma and beta)
                Ok(layers * 2)
            }

            AnsatzType::AdiabaticInspired { time_steps, .. } => {
                // For adiabatic: one parameter per time step
                Ok(*time_steps)
            }

            AnsatzType::Custom { structure } => {
                // Count unique parameter references
                let mut max_param_index = 0;
                for gate in structure {
                    if let Some(param_ref) = Self::extract_parameter_ref(gate) {
                        max_param_index = max_param_index.max(param_ref.index);
                    }
                }
                Ok(max_param_index + 1)
            }
        }
    }

    /// Extract parameter reference from a gate
    const fn extract_parameter_ref(gate: &QuantumGate) -> Option<&ParameterRef> {
        match gate {
            QuantumGate::RX { angle, .. }
            | QuantumGate::RY { angle, .. }
            | QuantumGate::RZ { angle, .. }
            | QuantumGate::ZZ { angle, .. } => Some(angle),
            _ => None,
        }
    }

    /// Initialize optimizer state
    fn initialize_optimizer_state(
        optimizer: &ClassicalOptimizer,
        num_params: usize,
    ) -> VqaResult<OptimizerState> {
        match optimizer {
            ClassicalOptimizer::GradientDescent { .. } => {
                Ok(OptimizerState::GradientDescent { momentum: None })
            }

            ClassicalOptimizer::Adam { .. } => Ok(OptimizerState::Adam {
                m: vec![0.0; num_params],
                v: vec![0.0; num_params],
                t: 0,
            }),

            ClassicalOptimizer::RMSprop { .. } => Ok(OptimizerState::RMSprop {
                s: vec![0.0; num_params],
            }),

            ClassicalOptimizer::NelderMead {
                initial_simplex_size,
                ..
            } => {
                // Initialize simplex
                let mut simplex = vec![vec![0.0; num_params]; num_params + 1];
                for i in 0..num_params {
                    simplex[i + 1][i] = *initial_simplex_size;
                }

                Ok(OptimizerState::NelderMead {
                    simplex,
                    function_values: vec![f64::INFINITY; num_params + 1],
                })
            }

            ClassicalOptimizer::BFGS { .. } => {
                // Initialize identity matrix for Hessian inverse
                let mut hessian_inverse = vec![vec![0.0; num_params]; num_params];
                for i in 0..num_params {
                    hessian_inverse[i][i] = 1.0;
                }

                Ok(OptimizerState::BFGS {
                    hessian_inverse,
                    previous_gradient: None,
                    previous_parameters: None,
                })
            }
        }
    }

    /// Initialize variational parameters
    fn initialize_parameters(&mut self) -> VqaResult<()> {
        let (min, max) = self.config.parameter_init_range;

        for param in &mut self.parameters {
            *param = self.rng.random_range(min..max);
        }

        Ok(())
    }

    /// Optimize the variational quantum annealing problem
    pub fn optimize(&mut self, problem: &IsingModel) -> VqaResult<VqaResults> {
        println!("Starting variational quantum annealing optimization...");

        // The exact parameter count of (e.g.) the hardware-efficient ansatz is
        // problem-dependent. Recompute it for this problem and resize/reinit the
        // parameter vector and optimizer state if it differs from the
        // construction-time estimate, so the circuit is parameterised exactly.
        let exact_params = Self::exact_parameter_count(&self.config.ansatz, problem.num_qubits)?;
        if exact_params != self.parameters.len() {
            self.parameters = vec![0.0; exact_params];
            self.initialize_parameters()?;
            self.optimizer_state =
                Self::initialize_optimizer_state(&self.config.optimizer, exact_params)?;
        }

        self.history.start_time = Instant::now();
        let mut best_energy = f64::INFINITY;
        let mut best_solution = vec![0; problem.num_qubits];
        let mut best_parameters = self.parameters.clone();

        for iteration in 0..self.config.max_iterations {
            let iteration_start = Instant::now();

            // Check runtime limit
            if let Some(max_runtime) = self.config.max_runtime {
                if self.history.start_time.elapsed() > max_runtime {
                    println!("Maximum runtime exceeded");
                    break;
                }
            }

            // Evaluate current parameters
            let current_params = self.parameters.clone();
            let (energy, solution) = self.evaluate_objective(problem, &current_params)?;

            // Update best solution
            if energy < best_energy {
                best_energy = energy;
                best_solution = solution;
                best_parameters = self.parameters.clone();
            }

            // Record history
            self.history.energies.push(energy);
            self.history.parameters.push(self.parameters.clone());

            // Compute gradients if needed
            let gradients = if self.config.use_gradients {
                let grads = self.compute_gradients(problem)?;
                self.history.gradients.push(grads.clone());
                Some(grads)
            } else {
                None
            };

            // Update parameters using classical optimizer
            self.update_parameters(gradients.as_ref().map(std::vec::Vec::as_slice))?;

            // Logging
            if iteration % self.config.log_frequency == 0 {
                let grad_norm = gradients
                    .as_ref()
                    .map_or(0.0, |g| g.iter().map(|&x| x.powi(2)).sum::<f64>().sqrt());

                println!(
                    "Iteration {}: Energy = {:.6}, Gradient norm = {:.6}, Time = {:.2?}",
                    iteration,
                    energy,
                    grad_norm,
                    iteration_start.elapsed()
                );
            }

            // Check convergence
            if self.check_convergence()? {
                println!("Converged at iteration {iteration}");
                break;
            }
        }

        let total_time = self.history.start_time.elapsed();

        // Calculate statistics
        let statistics = self.calculate_statistics();

        Ok(VqaResults {
            best_solution,
            best_energy,
            optimal_parameters: best_parameters,
            energy_history: self.history.energies.clone(),
            parameter_history: self.history.parameters.clone(),
            gradient_norms: self
                .history
                .gradients
                .iter()
                .map(|g| g.iter().map(|&x| x.powi(2)).sum::<f64>().sqrt())
                .collect(),
            iterations_completed: self.history.energies.len(),
            converged: self.check_convergence()?,
            total_time,
            statistics,
        })
    }

    /// Evaluate the objective function for given parameters
    fn evaluate_objective(
        &mut self,
        problem: &IsingModel,
        parameters: &[f64],
    ) -> VqaResult<(f64, Vec<i8>)> {
        self.history.function_evals += 1;

        // Construct the parameterized quantum circuit
        let circuit = self.build_quantum_circuit(problem, parameters)?;

        // Execute quantum annealing with the parameterized problem
        let modified_problem = self.apply_circuit_to_problem(problem, &circuit)?;

        // Perform quantum annealing
        let mut simulator = QuantumAnnealingSimulator::new(self.config.annealing_params.clone())
            .map_err(|e| VqaError::OptimizationFailed(e.to_string()))?;

        let mut best_energy = f64::INFINITY;
        let mut best_solution = vec![0; problem.num_qubits];

        // Multiple shots for statistical averaging
        for _ in 0..self.config.num_shots {
            let result = simulator
                .solve(&modified_problem)
                .map_err(|e| VqaError::OptimizationFailed(e.to_string()))?;

            if result.best_energy < best_energy {
                best_energy = result.best_energy;
                best_solution = result.best_spins;
            }
        }

        Ok((best_energy, best_solution))
    }

    /// Build the quantum circuit for current parameters
    fn build_quantum_circuit(
        &self,
        problem: &IsingModel,
        parameters: &[f64],
    ) -> VqaResult<QuantumCircuit> {
        let num_qubits = problem.num_qubits;
        let mut circuit = QuantumCircuit::new(num_qubits);

        match &self.config.ansatz {
            AnsatzType::HardwareEfficient {
                depth,
                entangling_gates,
            } => {
                self.build_hardware_efficient_circuit(
                    &mut circuit,
                    *depth,
                    entangling_gates,
                    parameters,
                )?;
            }

            AnsatzType::QaoaInspired { layers, mixer_type } => {
                self.build_qaoa_inspired_circuit(
                    &mut circuit,
                    problem,
                    *layers,
                    mixer_type,
                    parameters,
                )?;
            }

            AnsatzType::AdiabaticInspired {
                time_steps,
                evolution_time,
            } => {
                self.build_adiabatic_inspired_circuit(
                    &mut circuit,
                    problem,
                    *time_steps,
                    *evolution_time,
                    parameters,
                )?;
            }

            AnsatzType::Custom { structure } => {
                self.build_custom_circuit(&mut circuit, structure, parameters)?;
            }
        }

        Ok(circuit)
    }

    /// Build hardware-efficient ansatz circuit
    fn build_hardware_efficient_circuit(
        &self,
        circuit: &mut QuantumCircuit,
        depth: usize,
        entangling_gates: &EntanglingGateType,
        parameters: &[f64],
    ) -> VqaResult<()> {
        let num_qubits = circuit.num_qubits;
        let mut param_idx = 0;

        for layer in 0..depth {
            // Single-qubit rotations. The resolved angle is stored in the
            // `scale` field (consistent with the QAOA/adiabatic builders) so the
            // mean-field transform can read the actual rotation angle.
            for qubit in 0..num_qubits {
                if param_idx < parameters.len() {
                    circuit.add_gate(QuantumGate::RY {
                        qubit,
                        angle: ParameterRef::scaled(param_idx, parameters[param_idx]),
                    });
                    param_idx += 1;
                }

                if param_idx < parameters.len() {
                    circuit.add_gate(QuantumGate::RZ {
                        qubit,
                        angle: ParameterRef::scaled(param_idx, parameters[param_idx]),
                    });
                    param_idx += 1;
                }
            }

            // Entangling gates
            match entangling_gates {
                EntanglingGateType::CNot => {
                    for qubit in 0..num_qubits - 1 {
                        circuit.add_gate(QuantumGate::CNOT {
                            control: qubit,
                            target: qubit + 1,
                        });
                    }
                }

                EntanglingGateType::CZ => {
                    for qubit in 0..num_qubits - 1 {
                        circuit.add_gate(QuantumGate::CZ {
                            control: qubit,
                            target: qubit + 1,
                        });
                    }
                }

                EntanglingGateType::ZZ => {
                    for qubit in 0..num_qubits - 1 {
                        if param_idx < parameters.len() {
                            circuit.add_gate(QuantumGate::ZZ {
                                qubit1: qubit,
                                qubit2: qubit + 1,
                                angle: ParameterRef::scaled(param_idx, parameters[param_idx]),
                            });
                            param_idx += 1;
                        }
                    }
                }

                EntanglingGateType::XY => {
                    // Implement XY gates as combination of other gates
                    for qubit in 0..num_qubits - 1 {
                        if param_idx < parameters.len() {
                            // Simplified XY implementation
                            circuit.add_gate(QuantumGate::CNOT {
                                control: qubit,
                                target: qubit + 1,
                            });
                            param_idx += 1;
                        }
                    }
                }
            }
        }

        Ok(())
    }

    /// Build QAOA-inspired circuit
    fn build_qaoa_inspired_circuit(
        &self,
        circuit: &mut QuantumCircuit,
        problem: &IsingModel,
        layers: usize,
        mixer_type: &MixerType,
        parameters: &[f64],
    ) -> VqaResult<()> {
        let num_qubits = circuit.num_qubits;

        for layer in 0..layers {
            let gamma_idx = layer * 2;
            let beta_idx = layer * 2 + 1;

            if gamma_idx >= parameters.len() || beta_idx >= parameters.len() {
                break;
            }

            let gamma = parameters[gamma_idx];
            let beta = parameters[beta_idx];

            // Problem Hamiltonian layer (ZZ interactions)
            for i in 0..num_qubits {
                for j in (i + 1)..num_qubits {
                    if let Ok(coupling) = problem.get_coupling(i, j) {
                        if coupling != 0.0 {
                            circuit.add_gate(QuantumGate::ZZ {
                                qubit1: i,
                                qubit2: j,
                                angle: ParameterRef::scaled(gamma_idx, gamma * coupling),
                            });
                        }
                    }
                }

                // Bias terms
                if let Ok(bias) = problem.get_bias(i) {
                    if bias != 0.0 {
                        circuit.add_gate(QuantumGate::RZ {
                            qubit: i,
                            angle: ParameterRef::scaled(gamma_idx, gamma * bias),
                        });
                    }
                }
            }

            // Mixer layer
            match mixer_type {
                MixerType::XRotation => {
                    for qubit in 0..num_qubits {
                        circuit.add_gate(QuantumGate::RX {
                            qubit,
                            angle: ParameterRef::scaled(beta_idx, beta),
                        });
                    }
                }

                MixerType::XY => {
                    // XY mixer for hard constraints
                    for qubit in 0..num_qubits - 1 {
                        circuit.add_gate(QuantumGate::CNOT {
                            control: qubit,
                            target: qubit + 1,
                        });
                    }
                }

                MixerType::MultiAngle => {
                    // Multi-angle mixer
                    for qubit in 0..num_qubits {
                        circuit.add_gate(QuantumGate::RX {
                            qubit,
                            angle: ParameterRef::scaled(beta_idx, beta),
                        });
                        circuit.add_gate(QuantumGate::RY {
                            qubit,
                            angle: ParameterRef::scaled(beta_idx, beta * 0.5),
                        });
                    }
                }
            }
        }

        Ok(())
    }

    /// Build adiabatic-inspired circuit
    fn build_adiabatic_inspired_circuit(
        &self,
        circuit: &mut QuantumCircuit,
        problem: &IsingModel,
        time_steps: usize,
        evolution_time: f64,
        parameters: &[f64],
    ) -> VqaResult<()> {
        let num_qubits = circuit.num_qubits;
        let dt = evolution_time / time_steps as f64;

        for step in 0..time_steps {
            if step >= parameters.len() {
                break;
            }

            let s = parameters[step]; // Annealing parameter s(t)

            // Transverse field (initial Hamiltonian)
            for qubit in 0..num_qubits {
                circuit.add_gate(QuantumGate::RX {
                    qubit,
                    angle: ParameterRef::scaled(step, -2.0 * (1.0 - s) * dt),
                });
            }

            // Problem Hamiltonian
            for i in 0..num_qubits {
                for j in (i + 1)..num_qubits {
                    if let Ok(coupling) = problem.get_coupling(i, j) {
                        if coupling != 0.0 {
                            circuit.add_gate(QuantumGate::ZZ {
                                qubit1: i,
                                qubit2: j,
                                angle: ParameterRef::scaled(step, -s * coupling * dt),
                            });
                        }
                    }
                }
            }
        }

        Ok(())
    }

    /// Build custom circuit from structure
    fn build_custom_circuit(
        &self,
        circuit: &mut QuantumCircuit,
        structure: &[QuantumGate],
        parameters: &[f64],
    ) -> VqaResult<()> {
        for gate in structure {
            // Create gate with current parameter values
            let parameterized_gate = match gate {
                QuantumGate::RX { qubit, angle } => {
                    let param_value = if angle.index < parameters.len() {
                        parameters[angle.index] * angle.scale
                    } else {
                        0.0
                    };
                    QuantumGate::RX {
                        qubit: *qubit,
                        angle: ParameterRef::scaled(angle.index, param_value),
                    }
                }

                QuantumGate::RY { qubit, angle } => {
                    let param_value = if angle.index < parameters.len() {
                        parameters[angle.index] * angle.scale
                    } else {
                        0.0
                    };
                    QuantumGate::RY {
                        qubit: *qubit,
                        angle: ParameterRef::scaled(angle.index, param_value),
                    }
                }

                QuantumGate::RZ { qubit, angle } => {
                    let param_value = if angle.index < parameters.len() {
                        parameters[angle.index] * angle.scale
                    } else {
                        0.0
                    };
                    QuantumGate::RZ {
                        qubit: *qubit,
                        angle: ParameterRef::scaled(angle.index, param_value),
                    }
                }

                QuantumGate::ZZ {
                    qubit1,
                    qubit2,
                    angle,
                } => {
                    let param_value = if angle.index < parameters.len() {
                        parameters[angle.index] * angle.scale
                    } else {
                        0.0
                    };
                    QuantumGate::ZZ {
                        qubit1: *qubit1,
                        qubit2: *qubit2,
                        angle: ParameterRef::scaled(angle.index, param_value),
                    }
                }

                // Non-parameterized gates
                _ => gate.clone(),
            };

            circuit.add_gate(parameterized_gate);
        }

        Ok(())
    }

    /// Apply the parameterized variational circuit to the problem Hamiltonian.
    ///
    /// This produces the **mean-field effective Ising model** induced by the
    /// ansatz: each qubit's reduced single-qubit state is evolved on the Bloch
    /// sphere under the circuit's single-qubit rotations, giving a per-qubit
    /// longitudinal magnetization `m_i = ⟨Z_i⟩ ∈ [-1, 1]`. The problem
    /// Hamiltonian is then reweighted by these expectations,
    ///
    /// ```text
    /// h_i  →  h_i · m_i
    /// J_ij →  J_ij · m_i · m_j
    /// ```
    ///
    /// which is the standard mean-field (product-state) energy
    /// `⟨ψ(θ)| H | ψ(θ)⟩` for an Ising Hamiltonian. Entangling gates (CNOT/CZ)
    /// generate correlations that a product state cannot represent; we account
    /// for their depolarizing effect on the single-qubit marginals by damping
    /// the magnetization of the involved qubits, the standard mean-field
    /// correction. The transform is the identity when every rotation angle is
    /// zero (then `m_i = 1`), and the parameters genuinely reshape the energy
    /// landscape, so the variational gradient is non-trivial.
    fn apply_circuit_to_problem(
        &self,
        problem: &IsingModel,
        circuit: &QuantumCircuit,
    ) -> VqaResult<IsingModel> {
        let magnetizations = Self::compute_magnetizations(problem.num_qubits, circuit);

        let mut modified = IsingModel::new(problem.num_qubits);

        // Reweight biases by single-qubit magnetization.
        for (qubit, bias) in problem.biases() {
            let new_bias = bias * magnetizations[qubit];
            if new_bias != 0.0 {
                modified.set_bias(qubit, new_bias)?;
            }
        }

        // Reweight couplings by the product of the two magnetizations.
        for coupling in problem.couplings() {
            let new_strength =
                coupling.strength * magnetizations[coupling.i] * magnetizations[coupling.j];
            if new_strength != 0.0 {
                modified.set_coupling(coupling.i, coupling.j, new_strength)?;
            }
        }

        Ok(modified)
    }

    /// Evolve each qubit's Bloch vector under the circuit's single-qubit
    /// rotations and return the longitudinal magnetization `⟨Z_i⟩` per qubit.
    ///
    /// Each qubit starts polarized along `+Z` (Bloch vector `(0, 0, 1)`).
    /// Rotations act as the usual SO(3) Bloch-sphere rotations; CNOT/CZ
    /// entangling gates depolarize the single-qubit marginal of the involved
    /// qubits (mean-field correction), shrinking `|⟨Z⟩|` towards 0.
    fn compute_magnetizations(num_qubits: usize, circuit: &QuantumCircuit) -> Vec<f64> {
        // Bloch vectors, one per qubit, initialized along +Z.
        let mut bloch: Vec<[f64; 3]> = vec![[0.0, 0.0, 1.0]; num_qubits];

        // Damping factor applied to a qubit's transverse/longitudinal components
        // when it participates in an entangling gate (loss of single-qubit purity
        // to two-qubit correlations). 1/√2 reflects a maximally-entangling step.
        const ENTANGLE_DAMP: f64 = std::f64::consts::FRAC_1_SQRT_2;

        for gate in &circuit.gates {
            match gate {
                QuantumGate::RX { qubit, angle } => {
                    if *qubit < num_qubits {
                        bloch[*qubit] = rotate_x(bloch[*qubit], angle.scale);
                    }
                }
                QuantumGate::RY { qubit, angle } => {
                    if *qubit < num_qubits {
                        bloch[*qubit] = rotate_y(bloch[*qubit], angle.scale);
                    }
                }
                QuantumGate::RZ { qubit, angle } => {
                    if *qubit < num_qubits {
                        bloch[*qubit] = rotate_z(bloch[*qubit], angle.scale);
                    }
                }
                QuantumGate::ZZ {
                    qubit1,
                    qubit2,
                    angle,
                } => {
                    // A ZZ(θ) interaction correlates the two qubits; in mean field
                    // it damps each qubit's transverse magnetization by cos(θ) and
                    // shrinks the longitudinal marginals proportionally.
                    let damp = angle.scale.cos().abs();
                    for &q in &[*qubit1, *qubit2] {
                        if q < num_qubits {
                            bloch[q][0] *= damp;
                            bloch[q][1] *= damp;
                            bloch[q][2] *= damp;
                        }
                    }
                }
                QuantumGate::CNOT { control, target } | QuantumGate::CZ { control, target } => {
                    for &q in &[*control, *target] {
                        if q < num_qubits {
                            bloch[q][0] *= ENTANGLE_DAMP;
                            bloch[q][1] *= ENTANGLE_DAMP;
                            bloch[q][2] *= ENTANGLE_DAMP;
                        }
                    }
                }
            }
        }

        // The longitudinal magnetization is the Z-component of the Bloch vector.
        bloch.iter().map(|v| v[2]).collect()
    }

    /// Compute gradients using finite differences
    fn compute_gradients(&mut self, problem: &IsingModel) -> VqaResult<Vec<f64>> {
        self.history.gradient_evals += 1;

        let mut gradients = vec![0.0; self.parameters.len()];
        let step = self.config.gradient_step;

        for i in 0..self.parameters.len() {
            // Create modified parameter vectors
            let mut params_plus = self.parameters.clone();
            let mut params_minus = self.parameters.clone();

            params_plus[i] += step;
            params_minus[i] -= step;

            let (energy_plus, _) = self.evaluate_objective(problem, &params_plus)?;
            let (energy_minus, _) = self.evaluate_objective(problem, &params_minus)?;

            // Compute gradient
            gradients[i] = (energy_plus - energy_minus) / (2.0 * step);
        }

        Ok(gradients)
    }

    /// Update parameters using classical optimizer
    fn update_parameters(&mut self, gradients: Option<&[f64]>) -> VqaResult<()> {
        match (&mut self.optimizer_state, &self.config.optimizer) {
            (
                OptimizerState::Adam { m, v, t },
                ClassicalOptimizer::Adam {
                    learning_rate,
                    beta1,
                    beta2,
                    epsilon,
                },
            ) => {
                if let Some(grads) = gradients {
                    *t += 1;

                    for i in 0..self.parameters.len() {
                        // Update biased first moment estimate
                        m[i] = (1.0 - beta1).mul_add(grads[i], beta1 * m[i]);

                        // Update biased second moment estimate
                        v[i] = (1.0 - beta2).mul_add(grads[i].powi(2), beta2 * v[i]);

                        // Compute bias-corrected estimates
                        let m_hat = m[i] / (1.0 - beta1.powi(*t as i32));
                        let v_hat = v[i] / (1.0 - beta2.powi(*t as i32));

                        // Update parameter
                        self.parameters[i] -= learning_rate * m_hat / (v_hat.sqrt() + epsilon);
                    }
                }
            }

            (
                OptimizerState::GradientDescent { .. },
                ClassicalOptimizer::GradientDescent { learning_rate },
            ) => {
                if let Some(grads) = gradients {
                    for i in 0..self.parameters.len() {
                        self.parameters[i] -= learning_rate * grads[i];
                    }
                }
            }

            _ => {
                // Implement other optimizers as needed
                return Err(VqaError::OptimizationFailed(
                    "Optimizer not implemented".to_string(),
                ));
            }
        }

        Ok(())
    }

    /// Check convergence criteria
    fn check_convergence(&self) -> VqaResult<bool> {
        if self.history.energies.len() < 2 {
            return Ok(false);
        }

        let recent_energies =
            &self.history.energies[self.history.energies.len().saturating_sub(5)..];
        let energy_range = recent_energies
            .iter()
            .copied()
            .fold(f64::NEG_INFINITY, f64::max)
            - recent_energies
                .iter()
                .copied()
                .fold(f64::INFINITY, f64::min);

        Ok(energy_range < self.config.convergence_tolerance)
    }

    /// Calculate optimization statistics
    fn calculate_statistics(&self) -> VqaStatistics {
        let average_energy = if self.history.energies.is_empty() {
            0.0
        } else {
            self.history.energies.iter().sum::<f64>() / self.history.energies.len() as f64
        };

        let energy_variance = if self.history.energies.len() > 1 {
            let mean = average_energy;
            self.history
                .energies
                .iter()
                .map(|&e| (e - mean).powi(2))
                .sum::<f64>()
                / (self.history.energies.len() - 1) as f64
        } else {
            0.0
        };

        // Calculate parameter statistics
        let parameter_stats = self.calculate_parameter_statistics();

        // Wall-clock time spent in optimization (dominated by the annealing
        // shots). Computed from the run's start instant rather than reported as
        // a fixed zero.
        let total_annealing_time = self.history.start_time.elapsed();

        // Step acceptance rate: fraction of optimization steps that strictly
        // lowered the objective, measured from the recorded energy trajectory.
        let step_acceptance_rate = if self.history.energies.len() > 1 {
            let improving = self
                .history
                .energies
                .windows(2)
                .filter(|w| w[1] < w[0])
                .count();
            improving as f64 / (self.history.energies.len() - 1) as f64
        } else {
            0.0
        };

        // Average step size: mean Euclidean distance between consecutive
        // parameter vectors actually visited during optimization.
        let average_step_size = if self.history.parameters.len() > 1 {
            let total: f64 = self
                .history
                .parameters
                .windows(2)
                .map(|w| {
                    w[0].iter()
                        .zip(w[1].iter())
                        .map(|(a, b)| (a - b).powi(2))
                        .sum::<f64>()
                        .sqrt()
                })
                .sum();
            total / (self.history.parameters.len() - 1) as f64
        } else {
            0.0
        };

        VqaStatistics {
            function_evaluations: self.history.function_evals,
            gradient_evaluations: self.history.gradient_evals,
            total_annealing_time,
            average_energy,
            energy_variance,
            parameter_stats,
            optimizer_stats: OptimizerStatistics {
                step_acceptance_rate,
                average_step_size,
                line_search_iterations: 0,
                optimizer_metrics: HashMap::new(),
            },
        }
    }

    /// Calculate parameter statistics
    fn calculate_parameter_statistics(&self) -> ParameterStatistics {
        let average_magnitude = if self.parameters.is_empty() {
            0.0
        } else {
            self.parameters.iter().map(|&p| p.abs()).sum::<f64>() / self.parameters.len() as f64
        };

        let parameter_variance = if self.parameters.len() > 1 {
            let mean = self.parameters.iter().sum::<f64>() / self.parameters.len() as f64;
            self.parameters
                .iter()
                .map(|&p| (p - mean).powi(2))
                .sum::<f64>()
                / (self.parameters.len() - 1) as f64
        } else {
            0.0
        };

        // Per-parameter maximum absolute change observed across consecutive
        // recorded parameter vectors during optimization.
        let max_parameter_change = if self.history.parameters.len() > 1 {
            let num_params = self.parameters.len();
            let mut max_change = vec![0.0_f64; num_params];
            for window in self.history.parameters.windows(2) {
                for (idx, slot) in max_change.iter_mut().enumerate() {
                    if let (Some(&prev), Some(&curr)) = (window[0].get(idx), window[1].get(idx)) {
                        *slot = slot.max((curr - prev).abs());
                    }
                }
            }
            max_change
        } else {
            Vec::new()
        };

        ParameterStatistics {
            average_magnitude,
            parameter_variance,
            num_updates: self.history.parameters.len(),
            max_parameter_change,
        }
    }
}

/// Quantum circuit representation
#[derive(Debug, Clone)]
pub struct QuantumCircuit {
    /// Number of qubits
    pub num_qubits: usize,

    /// Sequence of quantum gates
    pub gates: Vec<QuantumGate>,
}

impl QuantumCircuit {
    /// Create a new quantum circuit
    #[must_use]
    pub const fn new(num_qubits: usize) -> Self {
        Self {
            num_qubits,
            gates: Vec::new(),
        }
    }

    /// Add a gate to the circuit
    pub fn add_gate(&mut self, gate: QuantumGate) {
        self.gates.push(gate);
    }

    /// Get the depth of the circuit
    #[must_use]
    pub fn depth(&self) -> usize {
        // Simplified depth calculation
        self.gates.len()
    }
}

/// Rotate a Bloch vector by `theta` about the X axis (right-handed SO(3)).
fn rotate_x(v: [f64; 3], theta: f64) -> [f64; 3] {
    let (s, c) = theta.sin_cos();
    [
        v[0],
        c.mul_add(v[1], -(s * v[2])),
        s.mul_add(v[1], c * v[2]),
    ]
}

/// Rotate a Bloch vector by `theta` about the Y axis (right-handed SO(3)).
fn rotate_y(v: [f64; 3], theta: f64) -> [f64; 3] {
    let (s, c) = theta.sin_cos();
    [
        s.mul_add(v[2], c * v[0]),
        v[1],
        c.mul_add(v[2], -(s * v[0])),
    ]
}

/// Rotate a Bloch vector by `theta` about the Z axis (right-handed SO(3)).
fn rotate_z(v: [f64; 3], theta: f64) -> [f64; 3] {
    let (s, c) = theta.sin_cos();
    [
        c.mul_add(v[0], -(s * v[1])),
        s.mul_add(v[0], c * v[1]),
        v[2],
    ]
}

/// Helper functions for creating common VQA configurations

/// Create a QAOA-style VQA configuration
#[must_use]
pub fn create_qaoa_vqa_config(layers: usize, max_iterations: usize) -> VqaConfig {
    VqaConfig {
        ansatz: AnsatzType::QaoaInspired {
            layers,
            mixer_type: MixerType::XRotation,
        },
        max_iterations,
        ..Default::default()
    }
}

/// Create a hardware-efficient VQA configuration
#[must_use]
pub fn create_hardware_efficient_vqa_config(depth: usize, max_iterations: usize) -> VqaConfig {
    VqaConfig {
        ansatz: AnsatzType::HardwareEfficient {
            depth,
            entangling_gates: EntanglingGateType::CNot,
        },
        max_iterations,
        ..Default::default()
    }
}

/// Create an adiabatic-inspired VQA configuration
#[must_use]
pub fn create_adiabatic_vqa_config(
    time_steps: usize,
    evolution_time: f64,
    max_iterations: usize,
) -> VqaConfig {
    VqaConfig {
        ansatz: AnsatzType::AdiabaticInspired {
            time_steps,
            evolution_time,
        },
        max_iterations,
        ..Default::default()
    }
}

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

    #[test]
    fn test_vqa_config_creation() {
        let config = create_qaoa_vqa_config(3, 50);

        match config.ansatz {
            AnsatzType::QaoaInspired { layers, .. } => {
                assert_eq!(layers, 3);
            }
            _ => panic!("Expected QAOA ansatz"),
        }

        assert_eq!(config.max_iterations, 50);
    }

    #[test]
    fn test_parameter_ref() {
        let param_ref = ParameterRef::new(5);
        assert_eq!(param_ref.index, 5);
        assert_eq!(param_ref.scale, 1.0);

        let scaled_ref = ParameterRef::scaled(3, 2.5);
        assert_eq!(scaled_ref.index, 3);
        assert_eq!(scaled_ref.scale, 2.5);
    }

    #[test]
    fn test_quantum_circuit() {
        let mut circuit = QuantumCircuit::new(3);
        assert_eq!(circuit.num_qubits, 3);
        assert_eq!(circuit.gates.len(), 0);

        circuit.add_gate(QuantumGate::RX {
            qubit: 0,
            angle: ParameterRef::new(0),
        });

        assert_eq!(circuit.gates.len(), 1);
        assert_eq!(circuit.depth(), 1);
    }

    #[test]
    fn test_parameter_counting() {
        let ansatz = AnsatzType::QaoaInspired {
            layers: 5,
            mixer_type: MixerType::XRotation,
        };

        let count = VariationalQuantumAnnealer::count_parameters(&ansatz)
            .expect("parameter counting should succeed");
        assert_eq!(count, 10); // 2 parameters per layer
    }
}