# 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
from IPython.display import HTML
import qiskit
import qiskit.visualization
from matplotlib import pyplot as plt
import numpy as np

n = 4
qreg = c2qa.QumodeRegister(1, num_qubits_per_qumode=n)
creg = qiskit.ClassicalRegister(n)
circuit = c2qa.CVCircuit(qreg, creg, probe_measure=True)

#creating a coherent state with λ = (1+1i) using the displacement operator.
#coherent states are sometimes referred to as "displaced ground states"
Lambda = 1+1j
circuit.cv_d(Lambda, qreg[0])

state, _, _ = c2qa.util.simulate(circuit)
c2qa.wigner.plot_wigner(circuit, state)
#The coherent state is displaced along the x axis by sqrt(2)*Re(λ),
#and on the y axis by sqrt(2)*Im(λ)

def measure_overlap(alpha1, alpha2):
    qreg = c2qa.QumodeRegister(1, num_qubits_per_qumode=5)
    circuit = c2qa.CVCircuit(qreg)

    circuit.cv_d(alpha1, qreg[0])
    state1, _, _ = c2qa.util.simulate(circuit)

    circuit.clear()

    circuit.cv_d(alpha2, qreg[0])
    state2, _, _ = c2qa.util.simulate(circuit)
    
    return state1.inner(state2)

ds = np.linspace(0,4,10)
overlaps = []
for d in ds:
    overlaps.append(measure_overlap(d/2, -d/2))

plt.plot(ds, overlaps, linestyle = "None", marker = "o")

xlist = np.linspace(0,4,100)
plt.plot(xlist, np.exp(-(xlist)**2/2))

qr = qiskit.QuantumRegister(1)
qmr = c2qa.QumodeRegister(1, num_qubits_per_qumode=4)

circuit = c2qa.CVCircuit(qmr, qr)

animate_steps = 10
t0 = 0
tf = 1
osc_frequency = 6 #6 GHz oscillator

#prepare initial state
circuit.cv_d(1, qmr[0])

stepsize = (tf-t0)/animate_steps
for t in range(animate_steps):
    circuit.cv_r(-osc_frequency*(stepsize), qmr[0])

anim = c2qa.animate.animate_wigner(circuit, qubit=qr[0], animation_segments=1);
HTML(anim.to_html5_video())