//! Optimized quantum gate operations using SIMD and memory-efficient algorithms
//!
//! This module provides highly optimized implementations of quantum gate operations
//! that leverage SIMD instructions, parallel processing, and memory-efficient algorithms
//! to enable simulation of quantum circuits with 30+ qubits.
//!
//! It uses the scirs2-core library for optimizations where available, and includes 
//! placeholder implementations for features that are still under development.

use num_complex::Complex64;
use ndarray::{Array1, Array2, ArrayView1, ArrayView2};
use rayon::prelude::*;

// Import features from scirs2-core if available
use scirs2_core::simd as scirs_simd;
use scirs2_core::parallel as scirs_parallel;
use scirs2_core::ndarray_ext as scirs_ndarray;

// For SIMD operations
#[cfg(feature = "simd")]
use wide::{f32x8, f64x4};

use crate::utils::flip_bit;

/// Represents a quantum state vector that can be efficiently operated on
pub struct OptimizedStateVector {
    /// The full state vector as a complex vector
    state: Vec<Complex64>,
    /// Number of qubits represented
    num_qubits: usize,
    /// Use parallel execution
    parallel: bool,
    /// Use SIMD operations
    simd: bool,
    /// Use chunked memory processing
    chunked: bool,
    /// Chunk size for memory-efficient processing
    chunk_size: usize,
}

impl OptimizedStateVector {
    /// Create a new optimized state vector for given number of qubits
    pub fn new(num_qubits: usize) -> Self {
        let dim = 1 << num_qubits;
        let mut state = vec![Complex64::new(0.0, 0.0); dim];
        state[0] = Complex64::new(1.0, 0.0); // Initialize to |0...0>
        
        Self {
            state,
            num_qubits,
            parallel: true,
            simd: true,
            chunked: num_qubits > 27, // Use chunked processing for 27+ qubits
            chunk_size: if num_qubits > 27 { 1 << 26 } else { dim }, // 64M chunks (512MB)
        }
    }
    
    /// Enable or disable parallel execution
    pub fn with_parallel(mut self, parallel: bool) -> Self {
        self.parallel = parallel;
        self
    }
    
    /// Enable or disable SIMD operations
    pub fn with_simd(mut self, simd: bool) -> Self {
        self.simd = simd;
        self
    }
    
    /// Enable or disable chunked memory processing
    pub fn with_chunked(mut self, chunked: bool, chunk_size: Option<usize>) -> Self {
        self.chunked = chunked;
        if let Some(size) = chunk_size {
            self.chunk_size = size;
        }
        self
    }
    
    /// Get a reference to the state vector
    pub fn state(&self) -> &[Complex64] {
        &self.state
    }
    
    /// Get a mutable reference to the state vector
    pub fn state_mut(&mut self) -> &mut [Complex64] {
        &mut self.state
    }
    
    /// Get the number of qubits
    pub fn num_qubits(&self) -> usize {
        self.num_qubits
    }
    
    /// Get the dimension of the state vector
    pub fn dimension(&self) -> usize {
        1 << self.num_qubits
    }
    
    /// Apply a single-qubit gate to the state vector
    ///
    /// # Arguments
    ///
    /// * `matrix` - The 2x2 matrix representation of the gate
    /// * `target` - The target qubit index
    pub fn apply_single_qubit_gate(&mut self, matrix: &[Complex64], target: usize) {
        if target >= self.num_qubits {
            panic!("Target qubit index out of range");
        }

        // TEMPORARILY DISABLED SIMD and parallel implementations due to issues
        // TODO: Re-enable when implementations are fixed
        self.apply_single_qubit_gate_sequential(matrix, target);
    }
    
    /// SIMD-accelerated implementation of single-qubit gate application
    #[cfg(feature = "simd")]
    fn apply_single_qubit_gate_simd(&mut self, matrix: &[Complex64], target: usize) {
        // Use the SIMD-optimized implementation from the simd module
        simd::apply_single_qubit_gate_simd(
            &mut self.state,
            matrix,
            target,
            self.num_qubits
        );
    }

    /// Fallback implementation when SIMD is not available
    #[cfg(not(feature = "simd"))]
    fn apply_single_qubit_gate_simd(&mut self, matrix: &[Complex64], target: usize) {
        // Fall back to parallel implementation
        self.apply_single_qubit_gate_parallel(matrix, target);
    }
    
