!pip3 install qiskit
import qiskit as q
 

!pip3 install pylatexenc

import math 
import numpy as np

statevec_sim = q.Aer.get_backend("statevector_simulator")

def statevec_exec(circ):
  statevec = q.execute(circ, backend=statevec_sim).result().get_statevector()
  return statevec

# qr = q.QuantumRegister(2, name="q")
# cr = q.ClassicalRegister(2, name="cr")

def build_circ(num_qbits, num_cbits):
  qr = q.QuantumRegister(num_qbits)
  cr = q.ClassicalRegister(num_cbits)
  final_circ = q.QuantumCircuit(qr, cr)
  return final_circ, qr, cr

# simple_circuit, qr, cr = build_circ(2)
# simple_circuit.draw()

def h_gates(qcirc, a, b): # 2 inputs and h gates in parallel
  qcirc.h(a)
  qcirc.h(b)   

#### start of by applying a rotation around y axis 

simple_circuit, qr, cr = build_circ(2, 2)

simple_circuit.ry(np.pi/6, qr[0])
simple_circuit.ry(np.pi/6, qr[1])
simple_circuit_statevec = statevec_exec(simple_circuit)

simple_circuit.measure([0,1], [0,1]) 


simple_circuit.draw()


q.visualization.plot_bloch_multivector(simple_circuit_statevec)

qasm_sim = q.Aer.get_backend('qasm_simulator')
orig_counts = q.execute(simple_circuit, backend=qasm_sim, shots=1024).result().get_counts()
q.visualization.plot_histogram([orig_counts], legend=['output: org qs'])

hadamard_front, qr1, cr1 = build_circ(2, 2)
hadamard_front.draw()

h_gates(hadamard_front, qr1[0], qr1[1])
hadamard_front.ry(math.pi/6, qr1[0])
hadamard_front.ry(math.pi/6, qr1[1])
hadamard_front_statevec = statevec_exec(hadamard_front)
hadamard_front.measure([0,1], [0,1])
hadamard_front.draw()

q.visualization.plot_bloch_multivector(hadamard_front_statevec)

#qasm_sim = q.Aer.get_backend('qasm_simulator')
hfront_counts = q.execute(hadamard_front, backend=qasm_sim, shots=1024).result().get_counts()
q.visualization.plot_histogram([hfront_counts], legend=['output: hfront qs'])

hadamard_front_back, qr2, cr2 = build_circ(2, 2)
h_gates(hadamard_front_back, qr2[0], qr2[1])
hadamard_front_back.ry(math.pi/6, [qr2[0]])
hadamard_front_back.ry(math.pi/6, [qr2[1]])
h_gates(hadamard_front_back, qr2[0], qr2[1])
#hadamard_front.barrier()
hadamard_front_back_statevec = statevec_exec(hadamard_front_back)
hadamard_front_back.measure([0,1], [0,1])
hadamard_front_back.draw()


q.visualization.plot_bloch_multivector(hadamard_front_back_statevec)

 
hfront_back_counts = q.execute(hadamard_front_back, backend=qasm_sim, shots=1024).result().get_counts()
q.visualization.plot_histogram([hfront_back_counts], legend=['output: hfront-back qs'])

num_input = 3 

balanced_circuit = q.QuantumCircuit(num_input+1) # num_input + one more where qubit is initialized to |1>

### we can create the balanced oracle by applying CNOT gate to every input and/or by wrapping the second CNOT gate with X gates 

x_str = "010" # this is for iterating over the string and add X gate when there's 1. 

for n in range(len(x_str)):
  if x_str[n]=="1":
    balanced_circuit.x(n)
balanced_circuit.draw()    

### add control not for all the inputs 

balanced_circuit.barrier()

for i in range(num_input):
  balanced_circuit.cx(i, num_input)
balanced_circuit.barrier()
balanced_circuit.draw()  

### add another x gate

for n in range(len(x_str)):
  if x_str[n]=="1":
    balanced_circuit.x(n)
balanced_circuit.draw()    

init_circuit, qr, cr = build_circ(num_input+1, num_input)

### all the input qubits  are initialized to 0, while the target qubit is initialized to 1. 

for n in range(num_input):
  init_circuit.h(n)

init_circuit.x(qr[num_input]) # state is now initialized to |1>
init_circuit.h(qr[num_input])
init_circuit.draw()  

complete_circuit = init_circuit.compose(balanced_circuit)
complete_circuit.draw()

## now we add h gates again to the input qubits 

for n in range(num_input):
  complete_circuit.h(n)

complete_circuit.barrier()

# Now we measure the input qubits 

for n in range(num_input):
  complete_circuit.measure(n, n)

complete_circuit.draw()  

from google.colab import drive
drive.mount('/content/drive')

path = '/content/drive/My Drive/'

style = {'backgroundcolor': 'lavender', 'dpi':200, 'subfontsize':10}

complete_circuit.draw(output='mpl',scale=0.8, style=style, filename=path+'DJ_Circ_balanced.png')

qasm_sim = q.Aer.get_backend('qasm_simulator')
balanced_counts = q.execute(complete_circuit, backend=qasm_sim, shots=1024).result().get_counts()
q.visualization.plot_histogram([balanced_counts], legend=['output: balanced'])

### let's use IBM Quantum Computer to simulate this, 

from qiskit import IBMQ

IBMQ.save_account('find your token')
IBMQ.load_account()

provider = IBMQ.get_provider('ibm-q')
qcomp = provider.get_backend('ibmq_athens') # check which computer has 0 jobs on queue
job = q.execute(complete_circuit, backend=qcomp)
q_result = job.result()

q.visualization.plot_histogram(q_result.get_counts(complete_circuit))