quantrs2-core 0.2.0

//! Quantum Shannon decomposition for arbitrary unitaries
//!
//! This module implements the quantum Shannon decomposition algorithm,
//! which decomposes any n-qubit unitary into a sequence of single-qubit
//! and CNOT gates with asymptotically optimal gate count.

use crate::{
    cartan::OptimizedCartanDecomposer,
    controlled::make_controlled,
    error::{QuantRS2Error, QuantRS2Result},
    gate::{single::*, GateOp},
    matrix_ops::{DenseMatrix, QuantumMatrix},
    qubit::QubitId,
    synthesis::{decompose_single_qubit_zyz, SingleQubitDecomposition},
};
use rustc_hash::FxHashMap;
use scirs2_core::ndarray::{s, Array2};
use scirs2_core::Complex;
use std::f64::consts::PI;

/// Shannon decomposition result for an n-qubit unitary
#[derive(Debug, Clone)]
pub struct ShannonDecomposition {
    /// The decomposed gate sequence
    pub gates: Vec<Box<dyn GateOp>>,
    /// Number of CNOT gates used
    pub cnot_count: usize,
    /// Number of single-qubit gates used
    pub single_qubit_count: usize,
    /// Total circuit depth
    pub depth: usize,
}

/// Shannon decomposer for quantum circuits
pub struct ShannonDecomposer {
    /// Tolerance for numerical comparisons
    tolerance: f64,
    /// Cache for small unitaries
    cache: FxHashMap<u64, ShannonDecomposition>,
    /// Maximum recursion depth
    max_depth: usize,
}

impl ShannonDecomposer {
    /// Create a new Shannon decomposer
    pub fn new() -> Self {
        Self {
            tolerance: 1e-10,
            cache: FxHashMap::default(),
            max_depth: 20,
        }
    }

    /// Create with custom tolerance
    pub fn with_tolerance(tolerance: f64) -> Self {
        Self {
            tolerance,
            cache: FxHashMap::default(),
            max_depth: 20,
        }
    }

    /// Decompose an n-qubit unitary matrix
    pub fn decompose(
        &mut self,
        unitary: &Array2<Complex<f64>>,
        qubit_ids: &[QubitId],
    ) -> QuantRS2Result<ShannonDecomposition> {
        let n = qubit_ids.len();
        let size = 1 << n;

        // Validate input
        if unitary.shape() != [size, size] {
            return Err(QuantRS2Error::InvalidInput(format!(
                "Unitary size {} doesn't match {} qubits",
                unitary.shape()[0],
                n
            )));
        }

        // Check unitarity
        let mat = DenseMatrix::new(unitary.clone())?;
        if !mat.is_unitary(self.tolerance)? {
            return Err(QuantRS2Error::InvalidInput(
                "Matrix is not unitary".to_string(),
            ));
        }

        // Base cases
        if n == 0 {
            return Ok(ShannonDecomposition {
                gates: vec![],
                cnot_count: 0,
                single_qubit_count: 0,
                depth: 0,
            });
        }

        if n == 1 {
            // Single-qubit gate
            let decomp = decompose_single_qubit_zyz(&unitary.view())?;
            let gates = self.single_qubit_to_gates(&decomp, qubit_ids[0]);
            let count = gates.len();

            return Ok(ShannonDecomposition {
                gates,
                cnot_count: 0,
                single_qubit_count: count,
                depth: count,
            });
        }

        if n == 2 {
            // Use specialized two-qubit decomposition
            return self.decompose_two_qubit(unitary, qubit_ids);
        }

        // For n > 2, use recursive Shannon decomposition
        self.decompose_recursive(unitary, qubit_ids, 0)
    }

