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

# importing Qiskit
from qiskit import IBMQ, BasicAer, Aer
from qiskit.providers.ibmq import least_busy
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister, execute

# import basic plot tools
from qiskit.tools.visualization import plot_histogram

provider = IBMQ.load_account()

#functions

# def OR(a, b, k):
#     # enter qubit numbers here 
#     """ function does the equivalent of a classical OR between nodes a and b and stores the result in k """
#     qc.barrier(q)
#     qc.cx(q[a], q[k])
#     qc.barrier(q)
#     qc.x(q[a])
#     qc.barrier(q)
#     qc.ccx(q[a], q[b], q[k])
#     qc.barrier(q)
#     qc.x(q[a])
#     qc.barrier(q)

def OR(qubit_1, qubit_2, k):
    # enter qubit numbers here 
    """ function does the equivalent of a classical OR between qubit numbers a and b and stores the result in qubit number k """
    qc.barrier(q)
    qc.x(q[qubit_1])
    qc.x(q[qubit_2])
    qc.barrier(q)
    qc.ccx(q[qubit_1], q[qubit_2], q[k])
    qc.x(q[k])
    qc.barrier(q)
    qc.x(q[qubit_1])
    qc.x(q[qubit_2])
    qc.barrier(q)
    

def are_not_equal(a_0, b_0, k):
    # enter node numbers here. For example, a is node 0, b is node 1 and c is node 2
    """ function outputs 1 if nodes a and b are not the same. Node numbering starts from 0
        as in the problem statement. k is the qubit number where the output is XOR-ed. qubit
        numbering also starts from 0 """
    qc.barrier(q)
    qc.cx(q[2*a_0], q[2*b_0])
    qc.cx(q[(2*a_0) + 1], q[(2*b_0) + 1])
    OR(2*b_0, (2*b_0)+1, k)
    qc.cx(q[2*a_0], q[2*b_0])
    qc.cx(q[(2*a_0) + 1], q[(2*b_0) + 1])
    qc.barrier(q)

# number of nodes 


def initialize_oracle_part():

    qc.h(q[0:6])
    qc.barrier(q)
    are_not_equal(0, 1, 6) # node a and b are not equal 
    are_not_equal(0, 2, 7)
    are_not_equal(1, 2, 8)
    qc.barrier(q)
    qc.mct([q[6], q[7], q[8]], q[10],[q[9]]) # answer is stored in 10. please keep 9 a clean qubit, it's used as ancilla here 
    qc.barrier(q)
    are_not_equal(0, 1, 6) # node a and b are not equal 
    are_not_equal(0, 2, 7)
    are_not_equal(1, 2, 8)
    
def diffusion():
    qc.h(q[0:6])
    qc.x(q[0:6])
    qc.h(q[5])
    qc.barrier()
    qc.mct([q[0], q[1], q[2], q[3], q[4]], q[5], [q[11], q[12], q[13]])
    qc.barrier()
    qc.h(q[5])
    qc.x(q[0:6])
    qc.h(q[0:6])
    
def grover(iter):
    for i in range(iter):
        initialize_oracle_part()
        diffusion()

qubit_num = 20  # max is 32 if you're using the simulator 
q = QuantumRegister(qubit_num)
c = ClassicalRegister(14)
qc = QuantumCircuit(q, c)
qc.x(q[10])
qc.h(q[10])

grover(2)
qc.measure(q[0:6], c[0:6])

qc.draw()

# running and getting results 

backend = Aer.get_backend('qasm_simulator')
job = execute(qc, backend, shots=1000)
# backend = provider.get_backend('ibmq_qasm_simulator')
# job = execute(qc, backend=backend, shots=8000, seed_simulator=12345, backend_options={"fusion_enable":True})
result = job.result()
count = result.get_counts()
print(count)

#sort count
count_sorted = sorted(count.items(), key=lambda x:x[1], reverse=True)

# collect answers with Top 9 probability
ans_list = count_sorted[0:9]

print(ans_list)

