quantrs2-core 0.1.3

// Advanced Error Mitigation Techniques
// Implements cutting-edge error mitigation methods including:
// - Virtual Distillation
// - Clifford Data Regression (CDR)
// - Probabilistic Error Cancellation (PEC) with ML optimization
// - Symmetry Verification
// - Quantum Subspace Expansion

use crate::error::QuantRS2Error;
use scirs2_core::ndarray::ndarray_linalg::Solve;
use scirs2_core::ndarray::{Array1, Array2};
use scirs2_core::random::prelude::*;
use scirs2_core::Complex64;
use std::collections::HashMap;

/// Virtual Distillation error mitigation technique
///
/// Virtual Distillation uses multiple copies of a quantum state to
/// create a "purified" version that suppresses errors. This technique
/// is particularly effective for variational algorithms.
///
/// Reference: Koczor, B. (2021). "Exponential Error Suppression for Near-Term Quantum Devices"
#[derive(Debug, Clone)]
pub struct VirtualDistillation {
    /// Number of state copies to use
    pub num_copies: usize,
    /// Permutation group for symmetrization
    pub use_symmetrization: bool,
    /// Cache for computed permutations
    permutation_cache: HashMap<usize, Vec<Vec<usize>>>,
}

impl VirtualDistillation {
    /// Create a new Virtual Distillation instance
    pub fn new(num_copies: usize) -> Self {
        Self {
            num_copies,
            use_symmetrization: true,
            permutation_cache: HashMap::new(),
        }
    }

    /// Apply virtual distillation to measurement results
    ///
    /// # Arguments
    /// * `raw_results` - Raw measurement outcomes from multiple circuit runs
    /// * `observable` - The observable to measure (Pauli string representation)
    ///
    /// # Returns
    /// Mitigated expectation value with reduced error
    pub fn mitigate_expectation(
        &mut self,
        raw_results: &[f64],
        observable: &Array1<Complex64>,
    ) -> Result<f64, QuantRS2Error> {
        if raw_results.len() < self.num_copies {
            return Err(QuantRS2Error::InvalidInput(format!(
                "Need at least {} measurements, got {}",
                self.num_copies,
                raw_results.len()
            )));
        }

        // Group measurements into copies
        let chunks: Vec<&[f64]> = raw_results
            .chunks(raw_results.len() / self.num_copies)
            .collect();

        // Compute virtual distillation estimator
        let mut mitigated_value = 0.0;
        let mut weight_sum = 0.0;

        for perm in self.generate_permutations(self.num_copies) {
            let mut product = 1.0;
            for (idx, &chunk_idx) in perm.iter().enumerate() {
                if chunk_idx < chunks.len() && idx < chunks[chunk_idx].len() {
                    product *= chunks[chunk_idx][idx];
                }
            }

            // Symmetrization weight
            let weight = if self.use_symmetrization {
                1.0 / Self::factorial(self.num_copies) as f64
            } else {
                1.0
            };

            mitigated_value += weight * product;
            weight_sum += weight;
        }

        Ok(mitigated_value / weight_sum.max(1e-10))
    }

    /// Generate all permutations of n elements
    fn generate_permutations(&mut self, n: usize) -> Vec<Vec<usize>> {
        if let Some(cached) = self.permutation_cache.get(&n) {
            return cached.clone();
        }

        let perms = Self::permute((0..n).collect());
        self.permutation_cache.insert(n, perms.clone());
        perms
    }

    /// Recursive permutation generation
    fn permute(elements: Vec<usize>) -> Vec<Vec<usize>> {
        if elements.len() <= 1 {
            return vec![elements];
        }

        let mut result = Vec::new();
        for i in 0..elements.len() {
            let mut remaining = elements.clone();
            let current = remaining.remove(i);

            for mut perm in Self::permute(remaining) {
                perm.insert(0, current);
                result.push(perm);
            }
        }
        result
    }

    /// Compute factorial
    const fn factorial(n: usize) -> usize {
        match n {
            0 | 1 => 1,
            _ => {
                let mut result = 1;
                let mut i = 2;
                while i <= n {
                    result *= i;
                    i += 1;
                }
                result
            }
        }
    }

