from qiskit import *
from random import randint
from qiskit_textbook.problems import dj_problem_oracle


#number of bits in function
qiskit_ex= bool(input("Enter a value if you want to test a random qiskit 4-bit oracle"))
if qiskit_ex==False:
    n=int(input("Enter the number of bits in the function: "))
    if (input("Is the function constant? "))=='True': constant=True
    else: constant=False
    print(constant)
else:
    n=4

#create circuit and initialize states
circuit=QuantumCircuit(n+1,n)
circuit.x(n)
circuit.h(range(n+1))
circuit.barrier()
circuit.draw()

#generate a constant function
try:
    if constant==True:
        oracle=QuantumCircuit(n+1)
        num=randint(0,2)
        if num==0:
            oracle.i(n)
        if num==1:
            oracle.x(n)
    #test oracle for 2 qubit
    if constant==False:
        oracle=QuantumCircuit(n+1)
        b =randint(1,2**n)
        b_str = format(b, '0'+str(n)+'b')
    
        for i in range(len(b_str)):
            if b_str[i]=='1':
                oracle.x(int(i))

        oracle.cx(range(0,n),n)

        for i in range(len(b_str)):
            if b_str[i]=='1':
                oracle.x(int(i))
            
        circuit+=oracle
        circuit.barrier()
except:
    oracle=QuantumCircuit(n+1)
    gate = dj_problem_oracle(2)
    oracle.append(oracle, range(n+1))
    circuit+=oracle
    circuit.barrier()
    circuit.draw()

#reverse initial H gates and add measurements
circuit.h(range(n))
circuit.measure(range(n),range(n))
circuit.draw()

#use Aer's qasm_simulater to execute
backend = BasicAer.get_backend('qasm_simulator')
#execute+ get results
result = execute(circuit, backend, shots = 1000).result()
#get counts of results
counts = result.get_counts()

print(counts)

constant=''
for i in range(n): constant+="0"
if max(counts)==constant: print("the function is constant")
else: print("the function is balanced")