    /// Recursive Shannon decomposition for n > 2 qubits
    fn decompose_recursive(
        &mut self,
        unitary: &Array2<Complex<f64>>,
        qubit_ids: &[QubitId],
        depth: usize,
    ) -> QuantRS2Result<ShannonDecomposition> {
        if depth > self.max_depth {
            return Err(QuantRS2Error::InvalidInput(
                "Maximum recursion depth exceeded".to_string(),
            ));
        }

        let n = qubit_ids.len();
        let half_size = 1 << (n - 1);

        // Split the unitary into blocks based on the first qubit
        // U = [A B]
        //     [C D]
        let a = unitary.slice(s![..half_size, ..half_size]).to_owned();
        let b = unitary.slice(s![..half_size, half_size..]).to_owned();
        let c = unitary.slice(s![half_size.., ..half_size]).to_owned();
        let d = unitary.slice(s![half_size.., half_size..]).to_owned();

        // Use block decomposition to find V, W such that:
        // U = (I ⊗ V) · Controlled-U_d · (I ⊗ W)
        // where U_d is diagonal in the computational basis
        let (v, w, u_diag) = self.block_diagonalize(&a, &b, &c, &d)?;

        let mut gates: Vec<Box<dyn GateOp>> = Vec::new();
        let mut cnot_count = 0;
        let mut single_qubit_count = 0;

        // Apply W to the lower qubits
        if !self.is_identity(&w) {
            let w_decomp = self.decompose_recursive(&w, &qubit_ids[1..], depth + 1)?;
            gates.extend(w_decomp.gates);
            cnot_count += w_decomp.cnot_count;
            single_qubit_count += w_decomp.single_qubit_count;
        }

        // Apply controlled diagonal gates
        let diag_gates = self.decompose_controlled_diagonal(&u_diag, qubit_ids)?;
        cnot_count += diag_gates.1;
        single_qubit_count += diag_gates.2;
        gates.extend(diag_gates.0);

        // Apply V† to the lower qubits
        if !self.is_identity(&v) {
            let v_dag = v.mapv(|z| z.conj()).t().to_owned();
            let v_decomp = self.decompose_recursive(&v_dag, &qubit_ids[1..], depth + 1)?;
            gates.extend(v_decomp.gates);
            cnot_count += v_decomp.cnot_count;
            single_qubit_count += v_decomp.single_qubit_count;
        }

        // Calculate depth (approximate)
        let depth = gates.len();

        Ok(ShannonDecomposition {
            gates,
            cnot_count,
            single_qubit_count,
            depth,
        })
    }

    /// Block diagonalize a 2x2 block matrix using SVD
    fn block_diagonalize(
        &self,
        a: &Array2<Complex<f64>>,
        b: &Array2<Complex<f64>>,
        c: &Array2<Complex<f64>>,
        d: &Array2<Complex<f64>>,
    ) -> QuantRS2Result<(
        Array2<Complex<f64>>,
        Array2<Complex<f64>>,
        Array2<Complex<f64>>,
    )> {
        let size = a.shape()[0];

        // For block diagonalization, we need to find V, W such that:
        // [A B] = [I 0] [Λ₁ 0 ] [I 0]
        // [C D]   [0 V] [0  Λ₂] [0 W]

        // This is equivalent to finding the CS decomposition
        // For now, use a simpler approach based on QR decomposition

        // If B = 0 and C = 0, already block diagonal
        let b_norm = b.iter().map(|z| z.norm_sqr()).sum::<f64>().sqrt();
        let c_norm = c.iter().map(|z| z.norm_sqr()).sum::<f64>().sqrt();

        if b_norm < self.tolerance && c_norm < self.tolerance {
            let identity = Array2::eye(size);
            let combined = self.combine_blocks(a, b, c, d);
            return Ok((identity.clone(), identity, combined));
        }

        // Use SVD-based approach for general case
        // This is a placeholder - full CS decomposition would be more efficient
        let combined = self.combine_blocks(a, b, c, d);

        // For simplicity, return identity matrices and the full unitary
        // A proper implementation would compute the actual CS decomposition
        let identity = Array2::eye(size);
        Ok((identity.clone(), identity, combined))
    }

    /// Combine 2x2 blocks into a single matrix
    fn combine_blocks(
        &self,
        a: &Array2<Complex<f64>>,
        b: &Array2<Complex<f64>>,
        c: &Array2<Complex<f64>>,
        d: &Array2<Complex<f64>>,
    ) -> Array2<Complex<f64>> {
        let size = a.shape()[0];
        let total_size = 2 * size;
        let mut result = Array2::zeros((total_size, total_size));

        result.slice_mut(s![..size, ..size]).assign(a);
        result.slice_mut(s![..size, size..]).assign(b);
        result.slice_mut(s![size.., ..size]).assign(c);
        result.slice_mut(s![size.., size..]).assign(d);

        result
    }

