quantrs2_core/

qaoa.rs

//! Quantum Approximate Optimization Algorithm (QAOA) implementation
//!
//! QAOA is a hybrid quantum-classical algorithm for solving combinatorial optimization problems.
//! This implementation leverages SciRS2 for enhanced performance.

use crate::complex_ext::QuantumComplexExt;
use crate::simd_ops;
use scirs2_core::ndarray::Array2;
use scirs2_core::Complex64;
use std::f64::consts::PI;

/// QAOA circuit parameters
#[derive(Debug, Clone)]
pub struct QAOAParams {
    /// Number of QAOA layers (p)
    pub layers: usize,
    /// Mixer angles (beta parameters)
    pub beta: Vec<f64>,
    /// Cost angles (gamma parameters)
    pub gamma: Vec<f64>,
}

impl QAOAParams {
    /// Create new QAOA parameters with given number of layers
    pub fn new(layers: usize) -> Self {
        Self {
            layers,
            beta: vec![0.0; layers],
            gamma: vec![0.0; layers],
        }
    }

    /// Initialize with random parameters
    pub fn random(layers: usize) -> Self {
        let mut beta = Vec::with_capacity(layers);
        let mut gamma = Vec::with_capacity(layers);

        for i in 0..layers {
            // Simple pseudo-random for reproducibility
            let rand_val = f64::midpoint((i as f64).mul_add(1.234, 5.678).sin(), 1.0);
            beta.push(rand_val * PI);
            gamma.push(rand_val * 2.0 * PI);
        }

        Self {
            layers,
            beta,
            gamma,
        }
    }

    /// Update parameters (for optimization)
    pub fn update(&mut self, new_beta: Vec<f64>, new_gamma: Vec<f64>) {
        assert_eq!(new_beta.len(), self.layers);
        assert_eq!(new_gamma.len(), self.layers);
        self.beta = new_beta;
        self.gamma = new_gamma;
    }
}

/// QAOA cost Hamiltonian types
#[derive(Clone)]
pub enum CostHamiltonian {
    /// Max-Cut problem: H_C = Σ_{<i,j>} (1 - Z_i Z_j) / 2
    MaxCut(Vec<(usize, usize)>),
    /// Weighted Max-Cut: H_C = Σ_{<i,j>} w_{ij} (1 - Z_i Z_j) / 2
    WeightedMaxCut(Vec<(usize, usize, f64)>),
    /// General Ising model: H_C = Σ_i h_i Z_i + Σ_{<i,j>} J_{ij} Z_i Z_j
    Ising {
        h: Vec<f64>,
        j: Vec<((usize, usize), f64)>,
    },
}

impl std::fmt::Debug for CostHamiltonian {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::MaxCut(edges) => f.debug_tuple("MaxCut").field(edges).finish(),
            Self::WeightedMaxCut(edges) => f.debug_tuple("WeightedMaxCut").field(edges).finish(),
            Self::Ising { h, j } => f.debug_struct("Ising").field("h", h).field("j", j).finish(),
        }
    }
}

/// QAOA mixer Hamiltonian types
#[derive(Debug, Clone)]
pub enum MixerHamiltonian {
    /// Standard X-mixer: H_B = Σ_i X_i
    TransverseField,
    /// Custom mixer
    Custom(Vec<Array2<Complex64>>),
}

/// QAOA circuit builder
#[derive(Debug, Clone)]
pub struct QAOACircuit {
    pub num_qubits: usize,
    cost_hamiltonian: CostHamiltonian,
    mixer_hamiltonian: MixerHamiltonian,
    params: QAOAParams,
}

impl QAOACircuit {
    /// Create a new QAOA circuit
    pub const fn new(
        num_qubits: usize,
        cost_hamiltonian: CostHamiltonian,
        mixer_hamiltonian: MixerHamiltonian,
        params: QAOAParams,
    ) -> Self {
        Self {
            num_qubits,
            cost_hamiltonian,
            mixer_hamiltonian,
            params,
        }
    }

    /// Apply the initial state preparation (usually |+>^n)
    pub fn prepare_initial_state(&self, state: &mut [Complex64]) {
        let n = state.len();
        let amplitude = Complex64::new(1.0 / (n as f64).sqrt(), 0.0);
        state.fill(amplitude);
    }

