# imports:
import numpy as np
from qiskit import IBMQ, Aer
from qiskit import QuantumCircuit, assemble, transpile

n=2
dj_circuit = QuantumCircuit(n+1, n)

bits = []

for qubit in range(n+1):
    dj_circuit.h(qubit)

# Put the n qubit in state |->
dj_circuit.x(n)
dj_circuit.draw(output = 'mpl')

def oracle(bits):
    for i in bits:
        return dj_circuit.x(i)

for qubit in range(n):
    dj_circuit.h(qubit)

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

# I: Create state |00>x|1>

n=2

dj_circuit = QuantumCircuit(n+1, n)

def full_circuit(bits):

    # II: Apply Hadamard to all
    for qubit in range(n+1):
        dj_circuit.h(qubit)

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


    # III: Call Oracle

    oracle(bits)  

    # IV: Apply Hadamard to all qubits in the first register

    for qubit in range(n):
        dj_circuit.h(qubit)

    # V: Measure first register
    for i in range(n):
        measurements = dj_circuit.measure(i, i)
    
    # get results
    
    aer_sim = Aer.get_backend('aer_simulator')
    qobj = assemble(dj_circuit, aer_sim)
    results = aer_sim.run(qobj).result()
    answer = results.get_counts()
    
    return list(answer)[0]


def oracle_type(bits):
    sample = full_circuit(bits)
    ans = ''
    if sample[0]=='0' and sample[1]=='0':
        ans = "constant"
    elif sample[0]=='1' and sample[1]=='1':
        ans = "constant"
    elif sample[0]=='0' and sample[1]=='1':
        ans = "balanced"
    elif sample[0]=='1' and sample[1]=='0':
        ans = "balanced"
    else:
        ans = "Neither constant nor balanced"
    return print(ans)

bits1 = [1,1]
oracle_type(bits1)