    /// Decompose controlled diagonal gates
    fn decompose_controlled_diagonal(
        &self,
        diagonal: &Array2<Complex<f64>>,
        qubit_ids: &[QubitId],
    ) -> QuantRS2Result<(Vec<Box<dyn GateOp>>, usize, usize)> {
        let mut gates: Vec<Box<dyn GateOp>> = Vec::new();
        let mut cnot_count = 0;
        let mut single_qubit_count = 0;

        // Extract diagonal elements
        let n = diagonal.shape()[0];
        let mut phases = Vec::with_capacity(n);

        for i in 0..n {
            let phase = diagonal[[i, i]].arg();
            phases.push(phase);
        }

        // Decompose into controlled phase gates
        // This is a simplified version - optimal decomposition would use Gray codes
        let control = qubit_ids[0];

        for (i, &phase) in phases.iter().enumerate() {
            if phase.abs() > self.tolerance {
                if i == 0 {
                    // Global phase on |0⟩ state
                    let gate: Box<dyn GateOp> = Box::new(RotationZ {
                        target: control,
                        theta: phase,
                    });
                    gates.push(gate);
                    single_qubit_count += 1;
                } else {
                    // Controlled phase
                    // For now, use simple controlled-RZ
                    // Optimal would use multi-controlled decomposition
                    let base_gate = Box::new(RotationZ {
                        target: qubit_ids[1],
                        theta: phase,
                    });

                    let controlled = Box::new(make_controlled(vec![control], *base_gate));
                    gates.push(controlled);
                    cnot_count += 2; // Controlled-RZ uses 2 CNOTs
                    single_qubit_count += 3; // And 3 single-qubit gates
                }
            }
        }

        Ok((gates, cnot_count, single_qubit_count))
    }

    /// Specialized two-qubit decomposition
    fn decompose_two_qubit(
        &self,
        unitary: &Array2<Complex<f64>>,
        qubit_ids: &[QubitId],
    ) -> QuantRS2Result<ShannonDecomposition> {
        // Check for identity matrix first
        if self.is_identity(unitary) {
            return Ok(ShannonDecomposition {
                gates: vec![],
                cnot_count: 0,
                single_qubit_count: 0,
                depth: 0,
            });
        }

        // Use Cartan (KAK) decomposition for optimal two-qubit decomposition
        let mut cartan_decomposer = OptimizedCartanDecomposer::new();
        let cartan_decomp = cartan_decomposer.decompose(unitary)?;
        let gates = cartan_decomposer.base.to_gates(&cartan_decomp, qubit_ids)?;

        // Count gates
        let mut cnot_count = 0;
        let mut single_qubit_count = 0;

        for gate in &gates {
            match gate.name() {
                "CNOT" => cnot_count += 1,
                _ => single_qubit_count += 1,
            }
        }

        let depth = gates.len();

        Ok(ShannonDecomposition {
            gates,
            cnot_count,
            single_qubit_count,
            depth,
        })
    }

    /// Convert single-qubit decomposition to gates
    fn single_qubit_to_gates(
        &self,
        decomp: &SingleQubitDecomposition,
        qubit: QubitId,
    ) -> Vec<Box<dyn GateOp>> {
        let mut gates = Vec::new();

        // First RZ rotation
        if decomp.theta1.abs() > self.tolerance {
            gates.push(Box::new(RotationZ {
                target: qubit,
                theta: decomp.theta1,
            }) as Box<dyn GateOp>);
        }

        // RY rotation
        if decomp.phi.abs() > self.tolerance {
            gates.push(Box::new(RotationY {
                target: qubit,
                theta: decomp.phi,
            }) as Box<dyn GateOp>);
        }

        // Second RZ rotation
        if decomp.theta2.abs() > self.tolerance {
            gates.push(Box::new(RotationZ {
                target: qubit,
                theta: decomp.theta2,
            }) as Box<dyn GateOp>);
        }

        // Global phase is ignored in gate sequence

        gates
    }

    /// Check if a matrix is approximately the identity
    fn is_identity(&self, matrix: &Array2<Complex<f64>>) -> bool {
        let n = matrix.shape()[0];

        for i in 0..n {
            for j in 0..n {
                let expected = if i == j {
                    Complex::new(1.0, 0.0)
                } else {
                    Complex::new(0.0, 0.0)
                };
                if (matrix[[i, j]] - expected).norm() > self.tolerance {
                    return false;
                }
            }
        }

        true
    }
}

/// Optimized Shannon decomposition with gate count reduction
pub struct OptimizedShannonDecomposer {
    base: ShannonDecomposer,
    /// Enable peephole optimization
    peephole: bool,
    /// Enable commutation-based optimization
    commutation: bool,
}

impl OptimizedShannonDecomposer {
    /// Create a new optimized decomposer
    pub fn new() -> Self {
        Self {
            base: ShannonDecomposer::new(),
            peephole: true,
            commutation: true,
        }
    }

    /// Decompose with optimization
    pub fn decompose(
        &mut self,
        unitary: &Array2<Complex<f64>>,
        qubit_ids: &[QubitId],
    ) -> QuantRS2Result<ShannonDecomposition> {
        // Get base decomposition
        let mut decomp = self.base.decompose(unitary, qubit_ids)?;

        if self.peephole {
            decomp = self.apply_peephole_optimization(decomp)?;
        }

        if self.commutation {
            decomp = self.apply_commutation_optimization(decomp)?;
        }

        Ok(decomp)
    }

