# 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
import numpy as np
from scipy.stats.contingency import margins
from matplotlib import pyplot as plt

qmr = c2qa.QumodeRegister(1, num_qubits_per_qumode = 4)
circuit = c2qa.CVCircuit(qmr)
circuit.cv_initialize(5, qmr[0])

state, _, _ = c2qa.util.simulate(circuit)
c2qa.wigner.plot_wigner(circuit, state)

circuit.cv_initialize(1/np.sqrt(2)*np.array([1,0,0,0,1]), qmr[0])
state, _, _ = c2qa.util.simulate(circuit)
c2qa.wigner.plot_wigner(circuit, state)

qmr = c2qa.QumodeRegister(1, num_qubits_per_qumode = 6)
qbr = qiskit.QuantumRegister(1)
cr = qiskit.ClassicalRegister(1)
circuit = c2qa.CVCircuit(qmr, qbr, cr)
circuit.initialize([1,0], qbr[0])
circuit.cv_initialize(0, qmr[0])

alpha = 2
circuit.h(qbr[0])
circuit.cv_c_d(alpha, qmr[0], qbr[0])
circuit.h(qbr[0])
circuit.measure(qbr, cr)

state, _, _ = c2qa.util.simulate(circuit)
c2qa.wigner.plot_wigner(circuit, state)

#set the number of qubits for simulation. Try setting this to 4 and observe the result!
n = 7
qmr = c2qa.QumodeRegister(1, num_qubits_per_qumode=n)
circuit = c2qa.CVCircuit(qmr)

#Parameter of the coherent state
alpha = 1-1j

#Prepare the coherent state by displacing the vacuum
circuit.cv_d(alpha, qmr[0])

#Simulate the circuit
state, _, _ = c2qa.util.simulate(circuit)

#Get the wigner distribution in units with hbar=1
ax_min = -6
ax_max = 6
steps = 200
w = c2qa.wigner.wigner(state, axes_max= ax_max, axes_min= ax_min, axes_steps = steps)

#Compute the marginals
x_dist, y_dist = margins(w.T)

#normalize the distributions
x_dist *= (ax_max-ax_min)/200
y_dist *= (ax_max-ax_min)/200

#This cell is mostly plotting
fig = plt.figure(figsize=(6,6))
gs = fig.add_gridspec(2, 2,  width_ratios=(3, 1), height_ratios=(1, 3),
                      left=0.1, right=0.9, bottom=0.1, top=0.9,
                      wspace=0.05, hspace=0.05)
ax = fig.add_subplot(gs[1,0])
ax_x = fig.add_subplot(gs[0,0])
ax_y = fig.add_subplot(gs[1,1])

absmax = max(np.max(w), np.abs(np.min(w)))
color_levels = np.linspace(-absmax, absmax, 100)

#Plot the wigner function of the coherent state in the main panel with a colormap
xaxis = np.linspace(ax_min, ax_max, steps)
yaxis = np.linspace(ax_min, ax_max, steps)
cont = ax.contourf(xaxis, yaxis, w, color_levels, cmap = "RdBu_r")

#compute the expected means and variance of the two distributions
mux = np.sqrt(2)*np.real(alpha)
muy = np.sqrt(2)*np.imag(alpha)
sigma = .5

#plot the distributions against the (normalized) marginals of the wigner function
ax_x.plot(xaxis, x_dist,  label = r"$|\psi_\alpha(x)|^2$", c= "b")
ax_x.plot(xaxis, 1/np.sqrt(2*np.pi*sigma)*np.exp(-(xaxis-mux)**2/(2*sigma)), c ="k", linestyle = "--", label="pred.")

ax_y.plot(y_dist.T, yaxis, label = r"$|\phi_\alpha(p)|^2$", c = "r")
ax_y.plot(1/np.sqrt(2*np.pi*sigma)*np.exp(-(yaxis-muy)**2/(2*sigma)), yaxis, c="k", linestyle = "--", label = "pred.")

ax.set(xlabel=r"$\hat x$", ylabel=r"$\hat p$")
ax_x.legend()
ax_y.legend()
ax_x.set_xticklabels([]);
ax_y.set_yticklabels([]);