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

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

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

# Definition of c_Increment, c_Decrement gates #

def increment(qc,qr):
    """controlled-increment gate, cf. PhysRevA.72.032329"""
    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):
    """controlled-decrement gate, cf. PhysRevA.72.032329"""
    for i in range(num_qubits - 1):
        qc.mct(qr[0:i+1], qr[i+1], qr_aux)
        
        
# Definition of QW cycle #

def quantum_walk(qc,qr):
    """implement DTQW on a previously defined circuit and register cf. PhysRevA.72.032329"""
    # coin operator
    qc.h(qr[0])
    
    #conditional shift
    increment(qc,qr)
    qc.x(qr[0])
    decrement(qc,qr)
    qc.x(qr[0])
    
    
# Definition of circuit analysis functions #
    
def get_tot_gates(qc):
    """get the total number of basic gates of a circuit"""
    tot_gates = 0

    for key in qc.decompose().count_ops():
        tot_gates = tot_gates + qc.decompose().count_ops()[key]
    
    return tot_gates

def get_cx_gates(qc):
    """get the total number of cx gates of a circuit"""
    cx_gates = qc.decompose().count_ops()['cx']
    
    return cx_gates

# Total number of qubits (lattice nodes + coin)
num_qubits = 4
num_steps = 1

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

# Define qCircuit
# We need (num_qubits - 3) aux qubits for mct-gates
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)

# BEGINNING OF QUANTUM CIRCUIT

# Initialization for symmetric state 1/sqrt(2) [|0> + i|1>] of coin

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

qc.barrier()

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

qc.x(qr[num_qubits - 1])

qc.barrier()

# Running the QW
for i in range(num_steps):
    quantum_walk(qc,qr)

qc.barrier()

# Measurement

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

qc.draw(output='mpl')

# Execute the circuit

job = execute(qc, backend=qasm_sim)

results_dict = job.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)