    /// Apply peephole optimization to reduce gate count
    fn apply_peephole_optimization(
        &self,
        mut decomp: ShannonDecomposition,
    ) -> QuantRS2Result<ShannonDecomposition> {
        // Look for patterns like:
        // - Adjacent inverse gates
        // - Mergeable rotations
        // - CNOT-CNOT = Identity

        let mut optimized_gates = Vec::new();
        let mut i = 0;

        while i < decomp.gates.len() {
            if i + 1 < decomp.gates.len() {
                // Check for cancellations
                if self.gates_cancel(&decomp.gates[i], &decomp.gates[i + 1]) {
                    // Skip both gates
                    i += 2;
                    decomp.cnot_count =
                        decomp
                            .cnot_count
                            .saturating_sub(if decomp.gates[i - 2].name() == "CNOT" {
                                2
                            } else {
                                0
                            });
                    decomp.single_qubit_count = decomp.single_qubit_count.saturating_sub(
                        if decomp.gates[i - 2].name() == "CNOT" {
                            0
                        } else {
                            2
                        },
                    );
                    continue;
                }

                // Check for mergeable rotations
                if let Some(merged) =
                    self.try_merge_rotations(&decomp.gates[i], &decomp.gates[i + 1])
                {
                    optimized_gates.push(merged);
                    i += 2;
                    decomp.single_qubit_count = decomp.single_qubit_count.saturating_sub(1);
                    continue;
                }
            }

            optimized_gates.push(decomp.gates[i].clone());
            i += 1;
        }

        decomp.gates = optimized_gates;
        decomp.depth = decomp.gates.len();

        Ok(decomp)
    }

    /// Apply commutation-based optimization
    const fn apply_commutation_optimization(
        &self,
        decomp: ShannonDecomposition,
    ) -> QuantRS2Result<ShannonDecomposition> {
        // Move commuting gates to reduce circuit depth
        // This is a simplified version - full implementation would use
        // a dependency graph and topological sorting

        Ok(decomp)
    }

    /// Check if two gates cancel each other
    fn gates_cancel(&self, gate1: &Box<dyn GateOp>, gate2: &Box<dyn GateOp>) -> bool {
        // Same gate on same qubits
        if gate1.name() == gate2.name() && gate1.qubits() == gate2.qubits() {
            match gate1.name() {
                "X" | "Y" | "Z" | "H" | "CNOT" | "SWAP" => true,
                _ => false,
            }
        } else {
            false
        }
    }

    /// Try to merge two rotation gates
    fn try_merge_rotations(
        &self,
        gate1: &Box<dyn GateOp>,
        gate2: &Box<dyn GateOp>,
    ) -> Option<Box<dyn GateOp>> {
        // Check if both are rotations on the same qubit and axis
        if gate1.qubits() != gate2.qubits() || gate1.qubits().len() != 1 {
            return None;
        }

        let qubit = gate1.qubits()[0];

        match (gate1.name(), gate2.name()) {
            ("RZ", "RZ") => {
                let theta1 = gate1.as_any().downcast_ref::<RotationZ>()?.theta;
                let theta2 = gate2.as_any().downcast_ref::<RotationZ>()?.theta;
                Some(Box::new(RotationZ {
                    target: qubit,
                    theta: theta1 + theta2,
                }))
            }
            ("RX", "RX") => {
                let theta1 = gate1.as_any().downcast_ref::<RotationX>()?.theta;
                let theta2 = gate2.as_any().downcast_ref::<RotationX>()?.theta;
                Some(Box::new(RotationX {
                    target: qubit,
                    theta: theta1 + theta2,
                }))
            }
            ("RY", "RY") => {
                let theta1 = gate1.as_any().downcast_ref::<RotationY>()?.theta;
                let theta2 = gate2.as_any().downcast_ref::<RotationY>()?.theta;
                Some(Box::new(RotationY {
                    target: qubit,
                    theta: theta1 + theta2,
                }))
            }
            _ => None,
        }
    }
}

/// Utility function for quick Shannon decomposition
pub fn shannon_decompose(
    unitary: &Array2<Complex<f64>>,
    qubit_ids: &[QubitId],
) -> QuantRS2Result<Vec<Box<dyn GateOp>>> {
    let mut decomposer = ShannonDecomposer::new();
    let decomp = decomposer.decompose(unitary, qubit_ids)?;
    Ok(decomp.gates)
}