    /// Parallel implementation of single-qubit gate application
    fn apply_single_qubit_gate_parallel(&mut self, matrix: &[Complex64], target: usize) {
        // Create a new vector for the results to avoid race conditions
        let mut new_state = vec![Complex64::new(0.0, 0.0); self.state.len()];

        if self.chunked {
            // Process in chunks to reduce memory pressure
            let chunk_size = self.chunk_size;
            let dim = self.state.len();
            let num_chunks = (dim + chunk_size - 1) / chunk_size;

            (0..num_chunks).into_par_iter().for_each(|chunk_idx| {
                let start = chunk_idx * chunk_size;
                let end = std::cmp::min(start + chunk_size, dim);

                for i in start..end {
                    let bit_val = (i >> target) & 1;
                    let paired_idx = flip_bit(i, target);

                    if bit_val == 0 {
                        // For bit value 0, apply first two elements of matrix
                        new_state[i] = matrix[0] * self.state[i] + matrix[1] * self.state[paired_idx];
                        new_state[paired_idx] = matrix[2] * self.state[i] + matrix[3] * self.state[paired_idx];
                    }
                    // Skip if bit_val is 1 as we've already processed this pair
                }
            });
        } else {
            // Process state vector by groups of basis pairs
            // This approach avoids race conditions by ensuring each thread processes
            // complete pairs of basis states affected by the gate

            // Number of basis state pairs affected by the gate: 2^(num_qubits-1)
            let num_pairs = 1 << (self.num_qubits - 1);

            (0..num_pairs).into_par_iter().for_each(|pair_idx| {
                // Compute the indices of the paired basis states
                // Each pair consists of two states differing only in the target qubit
                let idx_mask = !(1 << target);  // Mask that zeros the target bit
                let base_idx = (pair_idx & idx_mask) | ((pair_idx & (idx_mask >> 1)) << 1);
                let idx0 = base_idx;  // Basis state with target bit = 0
                let idx1 = base_idx | (1 << target);  // Basis state with target bit = 1

                // Apply the unitary matrix to the basis state pair
                new_state[idx0] = matrix[0] * self.state[idx0] + matrix[1] * self.state[idx1];
                new_state[idx1] = matrix[2] * self.state[idx0] + matrix[3] * self.state[idx1];
            });
        }

        // Update the state vector
        self.state = new_state;
    }
    
    /// Sequential implementation of single-qubit gate application
    fn apply_single_qubit_gate_sequential(&mut self, matrix: &[Complex64], target: usize) {
        let dim = self.state.len();
        let mut new_state = vec![Complex64::new(0.0, 0.0); dim];

        // For each pair of states that differ only in the target bit
        for i in 0..dim {
            let bit_val = (i >> target) & 1;

            // Only process each pair once (when target bit is 0)
            if bit_val == 0 {
                let paired_idx = flip_bit(i, target);

                // |i⟩ has target bit 0, |paired_idx⟩ has target bit 1
                let a0 = self.state[i];         // Amplitude for |i⟩
                let a1 = self.state[paired_idx]; // Amplitude for |paired_idx⟩

                // Apply the 2x2 unitary matrix:
                // [ matrix[0] matrix[1] ] [ a0 ] = [ new_a0 ]
                // [ matrix[2] matrix[3] ] [ a1 ]   [ new_a1 ]

                new_state[i] = matrix[0] * a0 + matrix[1] * a1;
                new_state[paired_idx] = matrix[2] * a0 + matrix[3] * a1;
            }
        }

        self.state = new_state;
    }
    
    /// Apply a controlled-NOT gate to the state vector
    ///
    /// # Arguments
    ///
    /// * `control` - The control qubit index
    /// * `target` - The target qubit index
    pub fn apply_cnot(&mut self, control: usize, target: usize) {
        if control >= self.num_qubits || target >= self.num_qubits {
            panic!("Qubit indices out of range");
        }

        if control == target {
            panic!("Control and target qubits must be different");
        }

        // TEMPORARILY DISABLED SIMD and parallel implementations due to issues
        // TODO: Re-enable when implementations are fixed
        self.apply_cnot_sequential(control, target);
    }
    
    /// SIMD-accelerated implementation of CNOT gate application
    #[cfg(feature = "simd")]
    fn apply_cnot_simd(&mut self, control: usize, target: usize) {
        // Use the SIMD-optimized implementation from the simd module
        simd::apply_cnot_simd(
            &mut self.state,
            control,
            target,
            self.num_qubits
        );
    }

    /// Fallback implementation when SIMD is not available
    #[cfg(not(feature = "simd"))]
    fn apply_cnot_simd(&mut self, control: usize, target: usize) {
        // Fall back to parallel implementation
        self.apply_cnot_parallel(control, target);
    }
    
