# The C2QA pacakge is currently not published to PyPI.
# To use the package locally, add the C2QA repository's root folder to the path prior to importing c2qa.
import os
import sys
module_path = os.path.abspath(os.path.join("../.."))
if module_path not in sys.path:
    sys.path.append(module_path)

# Cheat to get MS Visual Studio Code Jupyter server to recognize Python venv
module_path = os.path.abspath(os.path.join("../../venv/Lib/site-packages"))
if module_path not in sys.path:
    sys.path.append(module_path)

import c2qa
import qiskit
from qiskit import QuantumCircuit

qmr0 = c2qa.QumodeRegister(num_qumodes=1, num_qubits_per_qumode=2)
qr0 = qiskit.QuantumRegister(size=1)
circuit0 = c2qa.CVCircuit(qmr0, qr0)

# Initialize qubit
circuit0.initialize([1,0], qr0[0])
# circuit0.initialize([0,1], qr0[0])

# Initialize the qumode to, for example, a Fock sate 1
circuit0.cv_initialize(1, qmr0[0])

# ... Your circuit here ...

state0, _, _ = c2qa.util.simulate(circuit0)
print(state0)

# c2qa.util.trace_out_qumodes() performs a partial trace over the cavities, giving the reduced density matrix of the qubits
print(c2qa.util.trace_out_qumodes(circuit0, state0))

# c2qa.util.trace_out_qubits() performs a partial trace over the qubits, giving the reduced density matrix of the cavities
print(c2qa.util.trace_out_qubits(circuit0, state0))