#[cfg(test)]
mod tests {
    use super::*;
    use scirs2_core::ndarray::Array2;
    use scirs2_core::Complex;

    #[test]
    fn test_shannon_single_qubit() {
        let mut decomposer = ShannonDecomposer::new();

        // Hadamard matrix
        let h = Array2::from_shape_vec(
            (2, 2),
            vec![
                Complex::new(1.0, 0.0),
                Complex::new(1.0, 0.0),
                Complex::new(1.0, 0.0),
                Complex::new(-1.0, 0.0),
            ],
        )
        .expect("Failed to create Hadamard matrix")
            / Complex::new(2.0_f64.sqrt(), 0.0);

        let qubit_ids = vec![QubitId(0)];
        let decomp = decomposer
            .decompose(&h, &qubit_ids)
            .expect("Failed to decompose Hadamard gate");

        // Should decompose into at most 3 single-qubit gates
        assert!(decomp.single_qubit_count <= 3);
        assert_eq!(decomp.cnot_count, 0);
    }

    #[test]
    fn test_shannon_two_qubit() {
        let mut decomposer = ShannonDecomposer::new();

        // CNOT matrix
        let cnot = Array2::from_shape_vec(
            (4, 4),
            vec![
                Complex::new(1.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(1.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(1.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(1.0, 0.0),
                Complex::new(0.0, 0.0),
            ],
        )
        .expect("Failed to create CNOT matrix");

        let qubit_ids = vec![QubitId(0), QubitId(1)];
        let decomp = decomposer
            .decompose(&cnot, &qubit_ids)
            .expect("Failed to decompose CNOT gate");

        // Should use at most 3 CNOTs for arbitrary two-qubit gate
        assert!(decomp.cnot_count <= 3);
    }

    #[test]
    fn test_optimized_decomposer() {
        let mut decomposer = OptimizedShannonDecomposer::new();

        // Identity matrix should result in empty circuit
        let identity = Array2::eye(4);
        let identity_complex = identity.mapv(|x| Complex::new(x, 0.0));

        let qubit_ids = vec![QubitId(0), QubitId(1)];
        let decomp = decomposer
            .decompose(&identity_complex, &qubit_ids)
            .expect("Failed to decompose identity matrix");

        // Optimizations should eliminate all gates for identity
        assert_eq!(decomp.gates.len(), 0);
    }

    #[test]
    fn test_merge_rz_rotations() {
        let decomposer = OptimizedShannonDecomposer::new();
        let qubit = QubitId(0);
        let g1 = Box::new(RotationZ {
            target: qubit,
            theta: 0.3,
        }) as Box<dyn GateOp>;
        let g2 = Box::new(RotationZ {
            target: qubit,
            theta: 0.4,
        }) as Box<dyn GateOp>;
        let merged = decomposer
            .try_merge_rotations(&g1, &g2)
            .expect("should merge RZ+RZ");
        let rz = merged
            .as_any()
            .downcast_ref::<RotationZ>()
            .expect("merged gate must be RotationZ");
        assert!(
            (rz.theta - 0.7).abs() < 1e-10,
            "merged theta should be 0.7, got {}",
            rz.theta
        );
    }

    #[test]
    fn test_merge_rx_rotations() {
        let decomposer = OptimizedShannonDecomposer::new();
        let qubit = QubitId(0);
        let g1 = Box::new(RotationX {
            target: qubit,
            theta: 0.5,
        }) as Box<dyn GateOp>;
        let g2 = Box::new(RotationX {
            target: qubit,
            theta: 0.3,
        }) as Box<dyn GateOp>;
        let merged = decomposer
            .try_merge_rotations(&g1, &g2)
            .expect("should merge RX+RX");
        let rx = merged
            .as_any()
            .downcast_ref::<RotationX>()
            .expect("merged gate must be RotationX");
        assert!(
            (rx.theta - 0.8).abs() < 1e-10,
            "merged theta should be 0.8, got {}",
            rx.theta
        );
    }

    #[test]
    fn test_no_merge_different_axes() {
        let decomposer = OptimizedShannonDecomposer::new();
        let qubit = QubitId(0);
        let g1 = Box::new(RotationZ {
            target: qubit,
            theta: 0.3,
        }) as Box<dyn GateOp>;
        let g2 = Box::new(RotationX {
            target: qubit,
            theta: 0.4,
        }) as Box<dyn GateOp>;
        assert!(
            decomposer.try_merge_rotations(&g1, &g2).is_none(),
            "RZ and RX should not merge"
        );
    }
}