    /// Apply the cost Hamiltonian evolution exp(-i γ H_C)
    pub fn apply_cost_evolution(&self, state: &mut [Complex64], gamma: f64) {
        match &self.cost_hamiltonian {
            CostHamiltonian::MaxCut(edges) => {
                for &(i, j) in edges {
                    self.apply_zz_rotation(state, i, j, gamma);
                }
            }
            CostHamiltonian::WeightedMaxCut(weighted_edges) => {
                for &(i, j, weight) in weighted_edges {
                    self.apply_zz_rotation(state, i, j, gamma * weight);
                }
            }
            CostHamiltonian::Ising { h, j } => {
                // Apply single-qubit Z rotations
                for (i, &h_i) in h.iter().enumerate() {
                    if h_i.abs() > 1e-10 {
                        self.apply_z_rotation(state, i, gamma * h_i);
                    }
                }
                // Apply two-qubit ZZ rotations
                for &((i, j), j_ij) in j {
                    if j_ij.abs() > 1e-10 {
                        self.apply_zz_rotation(state, i, j, gamma * j_ij);
                    }
                }
            }
        }
    }

    /// Apply the mixer Hamiltonian evolution exp(-i β H_B)
    pub fn apply_mixer_evolution(&self, state: &mut [Complex64], beta: f64) {
        match &self.mixer_hamiltonian {
            MixerHamiltonian::TransverseField => {
                // Apply X rotation to each qubit
                for i in 0..self.num_qubits {
                    self.apply_x_rotation(state, i, beta);
                }
            }
            MixerHamiltonian::Custom(matrices) => {
                // Apply custom mixer evolution
                // For each term in the mixer Hamiltonian, apply the corresponding unitary evolution
                for matrix in matrices {
                    self.apply_custom_mixer_term(state, matrix, beta);
                }
            }
        }
    }

    /// Apply a custom mixer term exp(-i β H_term) to the state
    ///
    /// This implementation uses Trotterization for general Hamiltonians.
    /// For optimal performance, mixer terms should be decomposed into
    /// products of Pauli operators before being passed to this function.
    fn apply_custom_mixer_term(
        &self,
        state: &mut [Complex64],
        hamiltonian: &Array2<Complex64>,
        beta: f64,
    ) {
        use scirs2_core::ndarray::Array1;

        // Get the dimension of the Hamiltonian
        let dim = hamiltonian.nrows();

        // Determine which qubits this Hamiltonian acts on based on its dimension
        let num_target_qubits = (dim as f64).log2() as usize;

        if num_target_qubits == 0 || dim != (1 << num_target_qubits) {
            // Invalid dimension for quantum operator
            return;
        }

        // For small Hamiltonians (1 or 2 qubits), compute matrix exponential directly
        if num_target_qubits <= 2 {
            // Compute exp(-i * beta * H) using eigendecomposition or direct matrix exponential
            let evolution_op = self.compute_matrix_exponential(hamiltonian, -beta);

            // Apply the evolution operator to the relevant part of the state
            // For simplicity, assume it acts on the first num_target_qubits
            self.apply_small_unitary(state, &evolution_op, num_target_qubits);
        } else {
            // For larger systems, we would need Trotter decomposition or other advanced techniques
            // For now, log a warning and apply first-order Trotter approximation

            // First-order Trotter: exp(-i beta H) ≈ exp(-i beta H/n)^n for large n
            let n_trotter_steps = 10;
            let small_beta = beta / (n_trotter_steps as f64);

            for _ in 0..n_trotter_steps {
                let evolution_op = self.compute_matrix_exponential(hamiltonian, -small_beta);
                self.apply_small_unitary(state, &evolution_op, num_target_qubits);
            }
        }
    }

    /// Compute matrix exponential exp(i * angle * matrix)
    ///
    /// Uses Taylor series expansion for general matrices.
    /// For better performance, consider using eigendecomposition for Hermitian matrices.
    fn compute_matrix_exponential(
        &self,
        matrix: &Array2<Complex64>,
        angle: f64,
    ) -> Array2<Complex64> {
        use scirs2_core::ndarray::{s, Array2};

        let dim = matrix.nrows();
        let i_angle = Complex64::new(0.0, angle);

        // Scaled matrix: i * angle * H
        let scaled_matrix = matrix.mapv(|x| x * i_angle);

        // Compute matrix exponential using Taylor series: exp(A) = I + A + A^2/2! + A^3/3! + ...
        let mut result = Array2::eye(dim);
        let mut term = Array2::eye(dim);

        // Use up to 20 terms in Taylor series (should be sufficient for typical QAOA angles)
        for k in 1..=20 {
            // term = term * scaled_matrix / k
            term = term.dot(&scaled_matrix) / (k as f64);
            result = result + &term;

            // Check convergence
            let term_norm: f64 = term.iter().map(|x| x.norm_sqr()).sum::<f64>().sqrt();
            if term_norm < 1e-12 {
                break;
            }
        }

        result
    }

