quantrs2-core 0.1.3

//! Quantum Counting and Amplitude Estimation
//!
//! This module implements quantum counting and amplitude estimation algorithms,
//! which are key components for many quantum algorithms including Shor's algorithm
//! and quantum database search.
//!
//! TODO: The current implementations are simplified versions. Full implementations
//! would require:
//! - Proper controlled unitary implementations
//! - Full QFT and inverse QFT
//! - Better phase extraction from measurement results
//! - Integration with circuit builder for more complex operations

use scirs2_core::ndarray::{Array1, Array2};
use scirs2_core::Complex64;
use std::f64::consts::PI;

/// Quantum Phase Estimation (QPE) algorithm
///
/// Estimates the phase φ in the eigenvalue e^(2πiφ) of a unitary operator U
pub struct QuantumPhaseEstimation {
    /// Number of precision bits
    precision_bits: usize,
    /// The unitary operator U
    unitary: Array2<Complex64>,
    /// Number of target qubits
    target_qubits: usize,
}

impl QuantumPhaseEstimation {
    /// Create a new QPE instance
    pub fn new(precision_bits: usize, unitary: Array2<Complex64>) -> Self {
        let target_qubits = (unitary.shape()[0] as f64).log2() as usize;

        Self {
            precision_bits,
            unitary,
            target_qubits,
        }
    }

    /// Apply controlled-U^(2^k) operation
    fn apply_controlled_u_power(&self, state: &mut [Complex64], control: usize, k: usize) {
        let power = 1 << k;
        let n = self.target_qubits;
        let dim = 1 << n;

        // Build U^power by repeated squaring
        let mut u_power = Array2::eye(dim);
        let mut temp = self.unitary.clone();
        let mut p = power;

        while p > 0 {
            if p & 1 == 1 {
                u_power = u_power.dot(&temp);
            }
            temp = temp.dot(&temp);
            p >>= 1;
        }

        // Apply controlled operation.
        // We iterate over each unique "precision register" configuration; for
        // those with the control qubit set, we apply U^power to the target
        // register.  We use only the "canonical" representative (target = 0)
        // to avoid applying the operation more than once per configuration.
        let total_qubits = self.precision_bits + self.target_qubits;
        let precision_dim = 1 << self.precision_bits;

        for prec in 0..precision_dim {
            // Check if the control qubit is |1⟩ within the precision register.
            // Precision qubits occupy the HIGH bits; qubit `control` corresponds
            // to bit (precision_bits - control - 1) inside `prec`.
            if (prec >> (self.precision_bits - control - 1)) & 1 == 1 {
                // Base index: precision register `prec` in the high bits, target = 0
                let base_idx = prec << n; // precision in MSBs, target bits = 0

                // Read all 2^n amplitudes for this precision configuration
                let mut amplitudes = vec![Complex64::new(0.0, 0.0); dim];
                for i in 0..dim {
                    amplitudes[i] = state[base_idx | i];
                }

                // Apply U^power to the target register amplitudes
                let result = u_power.dot(&Array1::from(amplitudes));

                // Write back
                for i in 0..dim {
                    state[base_idx | i] = result[i];
                }
            }
        }
    }

    /// Apply inverse QFT to precision qubits
    fn apply_inverse_qft(&self, state: &mut [Complex64]) {
        let n = self.precision_bits;
        let total_qubits = n + self.target_qubits;

        // Implement inverse QFT on the first n qubits
        for j in (0..n).rev() {
            // Apply Hadamard to qubit j
            self.apply_hadamard(state, j, total_qubits);

            // Apply controlled phase rotations
            for k in (0..j).rev() {
                let angle = -PI / (1 << (j - k)) as f64;
                self.apply_controlled_phase(state, k, j, angle, total_qubits);
            }
        }

        // Swap qubits to reverse order
        for i in 0..n / 2 {
            self.swap_qubits(state, i, n - 1 - i, total_qubits);
        }
    }

    /// Apply Hadamard gate to a specific qubit
    fn apply_hadamard(&self, state: &mut [Complex64], qubit: usize, total_qubits: usize) {
        let h = 1.0 / std::f64::consts::SQRT_2;
        let dim = 1 << total_qubits;

        for i in 0..dim {
            if (i >> (total_qubits - qubit - 1)) & 1 == 0 {
                let j = i | (1 << (total_qubits - qubit - 1));
                let a = state[i];
                let b = state[j];
                state[i] = h * (a + b);
                state[j] = h * (a - b);
            }
        }
    }

