# Imports
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np

from qiskit import QuantumProgram, QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit import available_backends, get_backend, execute, register, least_busy
from qiskit.tools.visualization import plot_histogram, circuit_drawer

# Connecting to the IBM Quantum Experience
import sys, getpass
try:
    sys.path.append("../../") # go to parent dir
    import Qconfig
    qx_config = {
        "APItoken": Qconfig.APItoken,
        "url": Qconfig.config['url']}
    print('Qconfig loaded from %s.' % Qconfig.__file__)
except:
    APItoken = getpass.getpass('Please input your token and hit enter: ')
    qx_config = {
        "APItoken": APItoken,
        "url":"https://quantumexperience.ng.bluemix.net/api"}
    print('Qconfig.py not found in qiskit-tutorial directory; Qconfig loaded using user input.')

register(qx_config['APItoken'], qx_config['url'])

device_shots = 1024
device_name = least_busy(available_backends({'simulator': False, 'local': False}))
device = get_backend(device_name)
device_coupling = device.configuration['coupling_map']

print("the best backend is " + device_name + " with coupling " + str(device_coupling))

# Creating registers
q = QuantumRegister(2)
c = ClassicalRegister(2)

# quantum circuit to make an entangled bell state 
bell = QuantumCircuit(q, c)
bell.h(q[0])
bell.cx(q[0], q[1])

# quantum circuit to measure q in the standard basis 
measureZZ = QuantumCircuit(q, c)
measureZZ.measure(q[0], c[0])
measureZZ.measure(q[1], c[1])
bellZZ = bell+measureZZ

# quantum circuit to measure q in the superposition basis 
measureXX = QuantumCircuit(q, c)
measureXX.h(q[0])
measureXX.h(q[1])
measureXX.measure(q[0], c[0])
measureXX.measure(q[1], c[1])
bellXX = bell+measureXX

# quantum circuit to measure ZX
measureZX = QuantumCircuit(q, c)
measureZX.h(q[0])
measureZX.measure(q[0], c[0])
measureZX.measure(q[1], c[1])
bellZX = bell+measureZX

# quantum circuit to measure XZ
measureXZ = QuantumCircuit(q, c)
measureXZ.h(q[1])
measureXZ.measure(q[0], c[0])
measureXZ.measure(q[1], c[1])
bellXZ = bell+measureXZ

circuits = [bellZZ,bellXX,bellZX,bellXZ]

circuit_drawer(bellZZ)

circuit_drawer(bellXX)

circuit_drawer(bellZX)

circuit_drawer(bellXZ)

job = execute(circuits, backend=device_name, coupling_map=device_coupling, shots=device_shots)
result = job.result()

observable_first ={'00': 1, '01': -1, '10': 1, '11': -1}
observable_second ={'00': 1, '01': 1, '10': -1, '11': -1}
observable_correlated ={'00': 1, '01': -1, '10': -1, '11': 1}

print('IZ = ' + str(result.average_data(bellZZ,observable_first)))
print('ZI = ' + str(result.average_data(bellZZ,observable_second)))
print('ZZ = ' + str(result.average_data(bellZZ,observable_correlated)))

print('IX = ' + str(result.average_data(bellXX,observable_first)))
print('XI = ' + str(result.average_data(bellXX,observable_second)))
print('XX = ' + str(result.average_data(bellXX,observable_correlated)))

print('ZX = ' + str(result.average_data(bellZX,observable_correlated)))
print('XZ = ' + str(result.average_data(bellXZ,observable_correlated)))

CHSH = lambda x : x[0]+x[1]+x[2]-x[3]
measure = [measureZZ, measureZX, measureXX, measureXZ]

local_backend = 'local_qasm_simulator'
local_shots = 8192

sim_chsh_circuits = []
sim_x = []

sim_steps = 30
for step in range(sim_steps):
    theta = 2.0*np.pi*step/30
    bell_middle = QuantumCircuit(q,c)
    bell_middle.ry(theta,q[0])
    for m in measure:
        sim_chsh_circuits.append(bell+bell_middle+m)
    sim_x.append(theta)

job = execute(sim_chsh_circuits, backend=local_backend, shots=local_shots)
result = job.result()

sim_chsh = []
circ = 0

for x in range(len(sim_x)):
    temp_chsh = []
    for m in range(len(measure)):
        temp_chsh.append(result.average_data(sim_chsh_circuits[circ].name,observable_correlated))
        circ += 1
    sim_chsh.append(CHSH(temp_chsh))

real_chsh_circuits = []
real_x = []