    /// Parallel implementation of CNOT gate application
    fn apply_cnot_parallel(&mut self, control: usize, target: usize) {
        let mut new_state = vec![Complex64::new(0.0, 0.0); self.state.len()];

        if self.chunked {
            // Process in chunks to reduce memory pressure
            let chunk_size = self.chunk_size;
            let dim = self.state.len();
            let num_chunks = (dim + chunk_size - 1) / chunk_size;

            (0..num_chunks).into_par_iter().for_each(|chunk_idx| {
                let start = chunk_idx * chunk_size;
                let end = std::cmp::min(start + chunk_size, dim);

                for i in start..end {
                    let control_bit = (i >> control) & 1;

                    if control_bit == 1 {
                        // If control bit is 1, flip the target bit
                        let flipped = flip_bit(i, target);
                        new_state[i] = self.state[flipped];
                    } else {
                        // If control bit is 0, leave the state unchanged
                        new_state[i] = self.state[i];
                    }
                }
            });
        } else {
            // Process state vector by groups of paired basis states
            // For CNOT, we need to process groups of 4 basis states at a time
            // (all combinations of control and target bits)

            // Compute the number of groups of 4 basis states
            let group_size = 1 << 2; // 4 states per group (00, 01, 10, 11)
            let num_groups = self.state.len() / group_size;

            if control > target {
                // When control > target, we need a different grouping strategy
                (0..num_groups).into_par_iter().for_each(|group_idx| {
                    // Calculate base index for this group
                    let mut base_idx = group_idx;
                    let c_mask = !(1 << control);
                    let t_mask = !(1 << target);
                    let c_bit = 1 << control;
                    let t_bit = 1 << target;

                    // Convert group_idx to base_idx with control and target bits zeroed
                    let control_high_bits = (group_idx & ((1 << control) - 1));
                    let control_low_bits = (group_idx >> control) << (control + 1);
                    let idx_without_control = control_low_bits | control_high_bits;

                    let target_high_bits = (idx_without_control & ((1 << target) - 1));
                    let target_low_bits = (idx_without_control >> target) << (target + 1);
                    base_idx = target_low_bits | target_high_bits;

                    // Process all 4 combinations of control and target bits
                    // 00: No change
                    new_state[base_idx] = self.state[base_idx];

                    // 01: No change
                    new_state[base_idx | t_bit] = self.state[base_idx | t_bit];

                    // 10: Flip target
                    new_state[base_idx | c_bit] = self.state[base_idx | c_bit | t_bit];

                    // 11: Flip target
                    new_state[base_idx | c_bit | t_bit] = self.state[base_idx | c_bit];
                });
            } else {
                // Simple case: for each pair of states with control bit = 1, swap the target bit values
                (0..self.state.len() >> 1).into_par_iter().for_each(|idx| {
                    let control_bit = (idx >> (control - 1)) & 1;

                    if control_bit == 0 {
                        // Control is 0, copy state as is
                        new_state[idx] = self.state[idx];
                    } else {
                        // Control is 1, flip target bit
                        let target_bit = (idx >> (target - 1)) & 1;
                        let new_target = target_bit ^ 1;

                        // Calculate new index with flipped target
                        let idx_mask = !(1 << (target - 1));
                        let new_idx = (idx & idx_mask) | (new_target << (target - 1));

                        new_state[idx] = self.state[new_idx];
                    }
                });
            }
        }

        // Update the state vector
        self.state = new_state;
    }
    
    /// Sequential implementation of CNOT gate application
    fn apply_cnot_sequential(&mut self, control: usize, target: usize) {
        let dim = self.state.len();
        let mut new_state = vec![Complex64::new(0.0, 0.0); dim];
        
        for i in 0..dim {
            let control_bit = (i >> control) & 1;
            if control_bit == 1 {
                let flipped = flip_bit(i, target);
                new_state[i] = self.state[flipped];
            } else {
                new_state[i] = self.state[i];
            }
        }

        self.state = new_state;
    }
    