    /// Apply a small unitary operator (1-2 qubits) to the quantum state
    ///
    /// Assumes the unitary acts on the first `num_qubits` qubits of the system.
    /// For acting on different qubits, the state would need to be permuted.
    fn apply_small_unitary(
        &self,
        state: &mut [Complex64],
        unitary: &Array2<Complex64>,
        num_operator_qubits: usize,
    ) {
        use scirs2_core::ndarray::Array1;

        let operator_dim = 1 << num_operator_qubits;
        let remaining_dim = state.len() / operator_dim;

        // Process the state in blocks
        for block in 0..remaining_dim {
            let mut local_state = vec![Complex64::new(0.0, 0.0); operator_dim];

            // Extract local state for this block
            for i in 0..operator_dim {
                local_state[i] = state[block * operator_dim + i];
            }

            // Apply unitary: |ψ'⟩ = U |ψ⟩
            let state_vec = Array1::from_vec(local_state);
            let new_state = unitary.dot(&state_vec);

            // Write back the transformed state
            for i in 0..operator_dim {
                state[block * operator_dim + i] = new_state[i];
            }
        }
    }

    /// Apply a single-qubit Z rotation
    fn apply_z_rotation(&self, state: &mut [Complex64], qubit: usize, angle: f64) {
        let phase = Complex64::from_polar(1.0, -angle / 2.0);
        let phase_conj = phase.conj();

        let qubit_mask = 1 << qubit;

        for (i, amp) in state.iter_mut().enumerate() {
            if i & qubit_mask == 0 {
                *amp *= phase;
            } else {
                *amp *= phase_conj;
            }
        }
    }

    /// Apply a single-qubit X rotation using SciRS2 SIMD operations
    fn apply_x_rotation(&self, state: &mut [Complex64], qubit: usize, angle: f64) {
        let cos_half = (angle / 2.0).cos();
        let sin_half = (angle / 2.0).sin();

        let n = state.len();
        // let _qubit_mask = 1 << qubit;
        let stride = 1 << (qubit + 1);

        // Process in chunks for better cache efficiency
        for chunk_start in (0..n).step_by(stride) {
            for i in 0..(stride / 2) {
                let idx0 = chunk_start + i;
                let idx1 = idx0 + (stride / 2);

                if idx1 < n {
                    let amp0 = state[idx0];
                    let amp1 = state[idx1];

                    state[idx0] = amp0 * cos_half + amp1 * Complex64::new(0.0, -sin_half);
                    state[idx1] = amp1 * cos_half + amp0 * Complex64::new(0.0, -sin_half);
                }
            }
        }
    }

    /// Apply a two-qubit ZZ rotation
    fn apply_zz_rotation(&self, state: &mut [Complex64], qubit1: usize, qubit2: usize, angle: f64) {
        let phase = Complex64::from_polar(1.0, -angle / 2.0);

        let mask1 = 1 << qubit1;
        let mask2 = 1 << qubit2;

        for (i, amp) in state.iter_mut().enumerate() {
            let parity = ((i & mask1) >> qubit1) ^ ((i & mask2) >> qubit2);
            if parity == 0 {
                *amp *= phase;
            } else {
                *amp *= phase.conj();
            }
        }
    }

    /// Run the full QAOA circuit
    pub fn execute(&self, state: &mut [Complex64]) {
        // Initial state preparation
        self.prepare_initial_state(state);

        // Apply QAOA layers
        for layer in 0..self.params.layers {
            self.apply_cost_evolution(state, self.params.gamma[layer]);
            self.apply_mixer_evolution(state, self.params.beta[layer]);
        }

        // Normalize the state using SIMD operations
        let _ = simd_ops::normalize_simd(state);
    }

    /// Compute the expectation value of the cost Hamiltonian
    pub fn compute_expectation(&self, state: &[Complex64]) -> f64 {
        match &self.cost_hamiltonian {
            CostHamiltonian::MaxCut(edges) => {
                let mut expectation = 0.0;
                for &(i, j) in edges {
                    expectation += self.compute_zz_expectation(state, i, j);
                }
                edges.len() as f64 / 2.0 - expectation / 2.0
            }
            CostHamiltonian::WeightedMaxCut(weighted_edges) => {
                let mut expectation = 0.0;
                let mut total_weight = 0.0;
                for &(i, j, weight) in weighted_edges {
                    expectation += weight * self.compute_zz_expectation(state, i, j);
                    total_weight += weight;
                }
                total_weight / 2.0 - expectation / 2.0
            }
            CostHamiltonian::Ising { h, j } => {
                let mut expectation = 0.0;

                // Single-qubit terms
                for (i, &h_i) in h.iter().enumerate() {
                    if h_i.abs() > 1e-10 {
                        let num_qubits = (state.len() as f64).log2() as usize;
                        expectation += h_i * simd_ops::expectation_z_simd(state, i, num_qubits);
                    }
                }

                // Two-qubit terms
                for &((i, j), j_ij) in j {
                    if j_ij.abs() > 1e-10 {
                        expectation += j_ij * self.compute_zz_expectation(state, i, j);
                    }
                }

                expectation
            }
        }
    }

