# initialization
import numpy as np

# importing Qiskit
from qiskit import IBMQ, Aer
from qiskit.providers.ibmq import least_busy
from qiskit import QuantumCircuit, assemble, transpile

# import basic plot tools
from qiskit.visualization import plot_histogram

import matplotlib as mpl

# set the length of the n-bit input string. 
n = 2 #----------This was the original case considered by David Duetzch-----------#

# set the length of the n-bit input string. 

const_oracle = QuantumCircuit(n+1)

output = np.random.randint(2)
if output == 1:
    const_oracle.x(n)

const_oracle.draw("mpl")

balanced_oracle = QuantumCircuit(n+1)

#------input a 2-bit string as of your choice!!!!!!!!
b_str = "11"

balanced_oracle = QuantumCircuit(n+1)

# Place X-gates
for qubit in range(len(b_str)):
    if b_str[qubit] == '1':
        balanced_oracle.x(qubit)
balanced_oracle.draw("mpl")


# Use barrier as divider
balanced_oracle.barrier()

# Controlled-NOT gates
for qubit in range(n):
    balanced_oracle.cx(qubit, n)

balanced_oracle.barrier()
balanced_oracle.draw("mpl")

# Place X-gates
for qubit in range(len(b_str)):
    if b_str[qubit] == '1':
        balanced_oracle.x(qubit)

# Show oracle
balanced_oracle.draw("mpl")

dj_circuit = QuantumCircuit(n+1, n)

# Apply H-gates
for qubit in range(n):
    dj_circuit.h(qubit)

# Put qubit in state |->
dj_circuit.x(n)
dj_circuit.h(n)
dj_circuit.draw("mpl")

# Add oracle
dj_circuit = QuantumCircuit.compose(dj_circuit,balanced_oracle)
dj_circuit.draw("mpl")

# Repeat H-gates
for qubit in range(n):
    dj_circuit.h(qubit)
dj_circuit.barrier()

# Measure
for i in range(n):
    dj_circuit.measure(i, i)

# Display circuit
dj_circuit.draw("mpl")

# use local simulator
aer_sim = Aer.get_backend('aer_simulator')
qobj = assemble(dj_circuit, aer_sim)
results = aer_sim.run(qobj).result()
answer = results.get_counts()

plot_histogram(answer)

#---------CONCLUDING REMARKS----------------

###--------THE ABOVE ALGORITHM CAN BE EXTENDED TO N-VARAIBLE CASE. IT IS CALLED DUETZCH-JOZSA ALGORITHM-----