import qiskit
from qiskit.visualization import plot_bloch_multivector 
import numpy as np
import pylatexenc 
from qiskit import QuantumCircuit, Aer, transpile, assemble
from qiskit.visualization import plot_histogram
from qiskit.providers.aer import AerSimulator

# Function to automatically instantiate the oracle circuit 
def balancedOracle(circuit):
    circuit.barrier()
    circuit.x([0,2])
    circuit.cx([0,2] , 3)
    circuit.x([0,2])
    circuit.barrier()
    return circuit

# number of qubits
n = 3

# initialize circuit
classicalCircuit = QuantumCircuit(n+1,1)

# add not gate to output bit
classicalCircuit.x(n)
# add the oracle to the circuit
classicalCircuit = balancedOracle(classicalCircuit)
# measure the output bit after the oracle
classicalCircuit.measure(n, 0)

display(classicalCircuit.draw('mpl'))

# initialize circuit
dj_Circuit = QuantumCircuit(n+1, n)

# add not gate to the output qubit
dj_Circuit.x(n)
# add hadamard gates to all qubits
dj_Circuit.h(range(n+1))
# add the oracle to the circuit
dj_Circuit = balancedOracle(dj_Circuit)
# finish the "hadamard sandwich" by applying hadamards to the input qubits
dj_Circuit.h(range(n))
# measure the values of the input qubits
dj_Circuit.measure(range(n), range(n))

display(dj_Circuit.draw('mpl'))

# initialize circuit
phaseKickback = QuantumCircuit(2,2)

# create superposition states
phaseKickback.x(1)
phaseKickback.h(range(2))
# entangle states
phaseKickback.cx(0,1)
phaseKickback.h(range(2))
# measure
phaseKickback.measure(range(2), range(2))

display(phaseKickback.draw('mpl'))

# simulate outputs
sim = AerSimulator()
compiled = transpile(phaseKickback, sim)
job = assemble(compiled)

result = sim.run(compiled).result()
counts = result.get_counts()
print(counts)
plot_histogram(counts)


display(dj_Circuit.draw('mpl'))

# simulate algorithm circuit
compiled = transpile(dj_Circuit, sim)
job = assemble(compiled)

result = sim.run(compiled).result()
counts = result.get_counts()
print(counts)
plot_histogram(counts)