    /// Compute <ZZ> expectation value for two qubits
    fn compute_zz_expectation(&self, state: &[Complex64], qubit1: usize, qubit2: usize) -> f64 {
        let mask1 = 1 << qubit1;
        let mask2 = 1 << qubit2;

        let mut expectation = 0.0;
        for (i, amp) in state.iter().enumerate() {
            let bit1 = (i & mask1) >> qubit1;
            let bit2 = (i & mask2) >> qubit2;
            let sign = if bit1 == bit2 { 1.0 } else { -1.0 };
            expectation += sign * amp.probability();
        }

        expectation
    }

    /// Get the most probable bitstring from the final state
    pub fn get_solution(&self, state: &[Complex64]) -> Vec<bool> {
        let mut max_prob = 0.0;
        let mut max_idx = 0;

        for (i, amp) in state.iter().enumerate() {
            let prob = amp.probability();
            if prob > max_prob {
                max_prob = prob;
                max_idx = i;
            }
        }

        // Convert index to bitstring
        (0..self.num_qubits)
            .map(|i| (max_idx >> i) & 1 == 1)
            .collect()
    }
}

/// QAOA optimizer using classical optimization
pub struct QAOAOptimizer {
    circuit: QAOACircuit,
    #[allow(dead_code)]
    max_iterations: usize,
    #[allow(dead_code)]
    tolerance: f64,
}

impl QAOAOptimizer {
    /// Create a new QAOA optimizer
    pub const fn new(circuit: QAOACircuit, max_iterations: usize, tolerance: f64) -> Self {
        Self {
            circuit,
            max_iterations,
            tolerance,
        }
    }

    /// Execute the circuit with current parameters and return the state
    pub fn execute_circuit(&mut self) -> Vec<Complex64> {
        let state_size = 1 << self.circuit.num_qubits;
        let mut state = vec![Complex64::new(0.0, 0.0); state_size];
        self.circuit.execute(&mut state);
        state
    }

    /// Get the solution from a quantum state
    pub fn get_solution(&self, state: &[Complex64]) -> Vec<bool> {
        self.circuit.get_solution(state)
    }

    /// Optimize QAOA parameters using Nelder-Mead method
    pub fn optimize(&mut self) -> (Vec<f64>, Vec<f64>, f64) {
        let initial_beta = self.circuit.params.beta.clone();
        let initial_gamma = self.circuit.params.gamma.clone();

        let mut best_beta = initial_beta;
        let mut best_gamma = initial_gamma;
        let mut best_expectation = f64::NEG_INFINITY;

        // Simple grid search for better optimization
        let num_points = 5;
        for beta_scale in 0..num_points {
            for gamma_scale in 0..num_points {
                let beta_val = (beta_scale as f64) * std::f64::consts::PI / (num_points as f64);
                let gamma_val =
                    (gamma_scale as f64) * 2.0 * std::f64::consts::PI / (num_points as f64);

                let beta_params = vec![beta_val; self.circuit.params.layers];
                let gamma_params = vec![gamma_val; self.circuit.params.layers];

                self.circuit.params.beta.clone_from(&beta_params);
                self.circuit.params.gamma.clone_from(&gamma_params);

                let state = self.execute_circuit();
                let expectation = self.circuit.compute_expectation(&state);

                if expectation > best_expectation {
                    best_expectation = expectation;
                    best_beta = beta_params;
                    best_gamma = gamma_params;
                }
            }
        }

        self.circuit.params.beta.clone_from(&best_beta);
        self.circuit.params.gamma.clone_from(&best_gamma);

        (best_beta, best_gamma, best_expectation)
    }
}

/// Specialized MaxCut problem solver using QAOA
#[derive(Debug, Clone)]
pub struct MaxCutQAOA {
    /// Graph represented as adjacency list
    pub graph: Vec<Vec<usize>>,
    /// Edge weights (if weighted graph)
    pub weights: Option<Vec<Vec<f64>>>,
    /// Number of vertices
    pub num_vertices: usize,
    /// QAOA circuit
    circuit: Option<QAOACircuit>,
}