    /// Apply controlled phase rotation
    fn apply_controlled_phase(
        &self,
        state: &mut [Complex64],
        control: usize,
        target: usize,
        angle: f64,
        total_qubits: usize,
    ) {
        let phase = Complex64::new(angle.cos(), angle.sin());

        for (i, amp) in state.iter_mut().enumerate() {
            let control_bit = (i >> (total_qubits - control - 1)) & 1;
            let target_bit = (i >> (total_qubits - target - 1)) & 1;

            if control_bit == 1 && target_bit == 1 {
                *amp *= phase;
            }
        }
    }

    /// Swap two qubits
    fn swap_qubits(&self, state: &mut [Complex64], q1: usize, q2: usize, total_qubits: usize) {
        let dim = 1 << total_qubits;

        for i in 0..dim {
            let bit1 = (i >> (total_qubits - q1 - 1)) & 1;
            let bit2 = (i >> (total_qubits - q2 - 1)) & 1;

            if bit1 != bit2 {
                let j = i ^ (1 << (total_qubits - q1 - 1)) ^ (1 << (total_qubits - q2 - 1));
                if i < j {
                    state.swap(i, j);
                }
            }
        }
    }

    /// Run the QPE algorithm
    pub fn estimate_phase(&self, eigenstate: Vec<Complex64>) -> f64 {
        let total_qubits = self.precision_bits + self.target_qubits;
        let state_dim = 1 << total_qubits;
        let mut state = vec![Complex64::new(0.0, 0.0); state_dim];

        // Initialize precision qubits to |0⟩ and target qubits to eigenstate
        for i in 0..(1 << self.target_qubits) {
            if i < eigenstate.len() {
                state[i] = eigenstate[i];
            }
        }

        // Apply Hadamard to all precision qubits
        for j in 0..self.precision_bits {
            self.apply_hadamard(&mut state, j, total_qubits);
        }

        // Apply controlled-U^{2^j} operations.
        // Standard QPE: ancilla qubit j (j = 0 is the first/MSB ancilla)
        // controls U^{2^(precision_bits - 1 - j)} so that after inverse QFT
        // the MSB carries the most significant bit of the phase.
        // Equivalently, iterating j from (precision_bits-1) down to 0:
        //   - ancilla qubit 0 controls U^{2^(n-1)}
        //   - ancilla qubit 1 controls U^{2^(n-2)}
        //   - ...
        //   - ancilla qubit n-1 controls U^{2^0} = U
        for j in 0..self.precision_bits {
            // power_k is the exponent for the j-th controlled-U application
            let power_k = j; // U^{2^j} controlled on ancilla qubit (precision_bits - 1 - j)
            let control_qubit = self.precision_bits - 1 - j;
            self.apply_controlled_u_power(&mut state, control_qubit, power_k);
        }

        // Apply inverse QFT
        self.apply_inverse_qft(&mut state);

        // Measure precision qubits
        let mut max_prob = 0.0;
        let mut measured_value = 0;

        for (i, amp) in state.iter().enumerate() {
            let precision_bits_value = i >> self.target_qubits;
            let prob = amp.norm_sqr();

            if prob > max_prob {
                max_prob = prob;
                measured_value = precision_bits_value;
            }
        }

        // Convert to phase estimate
        measured_value as f64 / (1 << self.precision_bits) as f64
    }
}

/// Quantum Counting algorithm
///
/// Counts the number of solutions to a search problem
pub struct QuantumCounting {
    /// Number of items in the search space
    pub n_items: usize,
    /// Number of precision bits for counting
    pub precision_bits: usize,
    /// Oracle function that marks solutions
    pub oracle: Box<dyn Fn(usize) -> bool>,
}

impl QuantumCounting {
    /// Create a new quantum counting instance
    pub fn new(n_items: usize, precision_bits: usize, oracle: Box<dyn Fn(usize) -> bool>) -> Self {
        Self {
            n_items,
            precision_bits,
            oracle,
        }
    }

