

import random
import matplotlib.pyplot as plt

def generate_n_bit_inputs(n):
    """Generate all possible n-bit inputs."""
    return [bin(i)[2:].zfill(n) for i in range(2**n)]

def constant_function(value):
    """Returns a constant function that always returns the given value."""
    return lambda x: value

def balanced_function(n):
    """Returns a balanced function for n-bit inputs."""
    inputs = generate_n_bit_inputs(n)
    half = len(inputs) // 2
    random.shuffle(inputs)
    lookup = {x: 0 if i < half else 1 for i, x in enumerate(inputs)}
    return lambda x: lookup[x]

def determine_function_type(f, n):
    """Determine if the function f is constant or balanced."""
    inputs = generate_n_bit_inputs(n)
    outputs = [f(x) for x in inputs]
    unique_outputs = set(outputs)
    if len(unique_outputs) == 1:
        return "Constant"
    elif outputs.count(0) == outputs.count(1):
        return "Balanced"
    else:
        return "Unknown"

def plot_function(f, n, title):
    """Plot the function outputs for all n-bit inputs."""
    inputs = generate_n_bit_inputs(n)
    outputs = [f(x) for x in inputs]
    
    # Convert binary inputs to integers for plotting
    x = [int(i, 2) for i in inputs]
    y = outputs
    
    plt.figure(figsize=(10, 5))
    plt.scatter(x, y, c='blue')
    plt.title(title)
    plt.xlabel('Input (as integer)')
    plt.ylabel('Output')
    plt.xticks(range(2**n))
    plt.yticks([0, 1])
    plt.grid(True)
    plt.show()

# Define n
n = 3

# Create a constant function that always returns 1
const_func = constant_function(1)
print("Constant Function Test:")
print(f"The function is: {determine_function_type(const_func, n)}")
plot_function(const_func, n, "Constant Function (Always 1)")

# Create a balanced function for n-bit inputs
bal_func = balanced_function(n)
print("\nBalanced Function Test:")
print(f"The function is: {determine_function_type(bal_func, n)}")
plot_function(bal_func, n, "Balanced Function")


# useful additional packages 
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

# importing Qiskit
from qiskit import BasicAer, IBMQ
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister, execute
from qiskit.compiler import transpile
from qiskit.tools.monitor import job_monitor

# import basic plot tools
from qiskit.tools.visualization import plot_histogram

n = 13 # the length of the first register for querying the oracle 


# Choose a type of oracle at random. With probability half it is constant, 
# and with the same probability it is balanced
oracleType, oracleValue = np.random.randint(2), np.random.randint(2)

if oracleType == 0:
    print("The oracle returns a constant value ", oracleValue)
else:
    print("The oracle returns a balanced function")
    a = np.random.randint(1,2**n) # this is a hidden parameter for balanced oracle. 

# Creating registers
# n qubits for querying the oracle and one qubit for storing the answer
qr = QuantumRegister(n+1) #all qubits are initialized to zero
# for recording the measurement on the first register
cr = ClassicalRegister(n)

circuitName = "DeutschJozsa"
djCircuit = QuantumCircuit(qr, cr)

# Create the superposition of all input queries in the first register by applying the Hadamard gate to each qubit.
for i in range(n):
    djCircuit.h(qr[i])

# Flip the second register and apply the Hadamard gate.
djCircuit.x(qr[n])
djCircuit.h(qr[n])
    
# Apply barrier to mark the beginning of the oracle
djCircuit.barrier()

if oracleType == 0:#If the oracleType is "0", the oracle returns oracleValue for all input. 
    if oracleValue == 1:
        djCircuit.x(qr[n])
    else:
        djCircuit.id(qr[n])
else: # Otherwise, it returns the inner product of the input with a (non-zero bitstring) 
    for i in range(n):
        if (a & (1 << i)):
            djCircuit.cx(qr[i], qr[n])
        
# Apply barrier to mark the end of the oracle
djCircuit.barrier()

# Apply Hadamard gates after querying the oracle
for i in range(n):
    djCircuit.h(qr[i])
    
# Measurement
djCircuit.barrier()
for i in range(n):
    djCircuit.measure(qr[i], cr[i])

#draw the circuit
djCircuit.draw(output='mpl',scale=0.5)

backend = BasicAer.get_backend('qasm_simulator')
shots = 1000
job = execute(djCircuit, backend=backend, shots=shots)
results = job.result() 
answer = results.get_counts()

plot_histogram(answer)