device = 'ibmq_16_melbourne'

from qiskit import IBMQ
IBMQ.load_accounts()
backend = IBMQ.get_backend(device)

num = backend.configuration()['n_qubits']
coupling_map = backend.configuration()['coupling_map']
print('number of qubits\n',num,'\n')
print('list of pairs of qubits that are directly connected\n',coupling_map)

if device in ['ibmqx2','ibmqx4']:
    pos = { 0: (1,1), 1: (1,0), 2: (0.5,0.5), 3: (0,0), 4: (0,1) }        
elif device in ['ibmqx5']:
    pos = { 0: (0,0), 1: (0,1),  2: (1,1),  3: (2,1),  4: (3,1),  5: (4,1),  6: (5,1),  7: (6,1),
                8: (7,1), 9: (7,0), 10: (6,0), 11: (5,0), 12: (4,0), 13: (3,0), 14: (2,0), 15: (1,0) }
elif device in ['ibmq_16_melbourne']:
    pos = { 0: (0,1), 1: (1,1),  2: (2,1),  3: (3,1),  4: (4,1),  5: (5,1),  6: (6,1),
                7: (7,0), 8: (6,0), 9: (5,0), 10: (4,0), 11: (3,0), 12: (2,0), 13: (1,0) }

import sys
sys.path.append('game_engines')
from universal import layout

grid = layout(num,coupling_map,pos)
grid.plot()

from qiskit import QuantumRegister, ClassicalRegister, QuantumCircuit
from qiskit.tools.visualization import circuit_drawer
import numpy as np

qr = QuantumRegister(2)
cr = ClassicalRegister(2)
qp = QuantumCircuit(qr,cr)

qp.rx( np.pi/2,qr[0])
qp.cx(qr[0],qr[1])

qp.measure(qr,cr)

circuit_drawer(qp)

from qiskit import execute
from qiskit import Aer
job = execute(qp,backend=Aer.get_backend('qasm_simulator'))
results = job.result().get_counts()
print(results)

pairs = grid.matching()
print(pairs)

import random

qr = QuantumRegister(num)
cr = ClassicalRegister(num)
qp = QuantumCircuit(qr,cr)

for pair in pairs:
    qp.ry( (1+2*random.random())*np.pi/4,qr[pair[0]]) # angle generated randonly between pi/4 and 3pi/4
    qp.cx(qr[pair[0]],qr[pair[1]])
    
qp.measure(qr,cr)

from qiskit import execute
job = execute(qp,backend=IBMQ.get_backend('ibmq_qasm_simulator'))
results = job.result().get_counts()

probs = grid.calculate_probs(results)

grid.plot(probs=probs)

job = execute(qp,backend=backend)
results = job.result().get_counts()
probs = grid.calculate_probs(results)
grid.plot(probs=probs)

def mitigate(probs):
    
    av_prob = {}
    for j in range(num): # for each qubit, work out which neighbour it agrees with most
        neighbours = []
        for pair in grid.links:
            if j in grid.links[pair]:
                neighbours.append(pair)
        (guessed_pair,val) = (None,1)
        for pair in neighbours:
            if probs[pair]<val:
                guessed_pair = pair
                val = probs[pair]
        # then find the average probability of a 1 for the two
        av_prob[j] = (probs[grid.links[guessed_pair][0]]+probs[grid.links[guessed_pair][1]])/2
    
    # replace the probabilities for all qubits with these averages
    for j in range(num):
        probs[j] = av_prob[j]
            
    return probs

probs = mitigate(probs)
grid.plot(probs=probs)

pair_labels = {}
colors = {}
for node in grid.pos:
    if type(node)==str:
        pair_labels[node] = node
        colors[node] = (0.5,0.5,0.5)

chosen_pairs = []
m = 0
while len(chosen_pairs)<len(pairs):
    m += 1
    print('\nMOVE',m)
    grid.plot(probs=probs,labels=pair_labels,colors=colors)
    pair = str.upper(input("    > Type the name of a pair of qubits whose numbers are the same (or very similar)...\n"))
    chosen_pairs.append( pair )
    colors[pair] = (0.5,0.5,0.5)
    for j in range(2):
        colors[grid.links[pair][j]] = (0.5,0.5,0.5)

grid.plot(probs=probs,labels=pair_labels,colors=colors)

success = True
for pair in chosen_pairs:
    success = success and ( (grid.links[pair] in pairs) or (grid.links[pair][::-1] in pairs) )

if success:
    input("\n    ** You got all the correct pairs! :) **\n\n          Press any key to continue\n")
else:
    input("\n    ** You didn't get all the correct pairs! :( **\n\n          Press any key to continue\n")

