# 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
from qiskit.circuit import Gate
from math import pi
qc = QuantumCircuit(2)

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

# Initialize your qubit (should have no effect on Fock state Wigner function)
circuit0.initialize([1,0], qr0[0])
# circuit0.initialize([0,1], qr0[0])

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

# ... Your circtuit here ...

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

# c2qa.wigner.plot_wigner_function() will perform a partial trace to separate the qumode from the qubit state for you, the call below is to log its output as an example.
print(c2qa.util.trace_out_qubits(circuit0, state0))

c2qa.wigner.plot_wigner(circuit0, state0)

qmr1 = c2qa.QumodeRegister(num_qumodes=1, num_qubits_per_qumode=2)
qr1 = qiskit.QuantumRegister(size=1)
circuit1 = c2qa.CVCircuit(qmr1, qr1)

# Initialize your qubit (should have no effect on Fock state Wigner function)
circuit1.initialize([1,0], qr1[0])
# circuit1.initialize([0,1], qr1[0])

# Initialize the qumode to a one Fock sate
circuit1.cv_initialize(1, qmr1[0])

# ... Your circtuit here ...

state1, _, _ = c2qa.util.simulate(circuit1)
print(state1)

# c2qa.wigner.plot_wigner_function() will perform a partial trace to separate the qumode from the qubit state for you, the call below is to log its output as an example.
print(c2qa.util.trace_out_qubits(circuit1, state1))

c2qa.wigner.plot_wigner(circuit1, state1)

import numpy

qmr2 = c2qa.QumodeRegister(num_qumodes=1, num_qubits_per_qumode=2)
qr2 = qiskit.QuantumRegister(size=1)
circuit2 = c2qa.CVCircuit(qmr2, qr2)

# Initialize your qubit (should have no effect on Fock state Wigner function)
circuit2.initialize([1,0], qr2[0])
# circuit1.initialize([0,1], qr1[0])

# Initialize the qumode to a zero Fock sate (ie. the vaccuum)
circuit2.cv_initialize(0, qmr2[0])

# Displace the vaccuum using the displacement gate
# Displace the quasi-probability distribution along the position axis x with a real number
# circuit2.cv_d(numpy.pi/2,qmr1[0])
# Displace the quasi-probability distribution along the momentum axis with an imaginary number
circuit2.cv_d(numpy.pi/2*1j,qmr2[0])


state2, _, _ = c2qa.util.simulate(circuit2)
print(state2)

# c2qa.wigner.plot_wigner_function() will perform a partial trace to separate the qumode from the qubit state for you, the call below is to log its output as an example.
print(c2qa.util.trace_out_qubits(circuit2, state2))

c2qa.wigner.plot_wigner(circuit2, state2)

# Not the expected behavior of displacing the vaccuum, which should simply shift the quasi-probability distribution without distorting it.

# Augment the number of qubits per mode
qmr3 = c2qa.QumodeRegister(num_qumodes=1, num_qubits_per_qumode=6)
qr3 = qiskit.QuantumRegister(size=1)
circuit3 = c2qa.CVCircuit(qmr3, qr3)

# Initialize your qubit (should have no effect on Fock state Wigner function)
circuit3.initialize([1,0], qr3[0])
# circuit1.initialize([0,1], qr1[0])

# Initialize the qumode to a zero Fock sate (ie. the vaccuum)
circuit3.cv_initialize(0, qmr3[0])

# Displace the vaccuum using the displacement gate
# Displace the quasi-probability distribution along the position axis x with a real number
# circuit2.cv_d(numpy.pi/2,qmr1[0])
# Displace the quasi-probability distribution along the momentum axis with an imaginary number
circuit3.cv_d(numpy.pi/2*1j,qmr3[0])

state3, _, _ = c2qa.util.simulate(circuit3)
occ =c2qa. util.stateread(state3, 1, 1, 6)

print(state3)

# c2qa.wigner.plot_wigner_function() will perform a partial trace to separate the qumode from the qubit state for you, the call below is to log its output as an example.
print(c2qa.util.trace_out_qubits(circuit3, state3))

c2qa.wigner.plot_wigner(circuit3, state3)

# This is the expected behavior of displacing the vaccuum:
# a simple shift the quasi-probability distribution without distorting it, creating a coherent state.

def cv_ecdX(qc, qbr, qmr, qm, alpha):
    qc.ry(-np.pi / 2, qbr[0])
    qc.cv_ecd(alpha/2, qmr[qm], qbr[0])
    qc.ry(np.pi / 2, qbr[0])
    return qc

def cv_ecdY(qc, qbr, qmr, qm, alpha):
    qc.rx(-np.pi / 2, qbr[0])
    qc.cv_ecd(alpha/2, qmr[qm], qbr[0])
    qc.rx(np.pi / 2, qbr[0])
    return qc

import numpy as np

# Augment the number of qubits per mode
qmr = c2qa.QumodeRegister(num_qumodes=2, num_qubits_per_qumode=6)
qbr = qiskit.QuantumRegister(size=1)
qc = c2qa.CVCircuit(qmr, qbr)

# Initialize your qubit (should have no effect on Fock state Wigner function)
qc.initialize([1,0], qbr[0])

# initialise in x (just h) or y eigenstate
qc.h(qbr[0])
qc.rz(np.pi/2,qbr[0])

# Initialize the qumode to a zero Fock sate (ie. the vaccuum)
qc.cv_initialize(0, qmr[0])
qc.cv_initialize(0, qmr[1])

alpha = 1j/2
qc = cv_ecdX(qc, qbr, qmr, 0, alpha)
qc = cv_ecdY(qc, qbr, qmr, 1, alpha)
qc = cv_ecdX(qc, qbr, qmr, 0, -alpha)
qc = cv_ecdY(qc, qbr, qmr, 1, -alpha)
alpha = 1/2
qc = cv_ecdX(qc, qbr, qmr, 0, alpha)
qc = cv_ecdY(qc, qbr, qmr, 1, alpha)
qc = cv_ecdX(qc, qbr, qmr, 0, -alpha)
qc = cv_ecdY(qc, qbr, qmr, 1, -alpha)

state, _, _ = c2qa.util.simulate(qc)

# # c2qa.wigner.plot_wigner_function() will perform a partial trace to separate the qumode from the qubit state for you, the call below is to log its output as an example.
# print(c2qa.util.trace_out_qubits(qc, state))

c2qa.wigner.plot_wigner(qc, state, method="iterative")





