# 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
import qiskit_aer
from matplotlib import pyplot as plt
from IPython.display import HTML

#setup registers
qmr = c2qa.QumodeRegister(1, num_qubits_per_qumode = 6)
qbr = qiskit.QuantumRegister(1)
cr = qiskit.ClassicalRegister(1)

#convenient labeling
qbit = qbr[0]
qumode = qmr[0]

#The circuit is initalized to |0>|0> by default
circuit = c2qa.CVCircuit(qmr, qbr, cr)

#put the qubit into a superposition and then execute a controlled displacement of the cavity
alpha = 2
circuit.h(qbr[0])
circuit.cv_c_d(alpha, qmr[0], qbr[0])

#apply a rotation around an angle pi/4
theta = np.pi/4
circuit.cv_r(theta, qmr[0])

#measure qubit in x basis to collapse into even or odd cat state
circuit.measure_x(qbr, cr)

#simulate and plot
state, _, _ = c2qa.util.simulate(circuit)
c2qa.wigner.plot_wigner(circuit, state)
#the resulting plot is a cat state that is rotated 45 degrees

circuit.clear() #reset

#Initialize Fock state and displace by alpha
alpha = 2
circuit.cv_initialize(3, qmr[0])
circuit.cv_d(alpha, qmr[0])

#Simulate and plot
state, _, _ = c2qa.util.simulate(circuit)
c2qa.wigner.plot_wigner(circuit, state)

#Resulting state is a displaced Fock state

circuit.clear()

s = 1
circuit.cv_initialize(0, qmr[0])
circuit.cv_sq(s, qmr[0])


state, _, _ = c2qa.util.simulate(circuit)
#the resulting plot is a coherent state which is squeezed in the position quadrature
c2qa.wigner.plot_wigner(circuit, state)
circuit.data.pop(1)

#plot a momentum-squeezed state for comparison
circuit.cv_sq(-s, qmr[0])
state, _, _ = c2qa.util.simulate(circuit)
c2qa.wigner.plot_wigner(circuit, state)

qmr = c2qa.QumodeRegister(2, num_qubits_per_qumode = 3)
qbr = qiskit.QuantumRegister(12)
circuit = c2qa.CVCircuit(qmr, qbr)

circuit.cv_sq2(1, qmr[0], qmr[1])

_, result, fock_counts = c2qa.util.simulate(circuit)
qiskit.visualization.plot_histogram(fock_counts)
#Notice the correlation between the state of the two qumodes! They always have the same photon number

qmr = c2qa.QumodeRegister(1, num_qubits_per_qumode=4)
circuit = c2qa.CVCircuit(qmr)

#Define the side lengths of the rectangle
a = 2
b = 1
animate_steps = 4

#Execute the sequence of displacements
circuit.cv_d(a, qmr[0])
circuit.cv_d(1j*b, qmr[0])
circuit.cv_d(-a, qmr[0])
circuit.cv_d(-1j*b, qmr[0])

#animate the result
anim = c2qa.animate.animate_wigner(circuit, animation_segments=10);
HTML(anim.to_html5_video())

qbr = qiskit.QuantumRegister(1)
creg = qiskit.ClassicalRegister(1)

def measure_phase(c): #choose a = b = sqrt(c) to get a square of area c
    shots = 1024
    circuit = c2qa.CVCircuit(qmr, qbr, creg, probe_measure=True)
    circuit.h(qbr[0]) 

    circuit.cv_c_d(np.sqrt(c)/2, qmr[0], qbr[0])
    circuit.cv_d(1j*np.sqrt(c)/2, qmr[0])
    circuit.cv_c_d(-np.sqrt(c)/2, qmr[0], qbr[0])
    circuit.cv_d(-1j*np.sqrt(c)/2, qmr[0])

    circuit.measure_x(qbr[0], creg[0])
    state, result, fock_counts = c2qa.util.simulate(circuit)
    counts = result.get_counts(circuit)
    return (counts.get('0',0)-counts.get('1',0))/shots

measure_phase(0)

c_vals = np.linspace(0, 2*np.pi, 20)
results = []
for c in c_vals:
    results.append(measure_phase(c))
plt.plot(c_vals, results, c = "b", linestyle = "None", marker = "o", label = "data")
plt.plot(c_vals, np.cos(c_vals), c = "k", linestyle = "--", label =r"$\cos(area)$")
plt.xlabel("area/4")
plt.ylabel(r"$\langle \sigma_x \rangle$")
plt.legend()