    /// Build the Grover operator
    fn build_grover_operator(&self) -> Array2<Complex64> {
        let n = self.n_items;
        let mut grover = Array2::zeros((n, n));

        // Oracle: flip phase of marked items
        for i in 0..n {
            if (self.oracle)(i) {
                grover[[i, i]] = Complex64::new(-1.0, 0.0);
            } else {
                grover[[i, i]] = Complex64::new(1.0, 0.0);
            }
        }

        // Diffusion operator: 2|s⟩⟨s| - I
        let s_amplitude = 1.0 / (n as f64).sqrt();
        let diffusion =
            Array2::from_elem((n, n), Complex64::new(2.0 * s_amplitude * s_amplitude, 0.0))
                - Array2::<Complex64>::eye(n);

        // Grover operator = -Diffusion × Oracle
        -diffusion.dot(&grover)
    }

    /// Count the number of solutions
    pub fn count(&self) -> f64 {
        // Build Grover operator
        let grover = self.build_grover_operator();

        // Use QPE to estimate the phase
        let qpe = QuantumPhaseEstimation::new(self.precision_bits, grover);

        // Prepare uniform superposition as eigenstate
        let n = self.n_items;
        let amplitude = Complex64::new(1.0 / (n as f64).sqrt(), 0.0);
        let eigenstate = vec![amplitude; n];

        // Estimate phase
        let phase = qpe.estimate_phase(eigenstate);

        // Convert phase to count
        // For Grover operator, eigenvalues are e^(±2iθ) where sin²(θ) = M/N
        let theta = phase * PI;
        let sin_theta = theta.sin();
        sin_theta * sin_theta * n as f64
    }
}

/// Quantum Amplitude Estimation
///
/// Estimates the amplitude of marked states in a superposition
pub struct QuantumAmplitudeEstimation {
    /// State preparation operator A
    pub state_prep: Array2<Complex64>,
    /// Oracle that identifies good states
    pub oracle: Array2<Complex64>,
    /// Number of precision bits
    pub precision_bits: usize,
}

impl QuantumAmplitudeEstimation {
    /// Create a new amplitude estimation instance
    pub const fn new(
        state_prep: Array2<Complex64>,
        oracle: Array2<Complex64>,
        precision_bits: usize,
    ) -> Self {
        Self {
            state_prep,
            oracle,
            precision_bits,
        }
    }

    /// Build the Q operator for amplitude estimation
    fn build_q_operator(&self) -> Array2<Complex64> {
        let n = self.state_prep.shape()[0];
        let identity = Array2::<Complex64>::eye(n);

        // Reflection about good states: I - 2P where P projects onto good states
        let reflection_good = &identity - &self.oracle * 2.0;

        // Reflection about initial state: 2A|0⟩⟨0|A† - I
        let zero_state = Array1::zeros(n);
        let mut zero_state_vec = zero_state.to_vec();
        zero_state_vec[0] = Complex64::new(1.0, 0.0);

        let initial = self.state_prep.dot(&Array1::from(zero_state_vec));
        let mut reflection_initial = Array2::zeros((n, n));

        for i in 0..n {
            for j in 0..n {
                reflection_initial[[i, j]] = 2.0 * initial[i] * initial[j].conj();
            }
        }
        reflection_initial -= &identity;

        // Q = -reflection_initial × reflection_good
        -reflection_initial.dot(&reflection_good)
    }

    /// Estimate the amplitude
    pub fn estimate(&self) -> f64 {
        // Build Q operator
        let q_operator = self.build_q_operator();

        // Use QPE to find eigenphase
        let qpe = QuantumPhaseEstimation::new(self.precision_bits, q_operator);

        // Prepare initial state A|0⟩
        let n = self.state_prep.shape()[0];
        let mut zero_state = vec![Complex64::new(0.0, 0.0); n];
        zero_state[0] = Complex64::new(1.0, 0.0);
        let initial_state = self.state_prep.dot(&Array1::from(zero_state));

        // Estimate phase
        let phase = qpe.estimate_phase(initial_state.to_vec());

        // Convert phase to amplitude
        // For Q operator, eigenvalues are e^(±2iθ) where sin(θ) = amplitude
        let theta = phase * PI;
        theta.sin().abs()
    }
}

/// Example: Count solutions to a simple search problem
pub fn quantum_counting_example() {
    println!("Quantum Counting Example");
    println!("=======================");

    // Count numbers divisible by 3 in range 0-15
    let oracle = Box::new(|x: usize| x % 3 == 0 && x > 0);

    let counter = QuantumCounting::new(16, 4, oracle);
    let count = counter.count();

    println!("Counting numbers divisible by 3 in range 1-15:");
    println!("Estimated count: {count:.1}");
    println!("Actual count: 5 (3, 6, 9, 12, 15)");
    println!("Error: {:.1}", (count - 5.0).abs());
}

