import qiskit as qk
import qiskit_aer as qaer

import numpy as np
import math

sSimulator = qaer.Aer.backends(name="statevector_simulator")[0]
mSimulator = qaer.Aer.backends(name="qasm_simulator")[0]

def create_blackbox():

    def F1(x):
        return 0

    def F2(x):
        return 1

    def F3(x):
        return x%2

    def F4(x):
        return (x+1)%2

    functions = [F1, F2, F3, F4]

    function = functions[ np.random.randint(0, 3) ]

    return function

blackbox = create_blackbox()

print("f(0): ", blackbox(0))
print("f(1): ", blackbox(1))

if (blackbox(0) == blackbox(1)):
    print("f is a constant function")
else:
    print("f is a balanced function")

def return_blackbox():

    q = qk.QuantumRegister(2)
    blackbox = qk.QuantumCircuit(q, name=r"$\hat{U}_B$")

    rand = np.random.randint(0, 3)

    match rand:
    
        case 0: ## f(0,1) -> (0,1)
            blackbox.cx(q[0], q[1])
            print("balanced")
            
        case 1: ## f(0,1) -> (1,0)
            blackbox.x(q[0])
            blackbox.cx(q[0], q[1])
            blackbox.x(q[0])
            print("balanced")
            
        case 2: ## f(0,1) -> (0,0)
            blackbox.id(q[0])
            blackbox.id(q[1])
            print("constant")

        case 3: ## f(0,1) -> (1,1)
            blackbox.id(q[0])
            blackbox.x(q[1])
            print("constant")

    return blackbox.to_instruction()
        

q = qk.QuantumRegister(2, name="q")
c = qk.ClassicalRegister(1, name="c")
deutsch = qk.QuantumCircuit(q, c, name="qc")

deutsch.id(q[0])
deutsch.x(q[1])

deutsch.h(q[0])
deutsch.h(q[1])

deutsch.barrier(label=r"$|\psi\rangle_{in}$")

blackbox = return_blackbox()

deutsch.append(blackbox, [0,1])

deutsch.barrier(label=r"$\hat{U}_B|\psi\rangle_{in}$")

deutsch.h(q[0])
deutsch.h(q[1])

deutsch.barrier()

## We only need to measure q0 to determine if f is balanced or constant
deutsch.measure(q[0], c)

deutsch.draw("mpl")

## We only need 1 shot
## We need to decompose the circuit in order for the simulator to read the blackbox instructions
counts = mSimulator.run(deutsch.decompose(), shots=1).result().get_counts(deutsch.decompose())

if ("0" in counts):
    print(r"0 state measured: f is a constant function")
elif ("1" in counts):
    print(r"1 state measured: f is a balanced function")