    /// Apply a two-qubit gate to the state vector
    ///
    /// # Arguments
    ///
    /// * `matrix` - The 4x4 matrix representation of the gate
    /// * `qubit1` - The first qubit index
    /// * `qubit2` - The second qubit index
    pub fn apply_two_qubit_gate(&mut self, matrix: &[Complex64], qubit1: usize, qubit2: usize) {
        if qubit1 >= self.num_qubits || qubit2 >= self.num_qubits {
            panic!("Qubit indices out of range");
        }
        
        if qubit1 == qubit2 {
            panic!("Qubit indices must be different");
        }
        
        if self.simd {
            // Use SIMD-optimized implementation when available
            self.apply_two_qubit_gate_simd(matrix, qubit1, qubit2);
        } else if self.parallel {
            // Use parallel implementation
            self.apply_two_qubit_gate_parallel(matrix, qubit1, qubit2);
        } else {
            // Use sequential implementation
            self.apply_two_qubit_gate_sequential(matrix, qubit1, qubit2);
        }
    }
    
    /// SIMD-accelerated implementation of two-qubit gate application
    /// 
    /// This is a placeholder for the full SIMD implementation.
    fn apply_two_qubit_gate_simd(&mut self, matrix: &[Complex64], qubit1: usize, qubit2: usize) {
        // Placeholder for SIMD implementation
        // For now, fall back to parallel implementation
        self.apply_two_qubit_gate_parallel(matrix, qubit1, qubit2);
    }
    
    /// Parallel implementation of two-qubit gate application
    fn apply_two_qubit_gate_parallel(&mut self, matrix: &[Complex64], qubit1: usize, qubit2: usize) {
        let state_copy = self.state.clone();
        
        if self.chunked {
            // Process in chunks to reduce memory pressure
            let chunks = self.state.par_chunks_mut(self.chunk_size);
            let state_copy_ref = &state_copy;
            
            chunks.enumerate().for_each(|(chunk_idx, chunk)| {
                let base_idx = chunk_idx * self.chunk_size;
                
                chunk.iter_mut().enumerate().for_each(|(i, amp)| {
                    let idx = base_idx + i;
                    
                    // Calculate indices for all four basis states
                    let b00 = idx & !(1 << qubit1) & !(1 << qubit2);
                    let b01 = b00 | (1 << qubit2);
                    let b10 = b00 | (1 << qubit1);
                    let b11 = b00 | (1 << qubit1) | (1 << qubit2);
                    
                    // Find which basis state this index corresponds to
                    let bit1 = (idx >> qubit1) & 1;
                    let bit2 = (idx >> qubit2) & 1;
                    
                    // Matrix indices for the amplitudes
                    let mat_idx = (bit1 << 1) | bit2;
                    
                    // Apply the 4x4 matrix
                    *amp = matrix[mat_idx * 4 + 0] * state_copy_ref[b00] +
                           matrix[mat_idx * 4 + 1] * state_copy_ref[b01] +
                           matrix[mat_idx * 4 + 2] * state_copy_ref[b10] +
                           matrix[mat_idx * 4 + 3] * state_copy_ref[b11];
                });
            });
        } else {
            // Process entire state vector at once
            self.state.par_iter_mut().enumerate().for_each(|(idx, amp)| {
                // Calculate indices for all four basis states
                let b00 = idx & !(1 << qubit1) & !(1 << qubit2);
                let b01 = b00 | (1 << qubit2);
                let b10 = b00 | (1 << qubit1);
                let b11 = b00 | (1 << qubit1) | (1 << qubit2);
                
                // Find which basis state this index corresponds to
                let bit1 = (idx >> qubit1) & 1;
                let bit2 = (idx >> qubit2) & 1;
                
                // Matrix indices for the amplitudes
                let mat_idx = (bit1 << 1) | bit2;
                
                // Apply the 4x4 matrix
                *amp = matrix[mat_idx * 4 + 0] * state_copy[b00] +
                       matrix[mat_idx * 4 + 1] * state_copy[b01] +
                       matrix[mat_idx * 4 + 2] * state_copy[b10] +
                       matrix[mat_idx * 4 + 3] * state_copy[b11];
            });
        }
    }
    
    /// Sequential implementation of two-qubit gate application
    fn apply_two_qubit_gate_sequential(&mut self, matrix: &[Complex64], qubit1: usize, qubit2: usize) {
        let dim = self.state.len();
        let mut new_state = vec![Complex64::new(0.0, 0.0); dim];
        
        for idx in 0..dim {
            // Calculate indices for all four basis states
            let b00 = idx & !(1 << qubit1) & !(1 << qubit2);
            let b01 = b00 | (1 << qubit2);
            let b10 = b00 | (1 << qubit1);
            let b11 = b00 | (1 << qubit1) | (1 << qubit2);
            
            // Find which basis state this index corresponds to
            let bit1 = (idx >> qubit1) & 1;
            let bit2 = (idx >> qubit2) & 1;
            
            // Matrix indices for the amplitudes
            let mat_idx = (bit1 << 1) | bit2;
            
            // Apply the 4x4 matrix
            new_state[idx] = matrix[mat_idx * 4 + 0] * self.state[b00] +
                            matrix[mat_idx * 4 + 1] * self.state[b01] +
                            matrix[mat_idx * 4 + 2] * self.state[b10] +
                            matrix[mat_idx * 4 + 3] * self.state[b11];
        }
        
        self.state.copy_from_slice(&new_state);
    }
    
