# execute this cell twice to see the output due to an issue with matplotlib

import sys
if sys.version_info < (3,5):
    raise Exception('Please use Python version 3.5 or greater.')
    
from pprint import pprint
import math

# importing the QISKit
from qiskit import QuantumProgram

# To use API
#import Qconfig

def gate_mu3(qcirc,theta,phi,lam,a,b):
    qcirc.cx(a,b)
    qcirc.cu3(theta,phi,lam,b,a)
    qcirc.cx(a,b)
    

n_nodes = 5 
n_step = 3

# Creating Programs
qp = QuantumProgram()

# Creating Registers
 
qr = qp.create_quantum_register('qr', n_nodes)

cr = qp.create_classical_register('cr', n_nodes)

# Creating Circuits

qc = qp.create_circuit('QWalk', [qr], [cr])


# Creating of two partitions with M1' and M2
for i in range(0,n_nodes-1,2):
    gate_mu3(qc,math.pi, math.pi, 0, qr[i], qr[i+1])
for i in range(1,n_nodes,2):
    gate_mu3(qc,math.pi/2, 0, 0, qr[i], qr[i+1])



# import state tomography functions
from qiskit.tools.visualization import plot_histogram, plot_state

import numpy as np

# execute the quantum circuit 
backend = 'local_unitary_simulator' # the device to run on

qobj = qp.compile(['QWalk'], backend=backend)

result = qp.run(qobj)

initial_state = np.zeros(2**n_nodes)
initial_state[1]=1.0 # state 0 = ....0000, state 1 = ...000001

QWalk = result.get_data('QWalk')['unitary']

#Applying QWalk n_step times
for i in range(0,n_step):
    if i > 0: initial_state = np.copy(state_QWalk) # Copy previous state
    state_QWalk = np.dot(QWalk,initial_state)      # Multiply on QWalk matrix
    rho_QWalk=np.outer(state_QWalk, state_QWalk.conj()) # Calculate density matrix
    print('step = ',i+1)                           # print number
    plot_state(rho_QWalk,'qsphere')                # draw Quantum Sphere
    