    /// Estimate error suppression factor
    ///
    /// Virtual distillation provides exponential error suppression:
    /// ε_suppressed ≈ ε^n where n is the number of copies
    pub const fn error_suppression_factor(&self, base_error: f64) -> f64 {
        // Using const-compatible exponentiation approximation
        let mut result = 1.0;
        let mut i = 0;
        while i < self.num_copies {
            result *= base_error;
            i += 1;
        }
        result
    }
}

/// Clifford Data Regression (CDR) error mitigation
///
/// CDR learns a noise model from Clifford circuit data and applies
/// it to non-Clifford circuits for improved accuracy.
///
/// Reference: Czarnik et al. (2021). "Error mitigation with Clifford quantum-circuit data"
#[derive(Debug, Clone)]
pub struct CliffordDataRegression {
    /// Number of training Clifford circuits
    pub num_training_circuits: usize,
    /// Polynomial degree for regression
    pub regression_degree: usize,
    /// Learned regression coefficients
    coefficients: Option<Array1<f64>>,
    /// Training data cache
    training_data: Vec<(Array1<f64>, f64)>,
}

impl CliffordDataRegression {
    /// Create a new CDR instance
    pub const fn new(num_training_circuits: usize, regression_degree: usize) -> Self {
        Self {
            num_training_circuits,
            regression_degree,
            coefficients: None,
            training_data: Vec::new(),
        }
    }

    /// Train the CDR model using Clifford circuit data
    ///
    /// # Arguments
    /// * `clifford_noisy` - Noisy expectation values from Clifford circuits
    /// * `clifford_ideal` - Ideal (simulable) expectation values
    pub fn train(
        &mut self,
        clifford_noisy: &[f64],
        clifford_ideal: &[f64],
    ) -> Result<(), QuantRS2Error> {
        if clifford_noisy.len() != clifford_ideal.len() {
            return Err(QuantRS2Error::InvalidInput(
                "Clifford data lengths must match".to_string(),
            ));
        }

        if clifford_noisy.len() < self.regression_degree + 1 {
            return Err(QuantRS2Error::InvalidInput(format!(
                "Need at least {} training points for degree {} regression",
                self.regression_degree + 1,
                self.regression_degree
            )));
        }

        // Build feature matrix for polynomial regression
        let n = clifford_noisy.len();
        let mut features = Array2::<f64>::zeros((n, self.regression_degree + 1));

        for i in 0..n {
            for j in 0..=self.regression_degree {
                features[[i, j]] = clifford_noisy[i].powi(j as i32);
            }
        }

        // Target values (ideal - noisy = correction)
        let targets: Array1<f64> = clifford_ideal
            .iter()
            .zip(clifford_noisy.iter())
            .map(|(ideal, noisy)| ideal - noisy)
            .collect();

        // Solve least squares: features^T * features * coef = features^T * targets
        let ftf = features.t().dot(&features);
        let fty = features.t().dot(&targets);

        // Use SciRS2 linear algebra for solving
        match ftf.solve_into(&fty) {
            Ok(coeffs) => {
                self.coefficients = Some(coeffs);
                Ok(())
            }
            Err(_) => Err(QuantRS2Error::LinalgError(
                "Failed to solve regression problem".to_string(),
            )),
        }
    }

    /// Apply learned correction to non-Clifford circuit results
    pub fn mitigate(&self, noisy_value: f64) -> Result<f64, QuantRS2Error> {
        let coeffs = self.coefficients.as_ref().ok_or_else(|| {
            QuantRS2Error::InvalidOperation("CDR model not trained yet".to_string())
        })?;

        // Compute polynomial correction
        let mut correction = 0.0;
        for (degree, &coeff) in coeffs.iter().enumerate() {
            correction += coeff * noisy_value.powi(degree as i32);
        }

        Ok(noisy_value + correction)
    }

    /// Get model quality metric (R² score)
    pub fn get_r_squared(
        &self,
        test_noisy: &[f64],
        test_ideal: &[f64],
    ) -> Result<f64, QuantRS2Error> {
        if test_noisy.len() != test_ideal.len() {
            return Err(QuantRS2Error::InvalidInput(
                "Test data lengths must match".to_string(),
            ));
        }

        let mut ss_res = 0.0;
        let mut ss_tot = 0.0;
        let mean_ideal: f64 = test_ideal.iter().sum::<f64>() / test_ideal.len() as f64;

        for (noisy, ideal) in test_noisy.iter().zip(test_ideal.iter()) {
            let mitigated = self.mitigate(*noisy)?;
            ss_res += (ideal - mitigated).powi(2);
            ss_tot += (ideal - mean_ideal).powi(2);
        }

        Ok(1.0 - ss_res / ss_tot.max(1e-10))
    }
}