    /// Calculate probability of measuring a specific bit string
    pub fn probability(&self, bit_string: &[u8]) -> f64 {
        if bit_string.len() != self.num_qubits {
            panic!("Bit string length must match number of qubits");
        }
        
        // Convert bit string to index
        let mut idx = 0;
        for (i, &bit) in bit_string.iter().enumerate() {
            if bit != 0 {
                idx |= 1 << i;
            }
        }
        
        // Return probability
        self.state[idx].norm_sqr()
    }
    
    /// Calculate probabilities for all basis states
    pub fn probabilities(&self) -> Vec<f64> {
        self.state.iter().map(|a| a.norm_sqr()).collect()
    }
}

/// Helper functions for memory-efficient quantum simulation
pub mod memory {
    use super::*;
    
    /// A memory-efficient representation of a quantum state for large qubit counts
    ///
    /// This implementation uses a chunked approach to state vector storage,
    /// allowing for simulation of larger qubit counts without exhausting memory.
    pub struct ChunkedStateVector {
        // Number of qubits in the system
        num_qubits: usize,
        // Number of chunks to split the state vector into
        num_chunks: usize,
        // The maximum number of qubits that can be represented in a single chunk
        chunk_qubits: usize,
        // The current active chunk
        current_chunk: usize,
        // The currently loaded chunk of the state vector
        chunk_data: Vec<Complex64>,
    }
    
    impl ChunkedStateVector {
        /// Create a new chunked state vector for a given number of qubits
        pub fn new(num_qubits: usize) -> Self {
            // Determine how many qubits we can fit in a chunk
            // This is a simplified heuristic; in practice, this would be more sophisticated
            let chunk_qubits = if num_qubits > 30 { 27 } else { num_qubits };
            let num_chunks = 1 << (num_qubits - chunk_qubits);
            
            // Initialize with first chunk containing |0...0> state
            let mut chunk_data = vec![Complex64::new(0.0, 0.0); 1 << chunk_qubits];
            chunk_data[0] = Complex64::new(1.0, 0.0);
            
            Self {
                num_qubits,
                num_chunks,
                chunk_qubits,
                current_chunk: 0,
                chunk_data,
            }
        }
        
        /// Get the number of qubits in the system
        pub fn num_qubits(&self) -> usize {
            self.num_qubits
        }
        
        /// Get the total dimension of the state vector
        pub fn dimension(&self) -> usize {
            1 << self.num_qubits
        }
        
        /// Get the number of chunks
        pub fn num_chunks(&self) -> usize {
            self.num_chunks
        }
        
        /// Get the currently loaded chunk
        pub fn current_chunk(&self) -> usize {
            self.current_chunk
        }
        
        /// Load a specific chunk into memory
        pub fn load_chunk(&mut self, chunk_idx: usize) {
            // In a real implementation, this would load the chunk from disk or other storage
            // For this placeholder, we simulate the loading by zeroing the chunk data
            if chunk_idx >= self.num_chunks {
                panic!("Chunk index out of range");
            }
            
            self.current_chunk = chunk_idx;
            self.chunk_data = vec![Complex64::new(0.0, 0.0); 1 << self.chunk_qubits];
            
            // If this is chunk 0 and we have not performed any operations yet,
            // initialize it to the |0...0> state
            if chunk_idx == 0 {
                self.chunk_data[0] = Complex64::new(1.0, 0.0);
            }
        }
        
        /// Save the current chunk
        pub fn save_chunk(&self) {
            // In a real implementation, this would save the chunk to disk or other storage
            // For this placeholder, we do nothing
        }
        
