# 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 numpy as np
import qiskit
import qiskit.visualization

qmr = c2qa.QumodeRegister(1, num_qubits_per_qumode=3) #create a qumoderegister with one qumode
circuit = c2qa.CVCircuit(qmr) #create a CV circuit

circuit.cv_initialize(4, qmr) #initialize the qumode in the state |4>

#simulate and plot the wigner function (more about this in wigner_distribution.ipynb!)
state,_,_ = c2qa.wigner.simulate(circuit)
plot = c2qa.wigner.plot_wigner(circuit, state) #5 rings correspond to 4 quanta
#try out different numbers and plot the wigner functions!

#Intialize a qumode register. To map the number of quanta onto qubit states, the number of
#classical bits needs to equal the number of qubits representing the qumode.
num_qubits = 4
qureg = c2qa.QumodeRegister(1, num_qubits_per_qumode=num_qubits)
creg = qiskit.ClassicalRegister(num_qubits)

#Initialize circuit as a cuperposition of |0> and |2>
circuit = c2qa.CVCircuit(qureg, creg, probe_measure=True)
circuit.cv_initialize(1/np.sqrt(2)*np.array([1,0,1]), qureg[0])

#Use the cv_measure method to map the Fock number onto creg!
circuit.cv_measure([qureg[0]], creg)

#Simulation can use the Aer simulator or the build-in simulate method
_, result, fock_counts = c2qa.util.simulate(circuit)

#Convert the counts dictionary from binary to base-10 and plot histogram
qiskit.visualization.plot_histogram(fock_counts)