use crate::{QoqoError, QOQO_VERSION};
use bincode::{deserialize, serialize};
use pyo3::exceptions::{PyIndexError, PyRuntimeError, PyTypeError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::PyByteArray;
use roqoqo::prelude::*;
use roqoqo::{Circuit, OperationIterator, ROQOQO_VERSION};
use std::collections::HashSet;
use crate::operations::{convert_operation_to_pyobject, convert_pyany_to_operation};
#[pymodule]
fn circuit(_py: Python, module: &PyModule) -> PyResult<()> {
module.add_class::<CircuitWrapper>()?;
Ok(())
}
#[pyclass(name = "Circuit", module = "qoqo")]
#[derive(Clone, Debug, PartialEq)]
pub struct CircuitWrapper {
pub internal: Circuit,
}
impl Default for CircuitWrapper {
fn default() -> Self {
Self::new()
}
}
impl CircuitWrapper {
pub fn from_pyany(input: Py<PyAny>) -> PyResult<Circuit> {
Python::with_gil(|py| -> PyResult<Circuit> {
let input = input.as_ref(py);
if let Ok(try_downcast) = input.extract::<CircuitWrapper>() {
Ok(try_downcast.internal)
} else {
let get_bytes = input.call_method0("to_bincode").map_err(|_| {
PyTypeError::new_err("Python object cannot be converted to qoqo Circuit: Cast to binary representation failed".to_string())
})?;
let bytes = get_bytes.extract::<Vec<u8>>().map_err(|_| {
PyTypeError::new_err("Python object cannot be converted to qoqo Circuit: Cast to binary representation failed".to_string())
})?;
deserialize(&bytes[..]).map_err(|err| {
PyTypeError::new_err(format!(
"Python object cannot be converted to qoqo Circuit: Deserialization failed: {}",
err
))
})
}
})
}
}
#[pymethods]
impl CircuitWrapper {
#[new]
pub fn new() -> Self {
Self {
internal: Circuit::new(),
}
}
pub fn substitute_parameters(
&self,
substitution_parameters: std::collections::HashMap<&str, f64>,
) -> PyResult<Self> {
let mut calculator = qoqo_calculator::Calculator::new();
for (key, val) in substitution_parameters.iter() {
calculator.set_variable(key, *val);
}
Ok(Self {
internal: self
.internal
.substitute_parameters(&calculator)
.map_err(|x| {
pyo3::exceptions::PyRuntimeError::new_err(format!(
"Parameter Substitution failed: {:?}",
x
))
})?,
})
}
pub fn remap_qubits(&self, mapping: std::collections::HashMap<usize, usize>) -> PyResult<Self> {
let new_internal = self.internal.remap_qubits(&mapping).map_err(|err| {
pyo3::exceptions::PyRuntimeError::new_err(format!("Qubit remapping failed: {:?}", err))
})?;
Ok(Self {
internal: new_internal,
})
}
pub fn overrotate(&self) -> PyResult<Self> {
Ok(Self {
internal: self.internal.overrotate().map_err(|_| {
PyRuntimeError::new_err("Error applying PragmaOverrotation in circuit")
})?,
})
}
pub fn count_occurences(&self, operations: Vec<&str>) -> usize {
let mut counter: usize = 0;
for op in self.internal.iter() {
if operations.iter().any(|x| op.tags().contains(x)) {
counter += 1
}
}
counter
}
pub fn get_operation_types(&self) -> HashSet<&str> {
let mut operations: HashSet<&str> = HashSet::new();
for op in self.internal.iter() {
let _ = operations.insert(op.hqslang());
}
operations
}
pub fn __copy__(&self) -> CircuitWrapper {
self.clone()
}
pub fn __deepcopy__(&self, _memodict: Py<PyAny>) -> CircuitWrapper {
self.clone()
}
fn _qoqo_versions(&self) -> (String, String) {
let mut rsplit = ROQOQO_VERSION.split('.').take(2);
let mut qsplit = QOQO_VERSION.split('.').take(2);
let rver = format!(
"{}.{}",
rsplit.next().expect("ROQOQO_VERSION badly formatted"),
rsplit.next().expect("ROQOQO_VERSION badly formatted")
);
let qver = format!(
"{}.{}",
qsplit.next().expect("QOQO_VERSION badly formatted"),
qsplit.next().expect("QOQO_VERSION badly formatted")
);
(rver, qver)
}
pub fn to_bincode(&self) -> PyResult<Py<PyByteArray>> {
let serialized = serialize(&self.internal)
.map_err(|_| PyValueError::new_err("Cannot serialize Circuit to bytes"))?;
let b: Py<PyByteArray> = Python::with_gil(|py| -> Py<PyByteArray> {
PyByteArray::new(py, &serialized[..]).into()
});
Ok(b)
}
#[staticmethod]
pub fn from_bincode(input: &PyAny) -> PyResult<Self> {
let bytes = input
.extract::<Vec<u8>>()
.map_err(|_| PyTypeError::new_err("Input cannot be converted to byte array"))?;
Ok(Self {
internal: deserialize(&bytes[..])
.map_err(|_| PyValueError::new_err("Input cannot be deserialized to Circuit"))?,
})
}
fn to_json(&self) -> PyResult<String> {
let serialized = serde_json::to_string(&self.internal)
.map_err(|_| PyValueError::new_err("Cannot serialize Circuit to json"))?;
Ok(serialized)
}
#[cfg(feature = "json_schema")]
#[staticmethod]
pub fn json_schema() -> String {
let schema = schemars::schema_for!(Circuit);
serde_json::to_string_pretty(&schema).expect("Unexpected failure to serialize schema")
}
#[cfg(feature = "json_schema")]
#[staticmethod]
pub fn current_version() -> String {
ROQOQO_VERSION.to_string()
}
#[cfg(feature = "json_schema")]
pub fn min_supported_version(&self) -> String {
let min_version: (u32, u32, u32) =
Circuit::minimum_supported_roqoqo_version(&self.internal);
format!("{}.{}.{}", min_version.0, min_version.1, min_version.2)
}
#[staticmethod]
pub fn from_json(json_string: &str) -> PyResult<Self> {
Ok(Self {
internal: serde_json::from_str(json_string)
.map_err(|_| PyValueError::new_err("Input cannot be deserialized to Circuit"))?,
})
}
pub fn get(&self, index: usize) -> PyResult<PyObject> {
let operation = self
.internal
.get(index)
.ok_or_else(|| PyIndexError::new_err(format!("Index {} out of range", index)))?
.clone();
convert_operation_to_pyobject(operation)
}
pub fn get_slice(&self, start: Option<usize>, stop: Option<usize>) -> PyResult<CircuitWrapper> {
let start = match start {
Some(x) => x,
_ => 0,
};
let stop = match stop {
Some(x) => x,
_ => self.internal.len(),
};
if start >= stop {
return Err(PyIndexError::new_err(format!(
"Stop index {} smaller than start index {}",
stop, start
)));
}
if start >= self.internal.len() {
return Err(PyIndexError::new_err(format!(
"Start index {} out of range",
start
)));
}
if stop > self.internal.len() {
return Err(PyIndexError::new_err(format!(
"Stop index {} out of range",
stop
)));
}
let mut tmp_iter = self.internal.iter();
if start > 0 {
tmp_iter.nth(start - 1);
}
let circuit_slice: Circuit = tmp_iter.take(stop - start + 1).cloned().collect();
Ok(CircuitWrapper {
internal: circuit_slice,
})
}
pub fn definitions(&self) -> PyResult<Vec<PyObject>> {
let mut defs: Vec<PyObject> = Vec::new();
for op in self
.internal
.definitions()
.iter()
.cloned()
.map(convert_operation_to_pyobject)
{
defs.push(op?)
}
Ok(defs)
}
pub fn operations(&self) -> PyResult<Vec<PyObject>> {
let mut ops: Vec<PyObject> = Vec::new();
for op in self
.internal
.operations()
.iter()
.cloned()
.map(convert_operation_to_pyobject)
{
ops.push(op?)
}
Ok(ops)
}
pub fn filter_by_tag(&self, tag: &str) -> PyResult<Vec<PyObject>> {
let mut tagged: Vec<PyObject> = Vec::new();
for op in self
.internal
.iter()
.filter(|x| x.tags().contains(&tag))
.cloned()
.map(convert_operation_to_pyobject)
{
tagged.push(op?)
}
Ok(tagged)
}
pub fn add(&mut self, op: &PyAny) -> PyResult<()> {
let operation = convert_pyany_to_operation(op).map_err(|x| {
PyTypeError::new_err(format!("Cannot convert python object to Operation {:?}", x))
})?;
self.internal.add_operation(operation);
Ok(())
}
fn __format__(&self, _format_spec: &str) -> PyResult<String> {
Ok(format!("{}", self.internal))
}
fn __repr__(&self) -> PyResult<String> {
Ok(format!("{}", self.internal))
}
fn __richcmp__(&self, other: Py<PyAny>, op: pyo3::class::basic::CompareOp) -> PyResult<bool> {
let other = Self::from_pyany(other);
match op {
pyo3::class::basic::CompareOp::Eq => match other {
Ok(circ) => Ok(self.internal == circ),
_ => Ok(false),
},
pyo3::class::basic::CompareOp::Ne => match other {
Ok(circ) => Ok(self.internal != circ),
_ => Ok(true),
},
_ => Err(pyo3::exceptions::PyNotImplementedError::new_err(
"Other comparison not implemented",
)),
}
}
fn __iter__(slf: PyRef<Self>) -> PyResult<OperationIteratorWrapper> {
Ok(OperationIteratorWrapper {
internal: slf.internal.clone().into_iter(),
})
}
fn __len__(&self) -> usize {
self.internal.len()
}
fn __getitem__(&self, index: usize) -> PyResult<PyObject> {
let operation = self
.internal
.get(index)
.ok_or_else(|| PyIndexError::new_err(format!("Index {} out of range", index)))?
.clone();
convert_operation_to_pyobject(operation)
}
fn __setitem__(&mut self, index: usize, value: &PyAny) -> PyResult<()> {
let operation = convert_pyany_to_operation(value)
.map_err(|_| PyTypeError::new_err("Cannot convert python object to Operation"))?;
let mut_reference = self
.internal
.get_mut(index)
.ok_or_else(|| PyIndexError::new_err(format!("Index {} out of range", index)))?;
*mut_reference = operation;
Ok(())
}
fn __iadd__(&mut self, other: Py<PyAny>) -> PyResult<()> {
Python::with_gil(|py| -> PyResult<()> {
let other_ref = other.as_ref(py);
match convert_pyany_to_operation(other_ref) {
Ok(x) => {
self.internal += x;
Ok(())
}
Err(_) => {
let other = convert_into_circuit(other_ref).map_err(|x| {
pyo3::exceptions::PyTypeError::new_err(format!(
"Right hand side cannot be converted to Operation or Circuit {:?}",
x
))
});
match other {
Ok(x) => {
self.internal += x;
Ok(())
}
Err(y) => Err(y),
}
}
}
})
}
fn __add__(lhs: Py<PyAny>, rhs: Py<PyAny>) -> PyResult<CircuitWrapper> {
Python::with_gil(|py| -> PyResult<CircuitWrapper> {
let (lhs_ref, rhs_ref) = (lhs.as_ref(py), rhs.as_ref(py));
let self_circ = convert_into_circuit(lhs_ref).map_err(|_| {
PyTypeError::new_err("Left hand side can not be converted to Circuit")
})?;
match convert_pyany_to_operation(rhs_ref) {
Ok(x) => Ok(CircuitWrapper {
internal: self_circ + x,
}),
Err(_) => {
let other = convert_into_circuit(rhs_ref).map_err(|_| {
pyo3::exceptions::PyTypeError::new_err(
"Right hand side cannot be converted to Operation or Circuit",
)
})?;
Ok(CircuitWrapper {
internal: self_circ + other,
})
}
}
})
}
}
pub fn convert_into_circuit(input: &PyAny) -> Result<Circuit, QoqoError> {
if let Ok(try_downcast) = input.extract::<CircuitWrapper>() {
return Ok(try_downcast.internal);
}
let get_bytes = input
.call_method0("to_bincode")
.map_err(|_| QoqoError::CannotExtractObject)?;
let bytes = get_bytes
.extract::<Vec<u8>>()
.map_err(|_| QoqoError::CannotExtractObject)?;
deserialize(&bytes[..]).map_err(|_| QoqoError::CannotExtractObject)
}
#[pyclass(name = "OperationIterator", module = "qoqo")]
#[derive(Debug)]
pub struct OperationIteratorWrapper {
internal: OperationIterator,
}
#[pymethods]
impl OperationIteratorWrapper {
fn __iter__(slf: PyRef<Self>) -> PyRef<Self> {
slf
}
fn __next__(mut slf: PyRefMut<Self>) -> Option<PyObject> {
slf.internal
.next()
.map(|op| convert_operation_to_pyobject(op).unwrap())
}
}