3 files changed, 275 insertions, 3 deletions

diff --git a/libpsi-core/src/core/circuit.rs b/libpsi-core/src/core/circuit.rs
index f938ef1..e797511 100644
--- a/libpsi-core/src/core/circuit.rs
+++ b/libpsi-core/src/core/circuit.rs
@@ -1,8 +1,9 @@
-use super::{QuantumRegister, QuantumState};
+use super::{CustomGate, QuantumRegister, QuantumState};
 use crate::{format_amplitude, format_probability, Vector};
+use std::sync::Arc;
 use core::fmt;
 
-#[derive(Clone, Copy)]
+#[derive(Clone)]
 pub enum GateOp {
     H(usize),
     X(usize),
@@ -16,10 +17,11 @@ pub enum GateOp {
     CCNOT(usize, usize, usize),
     CSWAP(usize, usize, usize),
     Measure(usize, usize),
+    Custom(Arc<CustomGate>, Vec<usize>),
 }
 
 impl GateOp {
-    pub fn name(&self) -> &'static str {
+    pub fn name(&self) -> &str {
         match self {
             GateOp::H(_) => "H",
             GateOp::X(_) => "X",
@@ -33,6 +35,7 @@ impl GateOp {
             GateOp::CCNOT(_, _, _) => "CCNOT",
             GateOp::CSWAP(_, _, _) => "CSWAP",
             GateOp::Measure(_, _) => "M",
+            GateOp::Custom(gate, _) => &gate.name,
         }
     }
 
@@ -42,6 +45,7 @@ impl GateOp {
             GateOp::CNOT(c, t) | GateOp::CZ(c, t) | GateOp::SWAP(c, t) => vec![*c, *t],
             GateOp::CCNOT(c1, c2, t) | GateOp::CSWAP(c1, c2, t) => vec![*c1, *c2, *t],
             GateOp::Measure(q, _) => vec![*q],
+            GateOp::Custom(_, targets) => targets.clone(),
         }
     }
 
@@ -55,6 +59,10 @@ impl GateOp {
     pub fn is_measurement(&self) -> bool {
         matches!(self, GateOp::Measure(_, _))
     }
+
+    pub fn is_custom(&self) -> bool {
+        matches!(self, GateOp::Custom(_, _))
+    }
 }
 
 pub struct QuantumCircuit {
@@ -128,6 +136,10 @@ impl QuantumCircuit {
                 GateOp::CCNOT(c1, c2, t) => register.apply_gate(&TOFFOLI, &[*c1, *c2, *t]),
                 GateOp::CSWAP(c, t1, t2) => register.apply_gate(&FREDKIN, &[*c, *t1, *t2]),
                 GateOp::Measure(_, _) => {}
+                GateOp::Custom(gate, targets) => {
+                    let quantum_gate = gate.to_quantum_gate();
+                    register.apply_gate(&quantum_gate, targets);
+                }
             }
         }
 
@@ -232,6 +244,18 @@ impl QuantumCircuit {
         self
     }
 
+    pub fn custom(&mut self, gate: &Arc<CustomGate>, targets: &[usize]) -> &mut Self {
+        self.operations.push(GateOp::Custom(Arc::clone(gate), targets.to_vec()));
+        self.computed_state = None;
+        self
+    }
+
+    pub fn apply_custom(&mut self, gate: CustomGate, targets: &[usize]) -> &mut Self {
+        self.operations.push(GateOp::Custom(Arc::new(gate), targets.to_vec()));
+        self.computed_state = None;
+        self
+    }
+
     pub fn reset(&mut self) -> &mut Self {
         self.operations.clear();
         self.computed_state = None;
@@ -272,6 +296,7 @@ impl fmt::Display for QuantumCircuit {
         for (i, op) in self.operations.iter().enumerate() {
             match op {
                 GateOp::Measure(q, c) => writeln!(f, "  {}: {} q{} → c{}", i, op.name(), q, c)?,
+                GateOp::Custom(gate, targets) => writeln!(f, "  {}: [{}] on {:?}", i, gate.name, targets)?,
                 _ => writeln!(f, "  {}: {} on {:?}", i, op.name(), op.quantum_targets())?,
             }
         }
diff --git a/libpsi-core/src/core/custom_gate.rs b/libpsi-core/src/core/custom_gate.rs
new file mode 100644
index 0000000..9e49ac9
--- /dev/null
+++ b/libpsi-core/src/core/custom_gate.rs
@@ -0,0 +1,245 @@
+use crate::{Complex, Matrix, QuantumGate};
+
+#[derive(Clone)]
+pub enum CustomGateDefinition {
+    Matrix(Matrix<Complex<f64>>),
+    Composite(Vec<(CompositeOp, Vec<usize>)>),
+}
+
+#[derive(Clone, Copy)]
+pub enum CompositeOp {
+    H,
+    X,
+    Y,
+    Z,
+    S,
+    T,
+    CNOT,
+    CZ,
+    SWAP,
+    CCNOT,
+    CSWAP,
+}
+
+#[derive(Clone)]
+pub struct CustomGate {
+    pub name: String,
+    pub num_qubits: usize,
+    pub definition: CustomGateDefinition,
+}
+
+impl CustomGate {
+    pub fn from_matrix(name: &str, matrix: Matrix<Complex<f64>>) -> Self {
+        let dim = matrix.rows;
+        let num_qubits = (dim as f64).log2() as usize;
+        assert_eq!(
+            1 << num_qubits,
+            dim,
+            "Matrix dimension must be a power of 2"
+        );
+        assert_eq!(matrix.rows, matrix.cols, "Matrix must be square");
+
+        CustomGate {
+            name: String::from(name),
+            num_qubits,
+            definition: CustomGateDefinition::Matrix(matrix),
+        }
+    }
+
+    pub fn from_composite(
+        name: &str,
+        num_qubits: usize,
+        ops: Vec<(CompositeOp, Vec<usize>)>,
+    ) -> Self {
+        CustomGate {
+            name: String::from(name),
+            num_qubits,
+            definition: CustomGateDefinition::Composite(ops),
+        }
+    }
+
+    pub fn to_quantum_gate(&self) -> QuantumGate<'static> {
+        match &self.definition {
+            CustomGateDefinition::Matrix(matrix) => {
+                let name: &'static str = Box::leak(self.name.clone().into_boxed_str());
+                QuantumGate {
+                    name,
+                    matrix: matrix.clone(),
+                    num_qubits: self.num_qubits,
+                }
+            }
+            CustomGateDefinition::Composite(ops) => {
+                let matrix = self.compute_composite_matrix(ops);
+                let name: &'static str = Box::leak(self.name.clone().into_boxed_str());
+                QuantumGate {
+                    name,
+                    matrix,
+                    num_qubits: self.num_qubits,
+                }
+            }
+        }
+    }
+
+    fn compute_composite_matrix(&self, ops: &[(CompositeOp, Vec<usize>)]) -> Matrix<Complex<f64>> {
+        use crate::gates::*;
+        use crate::Complex;
+
+        let dim = 1 << self.num_qubits;
+        let mut result = Matrix::new(dim, dim, vec![Complex::new(0.0, 0.0); dim * dim]);
+        for i in 0..dim {
+            result.data[i * dim + i] = Complex::new(1.0, 0.0);
+        }
+
+        for (op, targets) in ops {
+            let gate: &QuantumGate = match op {
+                CompositeOp::H => &HADAMARD,
+                CompositeOp::X => &PAULI_X,
+                CompositeOp::Y => &PAULI_Y,
+                CompositeOp::Z => &PAULI_Z,
+                CompositeOp::S => &S_GATE,
+                CompositeOp::T => &T_GATE,
+                CompositeOp::CNOT => &CNOT,
+                CompositeOp::CZ => &CZ,
+                CompositeOp::SWAP => &SWAP,
+                CompositeOp::CCNOT => &TOFFOLI,
+                CompositeOp::CSWAP => &FREDKIN,
+            };
+
+            let full_gate = build_full_operator(&gate.matrix, targets, self.num_qubits);
+            result = matrix_multiply(&full_gate, &result);
+        }
+
+        result
+    }
+}
+
+fn build_full_operator(
+    gate_matrix: &Matrix<Complex<f64>>,
+    targets: &[usize],
+    total_qubits: usize,
+) -> Matrix<Complex<f64>> {
+    let dim = 1 << total_qubits;
+    let gate_dim = gate_matrix.rows;
+    let num_gate_qubits = targets.len();
+
+    let mut result = Matrix::new(dim, dim, vec![Complex::new(0.0, 0.0); dim * dim]);
+
+    for i in 0..dim {
+        for j in 0..dim {
+            let mut gate_i = 0usize;
+            let mut gate_j = 0usize;
+            let mut match_non_targets = true;
+
+            for q in 0..total_qubits {
+                let bit_i = (i >> (total_qubits - 1 - q)) & 1;
+                let bit_j = (j >> (total_qubits - 1 - q)) & 1;
+
+                if let Some(pos) = targets.iter().position(|&t| t == q) {
+                    gate_i |= bit_i << (num_gate_qubits - 1 - pos);
+                    gate_j |= bit_j << (num_gate_qubits - 1 - pos);
+                } else if bit_i != bit_j {
+                    match_non_targets = false;
+                    break;
+                }
+            }
+
+            if match_non_targets {
+                result.data[i * dim + j] = gate_matrix.data[gate_i * gate_dim + gate_j];
+            }
+        }
+    }
+
+    result
+}
+
+fn matrix_multiply(a: &Matrix<Complex<f64>>, b: &Matrix<Complex<f64>>) -> Matrix<Complex<f64>> {
+    let n = a.rows;
+    let mut result = Matrix::new(n, n, vec![Complex::new(0.0, 0.0); n * n]);
+
+    for i in 0..n {
+        for j in 0..n {
+            let mut sum = Complex::new(0.0, 0.0);
+            for k in 0..n {
+                sum = sum + a.data[i * n + k] * b.data[k * n + j];
+            }
+            result.data[i * n + j] = sum;
+        }
+    }
+
+    result
+}
+
+pub struct CustomGateBuilder {
+    name: String,
+    num_qubits: usize,
+    ops: Vec<(CompositeOp, Vec<usize>)>,
+}
+
+impl CustomGateBuilder {
+    pub fn new(name: &str, num_qubits: usize) -> Self {
+        CustomGateBuilder {
+            name: String::from(name),
+            num_qubits,
+            ops: Vec::new(),
+        }
+    }
+
+    pub fn h(mut self, target: usize) -> Self {
+        self.ops.push((CompositeOp::H, vec![target]));
+        self
+    }
+
+    pub fn x(mut self, target: usize) -> Self {
+        self.ops.push((CompositeOp::X, vec![target]));
+        self
+    }
+
+    pub fn y(mut self, target: usize) -> Self {
+        self.ops.push((CompositeOp::Y, vec![target]));
+        self
+    }
+
+    pub fn z(mut self, target: usize) -> Self {
+        self.ops.push((CompositeOp::Z, vec![target]));
+        self
+    }
+
+    pub fn s(mut self, target: usize) -> Self {
+        self.ops.push((CompositeOp::S, vec![target]));
+        self
+    }
+
+    pub fn t(mut self, target: usize) -> Self {
+        self.ops.push((CompositeOp::T, vec![target]));
+        self
+    }
+
+    pub fn cnot(mut self, control: usize, target: usize) -> Self {
+        self.ops.push((CompositeOp::CNOT, vec![control, target]));
+        self
+    }
+
+    pub fn cz(mut self, control: usize, target: usize) -> Self {
+        self.ops.push((CompositeOp::CZ, vec![control, target]));
+        self
+    }
+
+    pub fn swap(mut self, a: usize, b: usize) -> Self {
+        self.ops.push((CompositeOp::SWAP, vec![a, b]));
+        self
+    }
+
+    pub fn ccnot(mut self, c1: usize, c2: usize, target: usize) -> Self {
+        self.ops.push((CompositeOp::CCNOT, vec![c1, c2, target]));
+        self
+    }
+
+    pub fn cswap(mut self, control: usize, t1: usize, t2: usize) -> Self {
+        self.ops.push((CompositeOp::CSWAP, vec![control, t1, t2]));
+        self
+    }
+
+    pub fn build(self) -> CustomGate {
+        CustomGate::from_composite(&self.name, self.num_qubits, self.ops)
+    }
+}
diff --git a/libpsi-core/src/core/mod.rs b/libpsi-core/src/core/mod.rs
index 6936487..1b950f3 100644
--- a/libpsi-core/src/core/mod.rs
+++ b/libpsi-core/src/core/mod.rs
@@ -1,9 +1,11 @@
 pub mod circuit;
 pub mod classical_components;
+pub mod custom_gate;
 pub mod gates;
 pub mod quantum_components;
 
 pub use circuit::*;
 pub use classical_components::*;
+pub use custom_gate::*;
 pub use gates::*;
 pub use quantum_components::*;