impl MaxCutQAOA {
    /// Create a new MaxCut QAOA solver
    pub fn new(graph: Vec<Vec<usize>>) -> Self {
        let num_vertices = graph.len();
        Self {
            graph,
            weights: None,
            num_vertices,
            circuit: None,
        }
    }

    /// Create with weighted edges
    #[must_use]
    pub fn with_weights(mut self, weights: Vec<Vec<f64>>) -> Self {
        assert_eq!(weights.len(), self.num_vertices);
        self.weights = Some(weights);
        self
    }

    /// Build the QAOA circuit for this MaxCut instance
    pub fn build_circuit(&mut self, layers: usize) -> &mut Self {
        let edges = self.extract_edges();

        let cost_hamiltonian = if let Some(ref weights) = self.weights {
            let weighted_edges = edges
                .iter()
                .map(|(i, j)| (*i, *j, weights[*i][*j]))
                .collect();
            CostHamiltonian::WeightedMaxCut(weighted_edges)
        } else {
            CostHamiltonian::MaxCut(edges)
        };

        let mixer_hamiltonian = MixerHamiltonian::TransverseField;
        let params = QAOAParams::random(layers);

        self.circuit = Some(QAOACircuit::new(
            self.num_vertices,
            cost_hamiltonian,
            mixer_hamiltonian,
            params,
        ));

        self
    }

    /// Extract edges from adjacency list
    fn extract_edges(&self) -> Vec<(usize, usize)> {
        let mut edges = Vec::new();
        for (i, neighbors) in self.graph.iter().enumerate() {
            for &j in neighbors {
                if i < j {
                    // Avoid duplicate edges
                    edges.push((i, j));
                }
            }
        }
        edges
    }

    /// Solve the MaxCut problem using QAOA
    pub fn solve(&mut self) -> (Vec<bool>, f64) {
        if self.circuit.is_none() {
            self.build_circuit(2); // Default to 2 layers
        }

        // SAFETY: We just called build_circuit above if None, so this is guaranteed to be Some
        let circuit = self
            .circuit
            .as_mut()
            .expect("Circuit should be built after build_circuit call");
        let mut optimizer = QAOAOptimizer::new(circuit.clone(), 100, 1e-6);

        let (_, _, best_expectation) = optimizer.optimize();
        let final_state = optimizer.execute_circuit();
        let solution = optimizer.get_solution(&final_state);

        (solution, best_expectation)
    }

    /// Evaluate a cut solution
    pub fn evaluate_cut(&self, solution: &[bool]) -> f64 {
        let mut cut_value = 0.0;

        for (i, neighbors) in self.graph.iter().enumerate() {
            for &j in neighbors {
                if i < j && solution[i] != solution[j] {
                    let weight = if let Some(ref weights) = self.weights {
                        weights[i][j]
                    } else {
                        1.0
                    };
                    cut_value += weight;
                }
            }
        }

        cut_value
    }

    /// Create a random graph for testing
    pub fn random_graph(num_vertices: usize, edge_probability: f64) -> Self {
        let mut graph = vec![Vec::new(); num_vertices];

        for i in 0..num_vertices {
            for j in i + 1..num_vertices {
                // Simple pseudo-random for reproducibility
                let rand_val = ((i * j) as f64).mul_add(1.234, 5.678).sin().abs();
                if rand_val < edge_probability {
                    graph[i].push(j);
                    graph[j].push(i);
                }
            }
        }

        Self::new(graph)
    }
}

/// Traveling Salesman Problem solver using QAOA
#[derive(Debug, Clone)]
pub struct TSPQAOA {
    /// Distance matrix between cities
    pub distances: Vec<Vec<f64>>,
    /// Number of cities
    pub num_cities: usize,
    /// QAOA circuit for TSP encoding
    circuit: Option<QAOACircuit>,
}

impl TSPQAOA {
    /// Create a new TSP QAOA solver
    pub fn new(distances: Vec<Vec<f64>>) -> Self {
        let num_cities = distances.len();
        assert!(distances.iter().all(|row| row.len() == num_cities));

        Self {
            distances,
            num_cities,
            circuit: None,
        }
    }