        /// Apply a single-qubit gate to the state vector
        ///
        /// This is a placeholder implementation that shows the concept but
        /// does not fully implement the chunked operations.
        pub fn apply_single_qubit_gate(&mut self, matrix: &[Complex64], target: usize) {
            // In a real implementation, this would load each chunk, apply the gate,
            // and save the result. The actual implementation would be much more complex
            // and would depend on the memory model and storage system used.
            
            // For this placeholder, we simply apply the gate to the current chunk
            let dim = self.chunk_data.len();
            let chunk_offset = self.current_chunk * dim;
            
            // Check if the target qubit affects this chunk
            if target < self.chunk_qubits {
                // Target qubit is within this chunk, apply gate directly
                let mut new_chunk = vec![Complex64::new(0.0, 0.0); dim];
                
                for i in 0..dim {
                    let bit_val = (i >> target) & 1;
                    let paired_idx = i ^ (1 << target);
                    
                    if bit_val == 0 {
                        new_chunk[i] = matrix[0] * self.chunk_data[i] + matrix[1] * self.chunk_data[paired_idx];
                        new_chunk[paired_idx] = matrix[2] * self.chunk_data[i] + matrix[3] * self.chunk_data[paired_idx];
                    }
                }
                
                self.chunk_data = new_chunk;
            } else {
                // Target qubit spans chunks, this requires cross-chunk operations
                // In a real implementation, this would be a complex operation involving
                // loading pairs of chunks, applying the gate, and saving the results
                
                // For this placeholder, we just note that this is not implemented
                println!("Warning: Cross-chunk qubit operations not implemented in placeholder");
            }
        }
    }
}

/// SIMD-accelerated operations for quantum simulation
pub mod simd {
    use super::*;
    
    // Import the wide crate for SIMD operations if available
    #[cfg(feature = "simd")]
    use wide::{f32x4, f32x8, f64x2, f64x4};
    
    /// Apply a single-qubit gate using SIMD operations
    ///
    /// This implementation uses SIMD operations to apply a single-qubit gate
    /// to the state vector more efficiently, processing multiple amplitudes at once.
    #[cfg(feature = "simd")]
    pub fn apply_single_qubit_gate_simd(
        state: &mut [Complex64],
        matrix: &[Complex64],
        target: usize,
        num_qubits: usize
    ) {
        let dim = state.len();
        let mut new_state = vec![Complex64::new(0.0, 0.0); dim];

        // SIMD-optimized implementation for applying the gate
        let m00_real = matrix[0].re;
        let m00_imag = matrix[0].im;
        let m01_real = matrix[1].re;
        let m01_imag = matrix[1].im;
        let m10_real = matrix[2].re;
        let m10_imag = matrix[2].im;
        let m11_real = matrix[3].re;
        let m11_imag = matrix[3].im;

        // Create SIMD vectors for the matrix elements
        let m00_real_vec = f64x4::splat(m00_real);
        let m00_imag_vec = f64x4::splat(m00_imag);
        let m01_real_vec = f64x4::splat(m01_real);
        let m01_imag_vec = f64x4::splat(m01_imag);
        let m10_real_vec = f64x4::splat(m10_real);
        let m10_imag_vec = f64x4::splat(m10_imag);
        let m11_real_vec = f64x4::splat(m11_real);
        let m11_imag_vec = f64x4::splat(m11_imag);

        // Process chunks that are aligned with the target qubit
        let chunk_size = 1 << (target + 1);
        let num_chunks = dim / chunk_size;

        for chunk_idx in 0..num_chunks {
            let chunk_offset = chunk_idx * chunk_size;
            let zero_offset = chunk_offset;
            let one_offset = chunk_offset + (1 << target);

            // Process amplitudes in groups of 4 for SIMD
            for i in 0..(1 << target) / 4 {
                let idx_offset = i * 4;

                // Load amplitudes for |0⟩ states
                let zero_indices = [
                    zero_offset + idx_offset,
                    zero_offset + idx_offset + 1,
                    zero_offset + idx_offset + 2,
                    zero_offset + idx_offset + 3,
                ];

                // Load amplitudes for |1⟩ states
                let one_indices = [
                    one_offset + idx_offset,
                    one_offset + idx_offset + 1,
                    one_offset + idx_offset + 2,
                    one_offset + idx_offset + 3,
                ];

                // Load state vector values
                let a0_real = f64x4::new([
                    state[zero_indices[0]].re,
                    state[zero_indices[1]].re,
                    state[zero_indices[2]].re,
                    state[zero_indices[3]].re,
                ]);

                let a0_imag = f64x4::new([
                    state[zero_indices[0]].im,
                    state[zero_indices[1]].im,
                    state[zero_indices[2]].im,
                    state[zero_indices[3]].im,
                ]);

                let a1_real = f64x4::new([
                    state[one_indices[0]].re,
                    state[one_indices[1]].re,
                    state[one_indices[2]].re,
                    state[one_indices[3]].re,
                ]);

                let a1_imag = f64x4::new([
                    state[one_indices[0]].im,
                    state[one_indices[1]].im,
                    state[one_indices[2]].im,
                    state[one_indices[3]].im,
                ]);

                // Compute matrix-vector multiplication using SIMD
                // b0 = m00 * a0 + m01 * a1
                // b1 = m10 * a0 + m11 * a1

                // Real part: m_real * a_real - m_imag * a_imag
                // Imag part: m_real * a_imag + m_imag * a_real

                // Compute new |0⟩ amplitudes
                let b0_real = m00_real_vec * a0_real - m00_imag_vec * a0_imag +
                              m01_real_vec * a1_real - m01_imag_vec * a1_imag;

                let b0_imag = m00_real_vec * a0_imag + m00_imag_vec * a0_real +
                              m01_real_vec * a1_imag + m01_imag_vec * a1_real;

                // Compute new |1⟩ amplitudes
                let b1_real = m10_real_vec * a0_real - m10_imag_vec * a0_imag +
                              m11_real_vec * a1_real - m11_imag_vec * a1_imag;

                let b1_imag = m10_real_vec * a0_imag + m10_imag_vec * a0_real +
                              m11_real_vec * a1_imag + m11_imag_vec * a1_real;

                // Store results
                let b0_real_arr: [f64; 4] = b0_real.into();
                let b0_imag_arr: [f64; 4] = b0_imag.into();
                let b1_real_arr: [f64; 4] = b1_real.into();
                let b1_imag_arr: [f64; 4] = b1_imag.into();

                for j in 0..4 {
                    new_state[zero_indices[j]] = Complex64::new(b0_real_arr[j], b0_imag_arr[j]);
                    new_state[one_indices[j]] = Complex64::new(b1_real_arr[j], b1_imag_arr[j]);
                }
            }

            // Process remaining elements (if any)
            let remaining_offset = (1 << target) / 4 * 4;
            for i in remaining_offset..(1 << target) {
                let zero_idx = zero_offset + i;
                let one_idx = one_offset + i;

                // Apply the gate conventionally for these elements
                let a0 = state[zero_idx];
                let a1 = state[one_idx];

                new_state[zero_idx] = matrix[0] * a0 + matrix[1] * a1;
                new_state[one_idx] = matrix[2] * a0 + matrix[3] * a1;
            }
        }

        // Process any remaining elements at the end
        let processed_elements = num_chunks * chunk_size;
        for i in processed_elements..dim {
            let bit_val = (i >> target) & 1;
            let paired_idx = i ^ (1 << target);

            if bit_val == 0 {
                new_state[i] = matrix[0] * state[i] + matrix[1] * state[paired_idx];
                new_state[paired_idx] = matrix[2] * state[i] + matrix[3] * state[paired_idx];
            }
        }

        state.copy_from_slice(&new_state);
    }
    
