qzkp 0.1.0

//! Single-qubit state-vector quantum simulator with noise models.
//!
//! Provides [`QuantumState`] for individual qubit manipulation and
//! [`QuantumCircuit`] for multi-qubit register operations, including
//! depolarizing noise and amplitude damping channels.

use num_complex::Complex64;
use rand::Rng;

const SQRT_2: f64 = std::f64::consts::SQRT_2;

/// A single-qubit pure state represented as two complex amplitudes.
///
/// Initialized to |0> by default: `amplitude_0 = 1`, `amplitude_1 = 0`.
#[derive(Clone, Debug)]
pub struct QuantumState {
    /// Amplitude of the |0> basis state.
    pub amplitude_0: Complex64,
    /// Amplitude of the |1> basis state.
    pub amplitude_1: Complex64,
}

impl QuantumState {
    /// Create a new qubit in the |0> state.
    pub fn new() -> Self {
        QuantumState {
            amplitude_0: Complex64::new(1.0, 0.0),
            amplitude_1: Complex64::new(0.0, 0.0),
        }
    }

    /// Apply the Pauli-X (NOT) gate: swaps |0> and |1> amplitudes.
    pub fn apply_x(&mut self) {
        std::mem::swap(&mut self.amplitude_0, &mut self.amplitude_1);
    }

    /// Apply the Hadamard gate: maps |0> to |+> and |1> to |->.
    pub fn apply_h(&mut self) {
        let a0 = self.amplitude_0;
        let a1 = self.amplitude_1;
        self.amplitude_0 = (a0 + a1) / SQRT_2;
        self.amplitude_1 = (a0 - a1) / SQRT_2;
    }

    /// Perform a projective measurement in the computational basis.
    ///
    /// Returns `0` or `1` with probabilities |amplitude_0|^2 and |amplitude_1|^2.
    pub fn measure(&self, rng: &mut impl Rng) -> u8 {
        let prob_0 = self.amplitude_0.norm_sqr();
        if rng.gen::<f64>() < prob_0 {
            0
        } else {
            1
        }
    }

    /// Apply a depolarizing noise channel with the given `error_rate`.
    ///
    /// With probability `error_rate`, one of {X, Y, Z} is applied uniformly at random.
    pub fn apply_depolarizing_noise(&mut self, error_rate: f64, rng: &mut impl Rng) {
        if rng.gen::<f64>() < error_rate {
            let noise_type: f64 = rng.gen();
            if noise_type < 0.33 {
                self.apply_x();
            } else if noise_type < 0.67 {
                // Y gate
                let temp_0 = self.amplitude_0;
                self.amplitude_0 = Complex64::new(0.0, 1.0) * self.amplitude_1;
                self.amplitude_1 = Complex64::new(0.0, -1.0) * temp_0;
            } else {
                // Z gate
                self.amplitude_1 = -self.amplitude_1;
            }
        }
    }

    /// Apply an amplitude damping channel with parameter `gamma`.
    ///
    /// Models energy dissipation: |1> decays toward |0> with probability
    /// proportional to `gamma * |amplitude_1|^2`.
    pub fn apply_amplitude_damping(&mut self, gamma: f64, rng: &mut impl Rng) {
        let prob_decay = gamma * self.amplitude_1.norm_sqr();

        if rng.gen::<f64>() < prob_decay {
            self.amplitude_0 += self.amplitude_1;
            self.amplitude_1 = Complex64::new(0.0, 0.0);
        } else {
            let scale = (1.0 - gamma).sqrt();
            self.amplitude_1 *= scale;
        }

        // Re-normalize
        let norm = (self.amplitude_0.norm_sqr() + self.amplitude_1.norm_sqr()).sqrt();
        if norm > 0.0 {
            self.amplitude_0 /= norm;
            self.amplitude_1 /= norm;
        }
    }
}

impl Default for QuantumState {
    fn default() -> Self {
        Self::new()
    }
}

/// A multi-qubit register of independent single-qubit states.
///
/// Each qubit is simulated independently (no entanglement).
/// A global `noise_level` is applied via [`QuantumCircuit::apply_noise`].
pub struct QuantumCircuit {
    /// The individual qubit states.
    pub states: Vec<QuantumState>,
    /// Depolarizing noise level applied to all qubits.
    pub noise_level: f64,
}

impl QuantumCircuit {
    /// Create a new circuit with `n_qubits` qubits, all initialized to |0>.
    pub fn new(n_qubits: usize, noise_level: f64) -> Self {
        QuantumCircuit {
            states: vec![QuantumState::new(); n_qubits],
            noise_level,
        }
    }