/// Symmetry Verification for error detection
///
/// Exploits known symmetries of the quantum system to detect and
/// flag measurements that violate physical constraints.
#[derive(Debug, Clone)]
pub struct SymmetryVerification {
    /// Symmetry operators (represented as Pauli strings)
    pub symmetry_operators: Vec<Array1<Complex64>>,
    /// Tolerance for symmetry violation
    pub tolerance: f64,
}

impl SymmetryVerification {
    /// Create a new symmetry verification instance
    pub const fn new(symmetry_operators: Vec<Array1<Complex64>>, tolerance: f64) -> Self {
        Self {
            symmetry_operators,
            tolerance,
        }
    }

    /// Check if a measurement satisfies all symmetries
    ///
    /// # Returns
    /// `(is_valid, violations)` where violations lists which symmetries were broken
    pub fn verify_measurement(&self, measurement_state: &Array1<Complex64>) -> (bool, Vec<usize>) {
        let mut violations = Vec::new();

        for (idx, symmetry_op) in self.symmetry_operators.iter().enumerate() {
            // Check if state is eigenstate of symmetry operator
            let expectation = Self::compute_expectation(measurement_state, symmetry_op);

            // For a perfect eigenstate, expectation should be ±1
            if (expectation.abs() - 1.0).abs() > self.tolerance {
                violations.push(idx);
            }
        }

        (violations.is_empty(), violations)
    }

    /// Compute expectation value <ψ|O|ψ>
    fn compute_expectation(state: &Array1<Complex64>, operator: &Array1<Complex64>) -> f64 {
        // Simplified: assuming operator is diagonal in computational basis
        state
            .iter()
            .zip(operator.iter())
            .map(|(s, o)| (s.conj() * s * o).re)
            .sum()
    }

    /// Post-select measurements that satisfy symmetries
    pub fn post_select_measurements(
        &self,
        measurements: &[Array1<Complex64>],
    ) -> Vec<Array1<Complex64>> {
        measurements
            .iter()
            .filter(|m| self.verify_measurement(m).0)
            .cloned()
            .collect()
    }
}

/// Quantum Subspace Expansion for error mitigation
///
/// Expands the computation into a larger Hilbert space that includes
/// error states, then projects back to extract error-mitigated results.
///
/// Reference: McClean et al. (2020). "Decoding quantum errors with subspace expansions"
#[derive(Debug, Clone)]
pub struct QuantumSubspaceExpansion {
    /// Basis states for the expanded subspace
    pub expansion_basis: Vec<Array1<Complex64>>,
    /// Number of qubits
    pub num_qubits: usize,
}

impl QuantumSubspaceExpansion {
    /// Create a new QSE instance
    pub const fn new(num_qubits: usize) -> Self {
        let expansion_basis = Vec::new(); // Will be populated with error-aware basis states
        Self {
            expansion_basis,
            num_qubits,
        }
    }

    /// Generate expansion basis including single-excitation operators
    pub fn generate_excitation_basis(&mut self, num_excitations: usize) {
        let hilbert_dim = 1 << self.num_qubits;
        self.expansion_basis.clear();

        // Ground state
        let mut ground = Array1::<Complex64>::zeros(hilbert_dim);
        ground[0] = Complex64::new(1.0, 0.0);
        self.expansion_basis.push(ground);

        // Single excitations (bit flips)
        for i in 0..self.num_qubits.min(num_excitations) {
            let mut excited = Array1::<Complex64>::zeros(hilbert_dim);
            let excitation_idx = 1 << i;
            excited[excitation_idx] = Complex64::new(1.0, 0.0);
            self.expansion_basis.push(excited);
        }
    }