    /// Build QAOA circuit for TSP using binary encoding
    /// Each qubit x_{i,t} represents whether city i is visited at time t
    pub fn build_circuit(&mut self, layers: usize) -> &mut Self {
        let num_qubits = self.num_cities * self.num_cities;

        // Build TSP Hamiltonian with constraints
        let (h_fields, j_couplings) = self.build_tsp_hamiltonian();

        let cost_hamiltonian = CostHamiltonian::Ising {
            h: h_fields,
            j: j_couplings,
        };
        let mixer_hamiltonian = MixerHamiltonian::TransverseField;
        let params = QAOAParams::random(layers);

        self.circuit = Some(QAOACircuit::new(
            num_qubits,
            cost_hamiltonian,
            mixer_hamiltonian,
            params,
        ));

        self
    }

    /// Build the TSP Hamiltonian with penalty terms for constraints
    fn build_tsp_hamiltonian(&self) -> (Vec<f64>, Vec<((usize, usize), f64)>) {
        let n = self.num_cities;
        let num_qubits = n * n;

        let mut h_fields = vec![0.0; num_qubits];
        let mut j_couplings = Vec::new();

        let penalty_strength = 10.0; // Penalty for constraint violations

        // Distance terms in the objective
        for i in 0..n {
            for t in 0..n {
                let t_next = (t + 1) % n;
                for j in 0..n {
                    if i != j {
                        let qubit_it = i * n + t;
                        let qubit_jt_next = j * n + t_next;

                        // Add coupling for distance
                        j_couplings.push(((qubit_it, qubit_jt_next), self.distances[i][j] / 4.0));
                    }
                }
            }
        }

        // Constraint: each city visited exactly once
        for i in 0..n {
            // Linear term for normalization
            for t in 0..n {
                let qubit = i * n + t;
                h_fields[qubit] -= penalty_strength;
            }

            // Quadratic penalty terms
            for t1 in 0..n {
                for t2 in t1 + 1..n {
                    let qubit1 = i * n + t1;
                    let qubit2 = i * n + t2;
                    j_couplings.push(((qubit1, qubit2), penalty_strength));
                }
            }
        }

        // Constraint: each time step has exactly one city
        for t in 0..n {
            // Linear term
            for i in 0..n {
                let qubit = i * n + t;
                h_fields[qubit] -= penalty_strength;
            }

            // Quadratic penalty terms
            for i1 in 0..n {
                for i2 in i1 + 1..n {
                    let qubit1 = i1 * n + t;
                    let qubit2 = i2 * n + t;
                    j_couplings.push(((qubit1, qubit2), penalty_strength));
                }
            }
        }

        (h_fields, j_couplings)
    }

    /// Solve TSP using QAOA
    pub fn solve(&mut self) -> (Vec<usize>, f64) {
        if self.circuit.is_none() {
            self.build_circuit(3); // TSP typically needs more layers
        }

        // SAFETY: We just called build_circuit above if None, so this is guaranteed to be Some
        let circuit = self
            .circuit
            .as_mut()
            .expect("Circuit should be built after build_circuit call");
        let mut optimizer = QAOAOptimizer::new(circuit.clone(), 200, 1e-6);

        let (_, _, _best_expectation) = optimizer.optimize();
        let final_state = optimizer.execute_circuit();
        let bit_solution = optimizer.get_solution(&final_state);

        // Convert bit solution to TSP route
        let route = self.decode_tsp_solution(&bit_solution);
        let distance = self.evaluate_route(&route);

        (route, distance)
    }

    /// Decode binary solution to TSP route
    fn decode_tsp_solution(&self, bits: &[bool]) -> Vec<usize> {
        let n = self.num_cities;
        let mut route = Vec::new();

        for t in 0..n {
            let mut city_at_time_t = None;
            let mut _max_confidence = 0;

            // Find which city is most likely at time t
            for i in 0..n {
                let qubit_idx = i * n + t;
                if qubit_idx < bits.len() && bits[qubit_idx] {
                    _max_confidence += 1;
                    city_at_time_t = Some(i);
                }
            }

            // If no clear assignment, assign the first available city
            if city_at_time_t.is_none() {
                for i in 0..n {
                    if !route.contains(&i) {
                        city_at_time_t = Some(i);
                        break;
                    }
                }
            }

            if let Some(city) = city_at_time_t {
                if !route.contains(&city) {
                    route.push(city);
                }
            }
        }

        // Fill in missing cities
        for i in 0..n {
            if !route.contains(&i) {
                route.push(i);
            }
        }

        route
    }

    /// Evaluate the total distance of a route
    pub fn evaluate_route(&self, route: &[usize]) -> f64 {
        let mut total_distance = 0.0;

        for i in 0..route.len() {
            let current_city = route[i];
            let next_city = route[(i + 1) % route.len()];
            total_distance += self.distances[current_city][next_city];
        }

        total_distance
    }

