import sys  
sys.path.insert(0, 'qrw/')
from qrw import QRWsearch
%config InlineBackend.figure_format = 'svg' # Makes the images look nice
search_circuit = QRWsearch(2,0)
# display the matrix the shift operator
search_circuit.S(display_matrix=True);

# display the matrix the shift operator
search_circuit.C(display_matrix=True);

# display the matrix the U
search_circuit.Uprime(display_matrix=True);

# display the operator U
search_circuit.Uprime()

search_circuit.draw_circuit()

search_circuit.plot_states_hist()

import sys  
sys.path.insert(0, 'quantum_decomp/')
import quantum_decomp as qd
gates = qd.matrix_to_gates(search_circuit.Uprime().data)
gates

from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
n_steps=1
q = QuantumRegister (3, name='q') #3 qubit register
c = ClassicalRegister (3, name='c') #3 bit classical register
circ = QuantumCircuit (q,c) #Main circuit
for i in range(1):
    circ.h(0)
    circ.ccx(0, 2, 1)
    circ.ccx(1, 2, 0)
    circ.x(0)
    circ.ccx(0, 1, 2)
    circ.x(0)
    circ.x(2)
    circ.ccx(1, 2, 0)
    circ.ccx(0, 2, 1)
    circ.x(1)
    circ.ccx(1, 2, 0)
    circ.x(1)
    circ.x(2)
    circ.ccx(1, 2, 0)
    circ.x(0)
    circ.ccx(0, 1, 2)
    circ.x(0)
    circ.x(2)
    circ.ccx(1, 2, 0)
    circ.x(2)
    circ.ccx(0, 2, 1)
    circ.ccx(1, 2, 0)
    circ.x(0)
    circ.ccx(0, 1, 2)
    circ.x(0)
    circ.x(1)
    circ.ccx(1, 2, 0)
    circ.x(1)
    circ.ccx(0, 2, 1)
    circ.ccx(1, 2, 0)
circ.barrier()
circ.measure ([q[0],q[1],q[2]], [c[0],c[1],c[2]])
circ.draw('mpl')

from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit.visualization import plot_histogram, plot_state_city
from qiskit import execute, Aer

n_steps=4
q = QuantumRegister (3, name='q') #3 qubit register
c = ClassicalRegister (3, name='c') #3 bit classical register
circ = QuantumCircuit (q,c) #Main circuit
for i in range(n_steps):
    circ.h(q[0]) # coin
    circ.ccx(q[0],q[1],q[2]) # Toffoli 
    circ.cx (q[0], q[1]) # CNOT 
    circ.x(q[0])
    circ.x(q[1])
    circ.ccx(q[0], q[1], q[2])
    circ.x(q[1])
    circ.cx(q[0],q[1])
circ.barrier()
circ.measure([q[0],q[1],q[2]], [c[0],c[1],c[2]])
circ.draw('mpl')

backend = Aer.get_backend('qasm_simulator')
shots = 1024 # number of times circuit is run, for sampling
results = execute(circ, backend=backend, shots=shots).result()
answer = results.get_counts()
plot_histogram(answer)