import numpy as np
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit import Aer, execute
from qiskit.circuit import Gate
from qiskit.visualization import plot_histogram, plot_bloch_multivector

# Define no. of input qubits
n = 4

circ = QuantumCircuit(n+1, n)
circ.reset(range(n+1))
circ.x(0)
circ.barrier()
circ.h(range(n+1))
circ.barrier()
circ.append(Gate(name="$U_f$", num_qubits=n+1, params=[]), range(n+1))
circ.barrier()
circ.h(range(1, n+1))
circ.barrier()
circ.measure(range(1, n+1), range(n))
circ.draw('mpl', reverse_bits=True)

# Constant Function
constant = QuantumCircuit(n+1, name='Constant')
constant.x(0)
constant.draw('mpl', reverse_bits=True)

# Balanced Function
balanced = QuantumCircuit(n+1, name='Balanced')
for i in range(1, n+1):
    balanced.cx(i, 0)
balanced.draw('mpl', reverse_bits=True)

circ = QuantumCircuit(n+1)
circ.x(0)
circ.barrier()
circ.h(range(n+1))
circ.append(constant.to_instruction(), range(n+1))
circ.h(range(1, n+1))
circ.draw(reverse_bits=True, plot_barriers=False)

simulator = Aer.get_backend('statevector_simulator')
result = execute(circ, simulator).result()
state = result.get_statevector(circ)
plot_bloch_multivector(state)

circ = QuantumCircuit(n+1)
circ.x(0)
circ.barrier()
circ.h(range(n+1))
circ.append(balanced.to_instruction(), range(n+1))
circ.h(range(1, n+1))
circ.draw(reverse_bits=True, plot_barriers=False)

simulator = Aer.get_backend('statevector_simulator')
result = execute(circ, simulator).result()
state = result.get_statevector(circ)
plot_bloch_multivector(state)

circ = QuantumCircuit(n+1, n)
circ.x(0)
circ.barrier()
circ.h(range(n+1))
circ.append(constant.to_instruction(), range(n+1))
circ.h(range(1, n+1))
circ.measure(range(1, n+1), range(n))
circ.draw(reverse_bits=True, plot_barriers=False)

simulator = Aer.get_backend('qasm_simulator')
result = execute(circ, simulator, shots=1000).result()
counts = result.get_counts(circ)
plot_histogram(counts)