    /// Create a random TSP instance
    pub fn random_instance(num_cities: usize) -> Self {
        let mut distances = vec![vec![0.0; num_cities]; num_cities];

        for i in 0..num_cities {
            for j in 0..num_cities {
                if i != j {
                    // Generate symmetric random distances
                    let dist = ((i * j + i + j) as f64)
                        .mul_add(1.234, 5.678)
                        .sin()
                        .abs()
                        .mul_add(100.0, 1.0);
                    distances[i][j] = dist;
                    distances[j][i] = dist;
                }
            }
        }

        Self::new(distances)
    }
}

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

    #[test]
    fn test_qaoa_params() {
        let params = QAOAParams::new(3);
        assert_eq!(params.layers, 3);
        assert_eq!(params.beta.len(), 3);
        assert_eq!(params.gamma.len(), 3);

        let random_params = QAOAParams::random(2);
        assert_eq!(random_params.layers, 2);
        assert!(random_params.beta.iter().all(|&x| x >= 0.0 && x <= PI));
    }

    #[test]
    fn test_maxcut_qaoa_simple() {
        // Simple triangle graph
        let graph = vec![
            vec![1, 2], // Vertex 0 connected to 1, 2
            vec![0, 2], // Vertex 1 connected to 0, 2
            vec![0, 1], // Vertex 2 connected to 0, 1
        ];

        let mut maxcut = MaxCutQAOA::new(graph);
        maxcut.build_circuit(1);

        let (solution, _expectation) = maxcut.solve();
        assert_eq!(solution.len(), 3);

        // Evaluate the cut - for a triangle, max cut should be 2
        let cut_value = maxcut.evaluate_cut(&solution);
        println!(
            "Triangle MaxCut solution: {:?}, value: {}",
            solution, cut_value
        );

        // Any valid cut of a triangle should have value 2
        let expected_cut_values = [0.0, 2.0]; // Either all same (0) or optimal (2)
        assert!(expected_cut_values.contains(&cut_value) || cut_value == 1.0);
    }

    #[test]
    fn test_maxcut_weighted() {
        let graph = vec![vec![1], vec![0]];

        let weights = vec![vec![0.0, 5.0], vec![5.0, 0.0]];

        let mut maxcut = MaxCutQAOA::new(graph).with_weights(weights);
        maxcut.build_circuit(1);

        let (solution, _expectation) = maxcut.solve();
        let cut_value = maxcut.evaluate_cut(&solution);

        // For two vertices with weight 5, optimal cut should be 5
        println!(
            "Weighted MaxCut solution: {:?}, value: {}",
            solution, cut_value
        );
        assert!(cut_value >= 0.0);
    }

    #[test]
    fn test_maxcut_random_graph() {
        let maxcut = MaxCutQAOA::random_graph(4, 0.5);
        assert_eq!(maxcut.num_vertices, 4);
        println!("Random graph: {:?}", maxcut.graph);
    }

    #[test]
    fn test_tsp_qaoa_simple() {
        // Simple 3-city TSP
        let distances = vec![
            vec![0.0, 1.0, 2.0],
            vec![1.0, 0.0, 1.5],
            vec![2.0, 1.5, 0.0],
        ];

        let mut tsp = TSPQAOA::new(distances);
        tsp.build_circuit(2);

        let (route, distance) = tsp.solve();
        assert_eq!(route.len(), 3);

        // Verify it's a valid permutation
        let mut sorted_route = route.clone();
        sorted_route.sort();
        assert_eq!(sorted_route, vec![0, 1, 2]);

        println!("TSP route: {:?}, distance: {}", route, distance);
        assert!(distance > 0.0);
    }

    #[test]
    fn test_tsp_evaluation() {
        let distances = vec![
            vec![0.0, 1.0, 3.0],
            vec![1.0, 0.0, 2.0],
            vec![3.0, 2.0, 0.0],
        ];

        let tsp = TSPQAOA::new(distances);

        // Test route 0 -> 1 -> 2 -> 0
        let route = vec![0, 1, 2];
        let distance = tsp.evaluate_route(&route);

        // Should be 1 + 2 + 3 = 6
        assert_eq!(distance, 6.0);
    }

    #[test]
    fn test_tsp_random_instance() {
        let tsp = TSPQAOA::random_instance(4);
        assert_eq!(tsp.num_cities, 4);
        assert_eq!(tsp.distances.len(), 4);

        // Check symmetry
        for i in 0..4 {
            for j in 0..4 {
                if i != j {
                    assert_eq!(tsp.distances[i][j], tsp.distances[j][i]);
                }
            }
        }
    }

    #[test]
    fn test_qaoa_circuit_execution() {
        let edges = vec![(0, 1), (1, 2)];
        let cost_hamiltonian = CostHamiltonian::MaxCut(edges);
        let mixer_hamiltonian = MixerHamiltonian::TransverseField;
        let params = QAOAParams::new(1);

        let circuit = QAOACircuit::new(3, cost_hamiltonian, mixer_hamiltonian, params);

        let state_size = 1 << 3;
        let mut state = vec![Complex64::new(0.0, 0.0); state_size];

        circuit.execute(&mut state);

        // Check state is normalized
        let norm_sq: f64 = state.iter().map(|c| c.norm_sqr()).sum();
        assert!((norm_sq - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_qaoa_optimizer_simple() {
        let edges = vec![(0, 1)];
        let cost_hamiltonian = CostHamiltonian::MaxCut(edges);
        let mixer_hamiltonian = MixerHamiltonian::TransverseField;
        let params = QAOAParams::random(1);

        let circuit = QAOACircuit::new(2, cost_hamiltonian, mixer_hamiltonian, params);
        let mut optimizer = QAOAOptimizer::new(circuit, 10, 1e-6);

        let (_beta, _gamma, expectation) = optimizer.optimize();

        // Should find some reasonable expectation value
        assert!(expectation.is_finite());
        println!("QAOA optimizer result: expectation = {}", expectation);
    }

    #[test]
    fn test_custom_mixer_hamiltonian() {
        use scirs2_core::ndarray::array;

        // Create a simple custom mixer: Pauli-X on each qubit
        // X = [[0, 1], [1, 0]]
        let pauli_x = array![
            [Complex64::new(0.0, 0.0), Complex64::new(1.0, 0.0)],
            [Complex64::new(1.0, 0.0), Complex64::new(0.0, 0.0)]
        ];

        // Use two X operators for a 2-qubit system
        let custom_mixers = vec![pauli_x.clone(), pauli_x];

        let edges = vec![(0, 1)];
        let cost_hamiltonian = CostHamiltonian::MaxCut(edges);
        let mixer_hamiltonian = MixerHamiltonian::Custom(custom_mixers);
        let params = QAOAParams::new(1);

        let circuit = QAOACircuit::new(2, cost_hamiltonian, mixer_hamiltonian, params);

        let state_size = 1 << 2;
        let mut state = vec![Complex64::new(0.0, 0.0); state_size];

        // Execute the circuit
        circuit.execute(&mut state);

        // Check state is normalized
        let norm_sq: f64 = state.iter().map(|c| c.norm_sqr()).sum();
        assert!((norm_sq - 1.0).abs() < 1e-10);

        println!("Custom mixer QAOA state: {:?}", state);
    }

    #[test]
    fn test_custom_mixer_vs_standard() {
        use scirs2_core::ndarray::array;

        // Create identical circuits with standard and custom mixers
        let edges = vec![(0, 1)];
        let cost_hamiltonian1 = CostHamiltonian::MaxCut(edges.clone());
        let cost_hamiltonian2 = CostHamiltonian::MaxCut(edges);

        // Standard X-mixer
        let mixer1 = MixerHamiltonian::TransverseField;

        // Custom X-mixer (should produce same result)
        let pauli_x = array![
            [Complex64::new(0.0, 0.0), Complex64::new(1.0, 0.0)],
            [Complex64::new(1.0, 0.0), Complex64::new(0.0, 0.0)]
        ];
        let custom_mixers = vec![pauli_x.clone(), pauli_x];
        let mixer2 = MixerHamiltonian::Custom(custom_mixers);

        let params = QAOAParams::new(1);

        let circuit1 = QAOACircuit::new(2, cost_hamiltonian1, mixer1, params.clone());
        let circuit2 = QAOACircuit::new(2, cost_hamiltonian2, mixer2, params);

        let state_size = 1 << 2;
        let mut state1 = vec![Complex64::new(0.0, 0.0); state_size];
        let mut state2 = vec![Complex64::new(0.0, 0.0); state_size];

        circuit1.execute(&mut state1);
        circuit2.execute(&mut state2);

        // The states should be similar (may not be exactly equal due to numerical differences)
        println!("Standard mixer state: {:?}", state1);
        println!("Custom mixer state: {:?}", state2);

        // Check both are normalized
        let norm1: f64 = state1.iter().map(|c| c.norm_sqr()).sum();
        let norm2: f64 = state2.iter().map(|c| c.norm_sqr()).sum();
        assert!((norm1 - 1.0).abs() < 1e-10);
        assert!((norm2 - 1.0).abs() < 1e-10);
    }
}