from qiskit import *
import numpy as np
import matplotlib as mpl
from qiskit.tools.visualization import plot_histogram, plot_state_city

state_sim = Aer.get_backend('statevector_simulator')
qasm_sim = Aer.get_backend('qasm_simulator')

qiskit.__qiskit_version__   #developed in q0.14.0, q-terra0.11.0

# Definition of c_Increment, c_Decrement gates

def increment(qc,qr):
    for i in range(num_qubits - 1):
        qc.mct(qr[0:num_qubits - 1 - i], qr[num_qubits - 1 - i] , qr_aux)
        
def decrement(qc,qr):
    for i in range(num_qubits - 1):
        qc.x(qr[0:num_qubits - 1 - i])
        qc.mct(qr[0:num_qubits - 1 - i], qr[num_qubits - 1 - i], qr_aux)
        qc.x(qr[0:num_qubits - 1 - i])
        
# Definition of QW cycle

def quantum_walk(qc,qr,num_steps):
    for i in range(num_steps):
        qc.h(qr[0])
        increment(qc,qr)
        decrement(qc,qr)

# Input total number of qubits = nodes + coin
num_qubits = 7

# Define qRegister and qCircuit
qr = QuantumRegister(num_qubits, 'qr')
cr = ClassicalRegister(num_qubits - 1, 'cr')

# We need (num_control - 2) aux qubits for mct
if num_qubits > 3:
    qr_aux = QuantumRegister(num_qubits - 3, 'aux')
    qc = QuantumCircuit(qr,qr_aux,cr)
else:
    qr_aux = None
    qc = QuantumCircuit(qr,cr)
    
# Initialization for symmetric state 1/sqrt(2) [|0> + i|1>] of coin

qc.h(qr[0])
qc.s(qr[0])

qc.barrier(qr,qr_aux)

# Initialization for middle state |1000...> of nodes 

qc.x(qr[num_qubits - 1])

qc.barrier(qr,qr_aux)

# Running the QW

quantum_walk(qc,qr,1)

qc.measure(qr[1:num_qubits],cr)

qc.draw(output='mpl')

# Execute the circuit

result = execute(qc, qasm_sim).result()
results_dict = result.get_counts()

# Convert the results to decimal value of cReg and plot

results_dec = {}
for key, value in results_dict.items():
    results_dec[str(int(key,2))] = value
    
plot_histogram(results_dec)

