1 files changed, 176 insertions, 3 deletions

diff --git a/libpsi-core/src/core/quantum_components.rs b/libpsi-core/src/core/quantum_components.rs
index e6931a0..75f6325 100644
--- a/libpsi-core/src/core/quantum_components.rs
+++ b/libpsi-core/src/core/quantum_components.rs
@@ -1,7 +1,6 @@
-use crate::{column_vector, complex, ColumnVector, Complex, Matrix, Vector, VectorMatrix};
+use crate::{column_vector, complex, ColumnVector, Complex, Float, Matrix, Vector, VectorMatrix};
 use core::{fmt, ops};
 
-// TODO(Hachem): Redo these macros to work with the new function definition.
 #[macro_export]
 macro_rules! count {
     () => { 0 };
@@ -43,6 +42,14 @@ impl QuantumState {
     }
 }
 
+fn identity_matrix<T: Float>(size: usize) -> Matrix<T> {
+    let mut data = vec![T::zero(); size * size];
+    for i in 0..size {
+        data[i * size + i] = T::one();
+    }
+    Matrix::new(size, size, data)
+}
+
 #[derive(Clone)]
 pub struct QuantumBit<'a> {
     state: QuantumState,
@@ -60,6 +67,41 @@ pub struct QuantumRegister<'a> {
 pub struct QuantumGate<'a> {
     pub name: &'a str,
     pub matrix: Matrix<Complex<f64>>,
+    pub num_qubits: usize,
+}
+
+impl<'a> QuantumGate<'a> {
+    pub fn new(name: &'a str, matrix: Matrix<Complex<f64>>, num_qubits: usize) -> Self {
+        let expected_dim = 1 << num_qubits;
+        assert_eq!(
+            matrix.rows, expected_dim,
+            "Gate matrix rows must be 2^num_qubits"
+        );
+        assert_eq!(
+            matrix.cols, expected_dim,
+            "Gate matrix cols must be 2^num_qubits"
+        );
+        QuantumGate {
+            name,
+            matrix,
+            num_qubits,
+        }
+    }
+
+    pub fn from_matrix(name: &'a str, matrix: Matrix<Complex<f64>>) -> Self {
+        assert_eq!(matrix.rows, matrix.cols, "Gate matrix must be square");
+        let dim = matrix.rows;
+        assert!(
+            dim > 0 && (dim & (dim - 1)) == 0,
+            "Matrix dimension must be a power of 2"
+        );
+        let num_qubits = (dim as f64).log2() as usize;
+        QuantumGate {
+            name,
+            matrix,
+            num_qubits,
+        }
+    }
 }
 
 impl<'a> QuantumBit<'a> {
@@ -111,7 +153,7 @@ impl<'a> QuantumRegister<'a> {
         self.state_vector = ColumnVector::from_matrix(&new_result);
     }
 
-    pub fn get_bits(&self) -> Vec<QuantumBit> {
+    pub fn get_bits(&self) -> Vec<QuantumBit<'_>> {
         self.qubits.clone()
     }
 
@@ -122,6 +164,137 @@ impl<'a> QuantumRegister<'a> {
     pub fn get_name(&self) -> &'a str {
         self.name
     }
+
+    pub fn num_qubits(&self) -> usize {
+        self.qubits.len()
+    }
+
+    pub fn apply_gate(&mut self, gate: &QuantumGate, targets: &[usize]) {
+        let n = self.num_qubits();
+
+        assert_eq!(
+            gate.num_qubits,
+            targets.len(),
+            "Number of target qubits must match gate's qubit count"
+        );
+        for &t in targets {
+            assert!(
+                t < n,
+                "Target qubit index {} out of range for {}-qubit register",
+                t,
+                n
+            );
+        }
+
+        let mut sorted_targets = targets.to_vec();
+        sorted_targets.sort();
+        for i in 1..sorted_targets.len() {
+            assert_ne!(
+                sorted_targets[i],
+                sorted_targets[i - 1],
+                "Duplicate target qubit indices are not allowed"
+            );
+        }
+
+        let full_operator = self.build_full_operator(gate, targets);
+
+        self.state_vector = self
+            .state_vector
+            .mul_matrix(&full_operator)
+            .expect("Matrix multiplication failed during gate application");
+    }
+
+    fn build_full_operator(&self, gate: &QuantumGate, targets: &[usize]) -> Matrix<Complex<f64>> {
+        let n = self.num_qubits();
+        let g = gate.num_qubits;
+        let dim = 1 << n;
+
+        let mut contiguous = true;
+        for i in 1..targets.len() {
+            if targets[i] != targets[i - 1] + 1 {
+                contiguous = false;
+                break;
+            }
+        }
+
+        if contiguous && g == n {
+            return gate.matrix.clone();
+        }
+
+        if contiguous {
+            return self.build_contiguous_operator(gate, targets[0]);
+        }
+
+        let mut result = Matrix::new(dim, dim, vec![complex!(0.0, 0.0); dim * dim]);
+
+        for col in 0..dim {
+            for row in 0..dim {
+                let mut target_row_bits = 0usize;
+                let mut target_col_bits = 0usize;
+
+                for (i, &t) in targets.iter().enumerate() {
+                    let qubit_pos = n - 1 - t;
+                    if (row >> qubit_pos) & 1 == 1 {
+                        target_row_bits |= 1 << (g - 1 - i);
+                    }
+                    if (col >> qubit_pos) & 1 == 1 {
+                        target_col_bits |= 1 << (g - 1 - i);
+                    }
+                }
+
+                let mut non_target_match = true;
+                for q in 0..n {
+                    if !targets.contains(&q) {
+                        let qubit_pos = n - 1 - q;
+                        if ((row >> qubit_pos) & 1) != ((col >> qubit_pos) & 1) {
+                            non_target_match = false;
+                            break;
+                        }
+                    }
+                }
+
+                if non_target_match {
+                    result.set(row, col, gate.matrix.get(target_row_bits, target_col_bits));
+                }
+            }
+        }
+
+        result
+    }
+
+    fn build_contiguous_operator(
+        &self,
+        gate: &QuantumGate,
+        start_idx: usize,
+    ) -> Matrix<Complex<f64>> {
+        let n = self.num_qubits();
+        let g = gate.num_qubits;
+
+        let mut result: Option<Matrix<Complex<f64>>> = None;
+
+        for i in 0..n {
+            let part: Matrix<Complex<f64>> = if i == start_idx {
+                gate.matrix.clone()
+            } else if i > start_idx && i < start_idx + g {
+                continue;
+            } else {
+                identity_matrix(2)
+            };
+
+            result = Some(match result {
+                None => part,
+                Some(r) => r.kronecker(&part),
+            });
+        }
+
+        result.unwrap_or_else(|| identity_matrix(1 << n))
+    }
+
+    pub fn apply_gates(&mut self, operations: &[(&QuantumGate, &[usize])]) {
+        for (gate, targets) in operations {
+            self.apply_gate(gate, targets);
+        }
+    }
 }
 
 impl<'a> ops::Index<usize> for QuantumRegister<'a> {