    /// Apply the Pauli-X gate to the specified qubit.
    pub fn x(&mut self, qubit: usize) {
        if qubit < self.states.len() {
            self.states[qubit].apply_x();
        }
    }

    /// Apply the Hadamard gate to the specified qubit.
    pub fn h(&mut self, qubit: usize) {
        if qubit < self.states.len() {
            self.states[qubit].apply_h();
        }
    }

    /// Apply depolarizing noise to all qubits at the circuit's noise level.
    pub fn apply_noise(&mut self, rng: &mut impl Rng) {
        for state in &mut self.states {
            state.apply_depolarizing_noise(self.noise_level, rng);
        }
    }

    /// Apply amplitude damping noise to all qubits with the given `gamma`.
    pub fn apply_amplitude_damping_noise(&mut self, gamma: f64, rng: &mut impl Rng) {
        for state in &mut self.states {
            state.apply_amplitude_damping(gamma, rng);
        }
    }

    /// Measure a single qubit in the computational basis.
    pub fn measure(&self, qubit: usize, rng: &mut impl Rng) -> u8 {
        if qubit < self.states.len() {
            self.states[qubit].measure(rng)
        } else {
            0
        }
    }

    /// Measure all qubits and return the results as a vector.
    pub fn measure_all(&self, rng: &mut impl Rng) -> Vec<u8> {
        self.states.iter().map(|state| state.measure(rng)).collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rand::rngs::StdRng;
    use rand::SeedableRng;

    fn seeded_rng() -> StdRng {
        StdRng::seed_from_u64(42)
    }

    #[test]
    fn qubit_starts_in_zero_state() {
        let q = QuantumState::new();
        assert!((q.amplitude_0.re - 1.0).abs() < 1e-10);
        assert!(q.amplitude_1.norm_sqr() < 1e-10);
    }

    #[test]
    fn x_gate_flips_to_one() {
        let mut q = QuantumState::new();
        q.apply_x();
        assert!(q.amplitude_0.norm_sqr() < 1e-10);
        assert!((q.amplitude_1.re - 1.0).abs() < 1e-10);
    }

    #[test]
    fn double_x_is_identity() {
        let mut q = QuantumState::new();
        q.apply_x();
        q.apply_x();
        assert!((q.amplitude_0.re - 1.0).abs() < 1e-10);
        assert!(q.amplitude_1.norm_sqr() < 1e-10);
    }

    #[test]
    fn hadamard_creates_superposition() {
        let mut q = QuantumState::new();
        q.apply_h();
        let expected = 1.0 / SQRT_2;
        assert!((q.amplitude_0.re - expected).abs() < 1e-10);
        assert!((q.amplitude_1.re - expected).abs() < 1e-10);
    }

    #[test]
    fn double_hadamard_is_identity() {
        let mut q = QuantumState::new();
        q.apply_h();
        q.apply_h();
        assert!((q.amplitude_0.re - 1.0).abs() < 1e-10);
        assert!(q.amplitude_1.norm_sqr() < 1e-10);
    }

    #[test]
    fn measurement_statistics_zero_state() {
        let q = QuantumState::new();
        let mut rng = seeded_rng();
        let ones: usize = (0..1000).map(|_| q.measure(&mut rng) as usize).sum();
        assert_eq!(ones, 0, "|0> should always measure 0");
    }

    #[test]
    fn measurement_statistics_superposition() {
        let mut q = QuantumState::new();
        q.apply_h();
        let mut rng = seeded_rng();
        let ones: usize = (0..10000).map(|_| q.measure(&mut rng) as usize).sum();
        let ratio = ones as f64 / 10000.0;
        assert!(
            (ratio - 0.5).abs() < 0.05,
            "H|0> should give ~50/50, got {}",
            ratio
        );
    }

    #[test]
    fn circuit_measure_all() {
        let qc = QuantumCircuit::new(5, 0.0);
        let mut rng = seeded_rng();
        let results = qc.measure_all(&mut rng);
        assert_eq!(results.len(), 5);
        assert!(
            results.iter().all(|&b| b == 0),
            "all-|0> should measure all 0"
        );
    }

    #[test]
    fn circuit_x_and_measure() {
        let mut qc = QuantumCircuit::new(3, 0.0);
        qc.x(1);
        let mut rng = seeded_rng();
        assert_eq!(qc.measure(0, &mut rng), 0);
        assert_eq!(qc.measure(1, &mut rng), 1);
        assert_eq!(qc.measure(2, &mut rng), 0);
    }
}