real_steps = 10
for step in range(real_steps):
    theta = 2.0*np.pi*step/10
    bell_middle = QuantumCircuit(q,c)
    bell_middle.ry(theta,q[0])
    for m in measure:
        real_chsh_circuits.append(bell+bell_middle+m)
    real_x.append(theta)

job = execute(real_chsh_circuits, backend=device_name, coupling_map=device_coupling, shots=device_shots)
result = job.result()

real_chsh = []
circ = 0

for x in range(len(real_x)):
    temp_chsh = []
    for m in range(len(measure)):
        temp_chsh.append(result.average_data(real_chsh_circuits[circ].name,observable_correlated))
        circ += 1
    real_chsh.append(CHSH(temp_chsh))

plt.plot(sim_x, sim_chsh, 'r-', real_x, real_chsh, 'bo')
plt.plot([0, 2*np.pi], [2, 2], 'b-')
plt.plot([0, 2*np.pi], [-2, -2], 'b-')
plt.grid()
plt.ylabel('CHSH', fontsize=20)
plt.xlabel(r'$Y(\theta)$', fontsize=20)
plt.show()

print(real_chsh)

device_shots = 8192 # need more statistics for 3 and 4 qubit experiments

# 2 - qubits 

# quantum circuit to make GHZ state
q2 = QuantumRegister(2)
c2 = ClassicalRegister(2)
ghz = QuantumCircuit(q2, c2)
ghz.h(q2[0])
ghz.cx(q2[0],q2[1])

# quantum circuit to measure q in standard basis 
measureZZ = QuantumCircuit(q2, c2)
measureZZ.measure(q2[0], c2[0])
measureZZ.measure(q2[1], c2[1])
ghzZZ = ghz+measureZZ

measureXX = QuantumCircuit(q2, c2)
measureXX.h(q2[0])
measureXX.h(q2[1])
measureXX.measure(q2[0], c2[0])
measureXX.measure(q2[1], c2[1])
ghzXX = ghz+measureXX

circuits2 = [ghzZZ, ghzXX]

circuit_drawer(ghzZZ)

circuit_drawer(ghzXX)

job2 = execute(circuits2, backend=device_name, coupling_map=device_coupling, shots=device_shots)
result2 = job2.result()

plot_histogram(result2.get_counts(ghzZZ))
plot_histogram(result2.get_counts(ghzXX))

# 3 - qubits 

# quantum circuit to make GHZ state 
q3 = QuantumRegister(3)
c3 = ClassicalRegister(3)
ghz3 = QuantumCircuit(q3, c3)
ghz3.h(q3[0])
ghz3.cx(q3[0],q3[1])
ghz3.cx(q3[1],q3[2])

# quantum circuit to measure q in standard basis 
measureZZZ = QuantumCircuit(q3, c3)
measureZZZ.measure(q3[0], c3[0])
measureZZZ.measure(q3[1], c3[1])
measureZZZ.measure(q3[2], c3[2])
ghzZZZ = ghz3+measureZZZ

measureXXX = QuantumCircuit(q3, c3)
measureXXX.h(q3[0])
measureXXX.h(q3[1])
measureXXX.h(q3[2])
measureXXX.measure(q3[0], c3[0])
measureXXX.measure(q3[1], c3[1])
measureXXX.measure(q3[2], c3[2])
ghzXXX = ghz3+measureXXX

circuits3 = [ghzZZZ, ghzXXX]

circuit_drawer(ghzZZZ)

circuit_drawer(ghzXXX)

job3 = execute(circuits3, backend=device_name, coupling_map=device_coupling, shots=device_shots)
result3 = job3.result()

plot_histogram(result3.get_counts(ghzZZZ))
plot_histogram(result3.get_counts(ghzXXX))

# 4 - qubits 

# quantum circuit to make GHZ state 
q4 = QuantumRegister(4)
c4 = ClassicalRegister(4)
ghz4 = QuantumCircuit(q4, c4)
ghz4.h(q4[0])
ghz4.cx(q4[0],q4[1])
ghz4.cx(q4[1],q4[2])
ghz4.h(q4[3])
ghz4.h(q4[2])
ghz4.cx(q4[3],q4[2])
ghz4.h(q4[3])
ghz4.h(q4[2])

# quantum circuit to measure q in standard basis 
measureZZZZ = QuantumCircuit(q4, c4)
measureZZZZ.measure(q4[0], c4[0])
measureZZZZ.measure(q4[1], c4[1])
measureZZZZ.measure(q4[2], c4[2])
measureZZZZ.measure(q4[3], c4[3])
ghzZZZZ = ghz4+measureZZZZ