/// Example: Estimate amplitude of marked states
pub fn amplitude_estimation_example() {
    println!("\nAmplitude Estimation Example");
    println!("============================");

    // Create a simple state preparation that creates equal superposition
    let n = 8;
    let state_prep = Array2::from_elem((n, n), Complex64::new(1.0 / (n as f64).sqrt(), 0.0));

    // Oracle marks states 2 and 5
    let mut oracle = Array2::zeros((n, n));
    oracle[[2, 2]] = Complex64::new(1.0, 0.0);
    oracle[[5, 5]] = Complex64::new(1.0, 0.0);

    let qae = QuantumAmplitudeEstimation::new(state_prep, oracle, 4);
    let amplitude = qae.estimate();

    println!("Estimating amplitude of marked states (2 and 5) in uniform superposition:");
    println!("Estimated amplitude: {amplitude:.3}");
    println!("Actual amplitude: {:.3}", (2.0 / n as f64).sqrt());
    println!("Error: {:.3}", (amplitude - (2.0 / n as f64).sqrt()).abs());
}

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

    #[test]
    fn test_phase_estimation_basic() {
        // U = diag(1, e^{iφ}) with φ = π/4.
        // The eigenvalue of |1⟩ is e^{iπ/4}; the true phase is φ/(2π) = 1/8.
        //
        // Note: the current QPE implementation is a matrix-simulation approach
        // that gives a numerically valid phase estimate in [0, 1].  The exact
        // value depends on qubit ordering conventions; we verify structural
        // correctness (valid range) and that the result is a multiple of 1/2^k.
        let phase = PI / 4.0;
        let u = Array2::from_shape_vec(
            (2, 2),
            vec![
                Complex64::new(1.0, 0.0),
                Complex64::new(0.0, 0.0),
                Complex64::new(0.0, 0.0),
                Complex64::new(phase.cos(), phase.sin()),
            ],
        )
        .expect("2x2 matrix from 4-element vector should succeed");

        let precision_bits = 4usize;
        let qpe = QuantumPhaseEstimation::new(precision_bits, u);

        // Test with eigenstate |1⟩
        let eigenstate = vec![Complex64::new(0.0, 0.0), Complex64::new(1.0, 0.0)];
        let estimated = qpe.estimate_phase(eigenstate);

        // Phase must lie in [0, 1]
        assert!(
            (0.0..=1.0).contains(&estimated),
            "estimated phase {estimated} is outside [0, 1]"
        );

        // The estimate must be a multiple of 1/(2^precision_bits)
        let grid = 1.0 / (1u64 << precision_bits) as f64;
        let residual = (estimated / grid).round() * grid - estimated;
        assert!(
            residual.abs() < 1e-9,
            "estimated {estimated} is not on the {precision_bits}-bit phase grid"
        );

        // The true phase φ/(2π) = 1/8.  The estimated value should correspond
        // to one of the two conjugate eigenphases: φ/(2π) or 1 - φ/(2π).
        let true_phase = phase / (2.0 * PI);
        let conjugate_phase = 1.0 - true_phase;
        // Allow one grid step of slack on either eigenphase
        let slack = grid + 1e-9;
        let near_true = (estimated - true_phase).abs() <= slack;
        let near_conj = (estimated - conjugate_phase).abs() <= slack;
        assert!(
            near_true || near_conj,
            "QPE estimate {estimated:.6} is not near true phase {true_phase:.6} \
             or conjugate {conjugate_phase:.6} (tolerance {slack:.6})"
        );
    }

    #[test]
    fn test_quantum_counting_simple() {
        // Search space of 4 items; oracle marks exactly item 2 → M = 1.
        // QuantumCounting.count() returns N·sin²(π·θ_estimated) where θ is
        // the QPE phase estimate for the Grover operator.
        let oracle = Box::new(|x: usize| x == 2);
        let counter = QuantumCounting::new(4, 4, oracle);
        let count = counter.count();

        // The count must be non-negative and bounded by the search space size
        assert!(count >= 0.0, "count {count} must be non-negative");
        assert!(count <= 4.0 + 1e-6, "count {count} must not exceed N=4");
    }
}