    /// Compute subspace expansion coefficients
    pub fn compute_coefficients(
        &self,
        noisy_state: &Array1<Complex64>,
    ) -> Result<Array1<f64>, QuantRS2Error> {
        let n = self.expansion_basis.len();
        if n == 0 {
            return Err(QuantRS2Error::InvalidOperation(
                "Expansion basis not initialized".to_string(),
            ));
        }

        // Build overlap matrix S_ij = <φ_i|φ_j>
        let mut overlap = Array2::<Complex64>::zeros((n, n));
        for i in 0..n {
            for j in 0..n {
                overlap[[i, j]] = self.expansion_basis[i]
                    .iter()
                    .zip(self.expansion_basis[j].iter())
                    .map(|(a, b)| a.conj() * b)
                    .sum();
            }
        }

        // Build state overlap vector b_i = <φ_i|ψ_noisy>
        let state_overlap: Array1<Complex64> = self
            .expansion_basis
            .iter()
            .map(|basis_state| {
                basis_state
                    .iter()
                    .zip(noisy_state.iter())
                    .map(|(a, b)| a.conj() * b)
                    .sum()
            })
            .collect();

        // Convert to real arrays for solving (assuming Hermitian)
        let overlap_real: Array2<f64> = overlap.map(|c| c.re);
        let state_overlap_real: Array1<f64> = state_overlap.map(|c| c.re);

        // Solve S * c = b for expansion coefficients
        match overlap_real.solve_into(&state_overlap_real) {
            Ok(coeffs) => Ok(coeffs),
            Err(_) => Err(QuantRS2Error::LinalgError(
                "Failed to compute subspace coefficients".to_string(),
            )),
        }
    }

    /// Reconstruct mitigated state from subspace expansion
    pub fn reconstruct_state(
        &self,
        coefficients: &Array1<f64>,
    ) -> Result<Array1<Complex64>, QuantRS2Error> {
        if coefficients.len() != self.expansion_basis.len() {
            return Err(QuantRS2Error::InvalidInput(
                "Coefficient count must match basis size".to_string(),
            ));
        }

        let mut mitigated_state = Array1::<Complex64>::zeros(self.expansion_basis[0].len());

        for (coeff, basis_state) in coefficients.iter().zip(self.expansion_basis.iter()) {
            mitigated_state = mitigated_state + basis_state * Complex64::new(*coeff, 0.0);
        }

        // Normalize
        let norm: f64 = mitigated_state
            .iter()
            .map(|c| (c.conj() * c).re)
            .sum::<f64>()
            .sqrt();

        if norm > 1e-10 {
            mitigated_state /= Complex64::new(norm, 0.0);
        }

        Ok(mitigated_state)
    }
}

/// Combined error mitigation strategy using multiple techniques
#[derive(Debug)]
pub struct HybridErrorMitigation {
    pub virtual_distillation: Option<VirtualDistillation>,
    pub clifford_regression: Option<CliffordDataRegression>,
    pub symmetry_verification: Option<SymmetryVerification>,
    pub subspace_expansion: Option<QuantumSubspaceExpansion>,
}

impl HybridErrorMitigation {
    /// Create a new hybrid mitigation strategy
    pub const fn new() -> Self {
        Self {
            virtual_distillation: None,
            clifford_regression: None,
            symmetry_verification: None,
            subspace_expansion: None,
        }
    }

    /// Enable virtual distillation with specified number of copies
    #[must_use]
    pub fn with_virtual_distillation(mut self, num_copies: usize) -> Self {
        self.virtual_distillation = Some(VirtualDistillation::new(num_copies));
        self
    }

    /// Enable Clifford data regression
    #[must_use]
    pub fn with_clifford_regression(mut self, num_training: usize, degree: usize) -> Self {
        self.clifford_regression = Some(CliffordDataRegression::new(num_training, degree));
        self
    }

    /// Enable symmetry verification
    #[must_use]
    pub fn with_symmetry_verification(
        mut self,
        operators: Vec<Array1<Complex64>>,
        tolerance: f64,
    ) -> Self {
        self.symmetry_verification = Some(SymmetryVerification::new(operators, tolerance));
        self
    }

    /// Enable quantum subspace expansion
    #[must_use]
    pub fn with_subspace_expansion(mut self, num_qubits: usize) -> Self {
        self.subspace_expansion = Some(QuantumSubspaceExpansion::new(num_qubits));
        self
    }