    /// Apply a CNOT gate using SIMD operations
    ///
    /// This implementation is optimized for controlled operations by leveraging
    /// SIMD instructions to process multiple amplitudes at once.
    #[cfg(feature = "simd")]
    pub fn apply_cnot_simd(
        state: &mut [Complex64],
        control: usize,
        target: usize,
        num_qubits: usize
    ) {
        let dim = state.len();
        let mut new_state = vec![Complex64::new(0.0, 0.0); dim];

        // Determine chunk size and alignment
        // For CNOT gates, we need to arrange chunks that contain
        // complete sets of control and target qubit patterns
        let chunk_size = 1 << (control.max(target) + 1);
        let num_chunks = dim / chunk_size;

        // Process the statevector in chunks
        for chunk_idx in 0..num_chunks {
            let base_idx = chunk_idx * chunk_size;

            // For each chunk, we'll process four elements at a time using SIMD
            let vec_size = 4;
            let chunk_iterations = chunk_size / vec_size;

            for i in 0..chunk_iterations {
                let idx_base = base_idx + i * vec_size;

                // Determine which indices need to be swapped based on control bit
                let mut swap_indices = [0usize; 4];
                let mut needs_swap = [false; 4];

                for j in 0..vec_size {
                    let idx = idx_base + j;
                    let control_bit = (idx >> control) & 1;
                    needs_swap[j] = control_bit == 1;
                    swap_indices[j] = if needs_swap[j] { idx ^ (1 << target) } else { idx };
                }

                // Process the current batch
                for j in 0..vec_size {
                    let idx = idx_base + j;
                    if needs_swap[j] {
                        new_state[idx] = state[swap_indices[j]];
                        new_state[swap_indices[j]] = state[idx];
                    } else {
                        new_state[idx] = state[idx];
                    }
                }
            }

            // Handle any remaining elements in this chunk
            let remaining_start = base_idx + chunk_iterations * vec_size;
            for idx in remaining_start..(base_idx + chunk_size) {
                let control_bit = (idx >> control) & 1;
                if control_bit == 1 {
                    let swap_idx = idx ^ (1 << target);
                    new_state[idx] = state[swap_idx];
                    new_state[swap_idx] = state[idx];
                } else {
                    new_state[idx] = state[idx];
                }
            }
        }

        // Handle any remaining elements at the end
        let processed_elements = num_chunks * chunk_size;
        for idx in processed_elements..dim {
            let control_bit = (idx >> control) & 1;
            if control_bit == 1 {
                let swap_idx = idx ^ (1 << target);
                new_state[idx] = state[swap_idx];
            } else {
                new_state[idx] = state[idx];
            }
        }

        state.copy_from_slice(&new_state);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::f64::consts::FRAC_1_SQRT_2;
    
    #[test]
    fn test_optimized_state_vector_init() {
        let sv = OptimizedStateVector::new(2);
        assert_eq!(sv.num_qubits(), 2);
        assert_eq!(sv.dimension(), 4);
        
        // Initial state should be |00>
        assert_eq!(sv.state()[0], Complex64::new(1.0, 0.0));
        assert_eq!(sv.state()[1], Complex64::new(0.0, 0.0));
        assert_eq!(sv.state()[2], Complex64::new(0.0, 0.0));
        assert_eq!(sv.state()[3], Complex64::new(0.0, 0.0));
    }
    
    #[test]
    fn test_hadamard_gate() {
        // Hadamard matrix
        let h_matrix = [
            Complex64::new(FRAC_1_SQRT_2, 0.0),
            Complex64::new(FRAC_1_SQRT_2, 0.0),
            Complex64::new(FRAC_1_SQRT_2, 0.0),
            Complex64::new(-FRAC_1_SQRT_2, 0.0),
        ];
        
        // Apply H to the 0th qubit of |00>
        let mut sv = OptimizedStateVector::new(2);
        sv.apply_single_qubit_gate(&h_matrix, 0);
        
        // Result should be |00> + |10> / sqrt(2)
        assert_eq!(sv.state()[0], Complex64::new(FRAC_1_SQRT_2, 0.0));
        assert_eq!(sv.state()[1], Complex64::new(0.0, 0.0));
        assert_eq!(sv.state()[2], Complex64::new(FRAC_1_SQRT_2, 0.0));
        assert_eq!(sv.state()[3], Complex64::new(0.0, 0.0));
        
        // Apply H to the 1st qubit
        sv.apply_single_qubit_gate(&h_matrix, 1);
        
        // Result should be (|00> + |01> + |10> + |11>) / 2
        assert_eq!(sv.state()[0], Complex64::new(0.5, 0.0));
        assert_eq!(sv.state()[1], Complex64::new(0.5, 0.0));
        assert_eq!(sv.state()[2], Complex64::new(0.5, 0.0));
        assert_eq!(sv.state()[3], Complex64::new(-0.5, 0.0));
    }
    
    #[test]
    fn test_cnot_gate() {
        // Set up state |+0> = (|00> + |10>) / sqrt(2)
        let mut sv = OptimizedStateVector::new(2);
        
        // Hadamard on qubit 0
        let h_matrix = [
            Complex64::new(FRAC_1_SQRT_2, 0.0),
            Complex64::new(FRAC_1_SQRT_2, 0.0),
            Complex64::new(FRAC_1_SQRT_2, 0.0),
            Complex64::new(-FRAC_1_SQRT_2, 0.0),
        ];
        sv.apply_single_qubit_gate(&h_matrix, 0);
        
        // Apply CNOT
        sv.apply_cnot(0, 1);
        
        // Result should be (|00> + |11>) / sqrt(2) = Bell state
        assert_eq!(sv.state()[0], Complex64::new(FRAC_1_SQRT_2, 0.0));
        assert_eq!(sv.state()[1], Complex64::new(0.0, 0.0));
        assert_eq!(sv.state()[2], Complex64::new(0.0, 0.0));
        assert_eq!(sv.state()[3], Complex64::new(FRAC_1_SQRT_2, 0.0));
    }
}