# 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 math

import c2qa
import c2qa.util
import c2qa.wigner
import matplotlib.pyplot
import numpy
import qiskit
import qiskit_aer

def calibration_circuit(dist, num_qumodes = 1, num_qubits_per_qumode = 4):
    qmr = c2qa.QumodeRegister(num_qumodes=num_qumodes, num_qubits_per_qumode=num_qubits_per_qumode)
    qr = qiskit.QuantumRegister(size=1)
    cr = qiskit.ClassicalRegister(size=1)
    circuit = c2qa.CVCircuit(qmr, qr, cr)

    circuit.initialize([1,0], qr[0])
    circuit.cv_initialize(0, qmr[0])

    circuit.h(qr[0])
    circuit.cv_c_d(dist, qmr[0], qr[0])
    circuit.cv_d(1j * dist, qmr[0])
    circuit.cv_c_d(-dist, qmr[0], qr[0])
    circuit.cv_d(-1j * dist, qmr[0])
    circuit.h(qr[0])
    circuit.measure(qr[0], cr[0])

    return circuit, qr[0], cr[0]

circuit, _, _  = calibration_circuit(numpy.sqrt(numpy.pi) / numpy.sqrt(2))
state, result, fock_counts = c2qa.util.simulate(circuit)
counts = result.get_counts(circuit)

print("Simulated statevector:")
print(state)
print("Simulated result counts:")
print(counts)

# Plot Fock state Wigner function
c2qa.wigner.plot_wigner(circuit, state)

# Run the simulations
backend = qiskit_aer.Aer.get_backend("qasm_simulator")
up = numpy.linspace(numpy.sqrt(2), 2, 20)
down = numpy.linspace(2, numpy.sqrt(2), 20)
steps = []

for _ in range(2):
    steps.extend(up[0:19])
    steps.extend(down[0:19])

x = []
y = []

for i, step in enumerate(steps):
    dist = numpy.sqrt(numpy.pi) / step
    circuit, _, _ = calibration_circuit(dist)

    state, result, fock_counts = c2qa.util.simulate(circuit, shots=32)
    counts = result.get_counts(circuit)

    x.append(i)
    y.append(counts.get("0", 0) - counts.get("1", 0))

    print(f"{i}: {counts}")

# Plot the results
matplotlib.pyplot.scatter(x, y)
matplotlib.pyplot.show()

circuit, qubit, cbit  = calibration_circuit(dist=numpy.sqrt(numpy.pi) / numpy.sqrt(2))

anim = c2qa.animate.animate_wigner(circuit, qubit=qubit, cbit=cbit, animation_segments=10, shots=10)

from IPython.display import HTML
HTML(anim.to_html5_video())