    /// Apply all enabled mitigation techniques sequentially
    pub fn mitigate_comprehensive(
        &mut self,
        raw_measurements: &[f64],
        observable: &Array1<Complex64>,
    ) -> Result<f64, QuantRS2Error> {
        let mut mitigated_value =
            raw_measurements.iter().sum::<f64>() / raw_measurements.len() as f64;

        // Step 1: Virtual Distillation (if enabled)
        if let Some(ref mut vd) = self.virtual_distillation {
            mitigated_value = vd.mitigate_expectation(raw_measurements, observable)?;
        }

        // Step 2: Clifford Data Regression (if enabled and trained)
        if let Some(ref cdr) = self.clifford_regression {
            if cdr.coefficients.is_some() {
                mitigated_value = cdr.mitigate(mitigated_value)?;
            }
        }

        Ok(mitigated_value)
    }
}

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

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

    #[test]
    fn test_virtual_distillation_basic() {
        let mut vd = VirtualDistillation::new(2);
        let measurements = vec![0.9, 0.85, 0.88, 0.92];
        let observable =
            Array1::from_vec(vec![Complex64::new(1.0, 0.0), Complex64::new(-1.0, 0.0)]);

        let result = vd.mitigate_expectation(&measurements, &observable);
        assert!(result.is_ok());

        // Mitigated value should be in reasonable range (virtual distillation produces product of measurements)
        let mitigated = result.expect("Virtual distillation should produce a valid result");
        assert!(
            mitigated > 0.7 && mitigated < 1.0,
            "Expected mitigated value in range (0.7, 1.0), got {}",
            mitigated
        );
    }

    #[test]
    fn test_clifford_data_regression() {
        let mut cdr = CliffordDataRegression::new(10, 2);

        // Synthetic training data with linear noise model
        let noisy: Vec<f64> = (0..10).map(|i| 0.9 * i as f64 / 10.0).collect();
        let ideal: Vec<f64> = (0..10).map(|i| i as f64 / 10.0).collect();

        let train_result = cdr.train(&noisy, &ideal);
        assert!(train_result.is_ok());

        // Test mitigation
        let mitigated = cdr.mitigate(0.45);
        assert!(mitigated.is_ok());
        assert!((mitigated.expect("CDR mitigation should succeed") - 0.5).abs() < 0.1);
    }

    #[test]
    fn test_symmetry_verification() {
        let symmetry_op = Array1::from_vec(vec![
            Complex64::new(1.0, 0.0),
            Complex64::new(1.0, 0.0),
            Complex64::new(-1.0, 0.0),
            Complex64::new(-1.0, 0.0),
        ]);

        let sv = SymmetryVerification::new(vec![symmetry_op], 0.1);

        // Valid eigenstate (all in +1 subspace)
        let valid_state = Array1::from_vec(vec![
            Complex64::new(0.707, 0.0),
            Complex64::new(0.707, 0.0),
            Complex64::new(0.0, 0.0),
            Complex64::new(0.0, 0.0),
        ]);

        let (is_valid, violations) = sv.verify_measurement(&valid_state);
        assert!(is_valid || violations.len() < sv.symmetry_operators.len());
    }

    #[test]
    fn test_quantum_subspace_expansion() {
        let mut qse = QuantumSubspaceExpansion::new(2);
        qse.generate_excitation_basis(2);

        assert_eq!(qse.expansion_basis.len(), 3); // Ground + 2 single excitations

        // Test with a simple noisy state
        let noisy_state = Array1::from_vec(vec![
            Complex64::new(0.9, 0.0),
            Complex64::new(0.1, 0.0),
            Complex64::new(0.05, 0.0),
            Complex64::new(0.0, 0.0),
        ]);

        let coeffs = qse.compute_coefficients(&noisy_state);
        assert!(coeffs.is_ok());
    }

    #[test]
    fn test_error_suppression_factor() {
        let vd = VirtualDistillation::new(3);
        let base_error = 0.1;
        let suppressed = vd.error_suppression_factor(base_error);

        // With 3 copies, error should be suppressed to ~0.1^3 = 0.001
        assert!((suppressed - 0.001).abs() < 1e-6);
    }
}