from qiskit import *
from qiskit.visualization import plot_bloch_multivector
from qiskit.visualization import plot_histogram
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

IBMQ.load_account()

matplotlib inline

my_provider = IBMQ.get_provider()
my_provider.backends()

simulator = my_provider.get_backend('ibmq_5_yorktown')

# Making the quantum circuit
def circuitI(gamma, beta, J, h1, h2):
    qc = QuantumCircuit(2, 2)
    qc.h([0, 1])
    qc.h(1)
    qc.cz(0,1)
    qc.rx(2*gamma*J, 1)
    qc.cz(0,1)
    qc.h(1)
    qc.i(0)
    qc.rz(2*gamma*h1, 0)
    qc.rz(2*gamma*h2, 1)
    qc.rx(2*beta, 0)
    qc.rx(2*beta, 1)
    qc.measure([0,1], [0,1])
    return qc

# Initializing the variables for D.
J = 1
h1 = 0
h2 = 0

#Setting the dimesions of the pasephase. The amount of simulated points = n**2
n = 61
p = np.pi/(n-1)
res = np.zeros([n, n, 4])


# Running a simulation for each point
for i in range(n):
    qc_list = []
    gamma = i*p
    for j in range(n):
        beta = j*p
        qc = circuitI(gamma, beta, J, h1, h2)       
        qc_list.append(qc)
    transpiled_circs = transpile(qc_list, backend=simulator)
    job = execute(transpiled_circs, backend = simulator, shots=1000 )
    for circ_index in range(len(transpiled_circs)):
        counts = job.result().get_counts(qc_list[circ_index])
        res[i][circ_index][:] = np.array([int(counts.get('00', 0)), int(counts.get('01', 0)), int(counts.get('10', 0)),
                            int(counts.get('11', 0))])


print(res)

#Draw the current circuit
qc.draw(output='mpl')

transpiled_circs[0].draw(output='mpl')

# Making a heatmap plot
sns.heatmap(res[:, :, 1]+res[:, :, 2])

