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

# Definition of matchgate

def gate_mu3(qcirc,theta,phi,lam,a,b):
    qcirc.cx(a,b)
    qcirc.cu3(theta,phi,lam,b,a)
    qcirc.cx(a,b)
    
# Number of qubits (should be odd)
n_nodes = 5 

# Number of steps
n_step = 2

# Histogram
hist = True
# Quantum Sphere
#hist = False

# 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])

# Initial state
qc.x(qr[0])

# Creating of two partitions with M1' and M2
# Repeating that n_step times

for k in range(0,n_step):
    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])

if hist:
    for i in range(0,n_nodes):
        qc.measure(qr[i], cr[i])


# To print the circuit  
# QASM_source = qp.get_qasm('QWalk')
# print(QASM_source)

# To use API
# qp.set_api(Qconfig.APItoken, Qconfig.config['url']) 

backend = 'local_qasm_simulator' 

if hist:
    shots = 4096
else:
    shots = 1   # backend 'trick': produces amplitudes instead of probabilities
    
qobj=qp.compile(['QWalk'], backend = backend, shots = shots ) # Compile quantum walk 

result = qp.run(qobj)
print(result)



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

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

import numpy as np

if hist:
    plot_histogram(result.get_counts('QWalk'))
else:    
    data_ampl = result.get_data('QWalk')
    state_walk = data_ampl['quantum_state']
    rho_walk = np.outer(state_walk,state_walk.conj())
    plot_state(rho_walk,'qsphere')