measureXXXX = QuantumCircuit(q4, c4)
measureXXXX.h(q4[0])
measureXXXX.h(q4[1])
measureXXXX.h(q4[2])
measureXXXX.h(q4[3])
measureXXXX.measure(q4[0], c4[0])
measureXXXX.measure(q4[1], c4[1])
measureXXXX.measure(q4[2], c4[2])
measureXXXX.measure(q4[3], c4[3])
ghzXXXX = ghz4+measureXXXX

circuits4 = [ghzZZZZ, ghzXXXX]

circuit_drawer(ghzZZZZ)

circuit_drawer(ghzXXXX)

job4 = execute(circuits4, backend=device_name, coupling_map=device_coupling, shots=device_shots)
result4 = job4.result()

plot_histogram(result4.get_counts(ghzZZZZ))
plot_histogram(result4.get_counts(ghzXXXX))

# quantum circuit to make GHZ state 
q3 = QuantumRegister(3)
c3 = ClassicalRegister(3)
ghz3 = QuantumCircuit(q3, c3)
ghz3.h(q3[0])
ghz3.cx(q3[0],q3[1])
ghz3.cx(q3[0],q3[2])

# quantum circuit to measure q in standard basis 
measureZZZ = QuantumCircuit(q3, c3)
measureZZZ.measure(q3[0], c3[0])
measureZZZ.measure(q3[1], c3[1])
measureZZZ.measure(q3[2], c3[2])
ghzZZZ = ghz3+measureZZZ

circuits5 = [ghzZZZ]

circuit_drawer(ghzZZZ)

job5 = execute(circuits5, backend=device_name, coupling_map=device_coupling, shots=device_shots)
result5 = job5.result()

plot_histogram(result5.get_counts(ghzZZZ))

MerminM = lambda x : x[0]*x[1]*x[2]*x[3]

observable ={'000': 1, '001': -1, '010': -1, '011': 1, '100': -1, '101': 1, '110': 1, '111': -1}

# quantum circuit to measure q XXX 
measureXXX = QuantumCircuit(q3, c3)
measureXXX.h(q3[0])
measureXXX.h(q3[1])
measureXXX.h(q3[2])
measureXXX.measure(q3[0], c3[0])
measureXXX.measure(q3[1], c3[1])
measureXXX.measure(q3[2], c3[2])
ghzXXX = ghz3+measureXXX

# quantum circuit to measure q XYY
measureXYY = QuantumCircuit(q3, c3)
measureXYY.s(q3[1]).inverse()
measureXYY.s(q3[2]).inverse()
measureXYY.h(q3[0])
measureXYY.h(q3[1])
measureXYY.h(q3[2])
measureXYY.measure(q3[0], c3[0])
measureXYY.measure(q3[1], c3[1])
measureXYY.measure(q3[2], c3[2])
ghzXYY = ghz3+measureXYY

# quantum circuit to measure q YXY
measureYXY = QuantumCircuit(q3, c3)
measureYXY.s(q3[0]).inverse()
measureYXY.s(q3[2]).inverse()
measureYXY.h(q3[0])
measureYXY.h(q3[1])
measureYXY.h(q3[2])
measureYXY.measure(q3[0], c3[0])
measureYXY.measure(q3[1], c3[1])
measureYXY.measure(q3[2], c3[2])
ghzYXY = ghz3+measureYXY

# quantum circuit to measure q YYX
measureYYX = QuantumCircuit(q3, c3)
measureYYX.s(q3[0]).inverse()
measureYYX.s(q3[1]).inverse()
measureYYX.h(q3[0])
measureYYX.h(q3[1])
measureYYX.h(q3[2])
measureYYX.measure(q3[0], c3[0])
measureYYX.measure(q3[1], c3[1])
measureYYX.measure(q3[2], c3[2])
ghzYYX = ghz3+measureYYX

circuits6 = [ghzXXX, ghzYYX, ghzYXY, ghzXYY]

circuit_drawer(ghzXXX)

circuit_drawer(ghzYYX)

circuit_drawer(ghzYXY)

circuit_drawer(ghzXYY)

job6 = execute(circuits6, backend=device_name, coupling_map=device_coupling, shots=device_shots)
result6 = job6.result()

temp=[]
temp.append(result6.average_data(ghzXXX,observable))
temp.append(result6.average_data(ghzYYX,observable))
temp.append(result6.average_data(ghzYXY,observable))
temp.append(result6.average_data(ghzXYY,observable))
print(MerminM(temp))

