# ! pip install numpy 
# ! pip install qiskit
# ! pip install matplotlib

from qiskit import QuantumCircuit, execute, Aer, BasicAer
from qiskit.visualization import plot_histogram

import numpy as np


n = 2

f1 = QuantumCircuit(n+1, name='f1')

# Applying the x gate to ancilla for when f_1 = 1 
f1.cx(1,2)


display(f1.draw())


f2 = QuantumCircuit(n+1, name='f2')

# Applying the x gate to ancilla for when f_1 = 1 
f2.cx(0,2)
f2.cx(1,2)

f2.ccx(0,1,2)

display(f2.draw())

dj = QuantumCircuit(n+1, n)

# Apply H-gates
for qubit in range(n):
    dj.h(qubit)
    
# Put ancilia qubit in state |->
dj.x(n)
dj.h(n)


dj.barrier()

# Add oracle corresponding to the fucntion
dj.append(f1, [*range(n+1)])

dj.barrier()

# Repeat H-gates
for qubit in range(n):
    dj.h(qubit)
dj.barrier()

# Measure
for i in range(n):
    dj.measure(i, i)

# Display circuit
dj.draw('mpl')

# use local simulator
backend = BasicAer.get_backend('qasm_simulator')
shots = 1024
results = execute(dj, backend=backend, shots=shots).result()
answer = results.get_counts()

plot_histogram(answer)

# Since we measure the state |00> with probability 0, the funciton is balanced as expected

dj = QuantumCircuit(n+1, n)

# Apply H-gates
for qubit in range(n):
    dj.h(qubit)
    
# Put ancilia qubit in state |->
dj.x(n)
dj.h(n)


dj.barrier()

# Add oracle corresponding to the fucntion
dj.append(f2, [*range(n+1)])

dj.barrier()

# Repeat H-gates
for qubit in range(n):
    dj.h(qubit)
dj.barrier()

# Measure
for i in range(n):
    dj.measure(i, i)

# Display circuit
dj.draw('mpl')

# use local simulator
backend = BasicAer.get_backend('qasm_simulator')
shots = 1024
results = execute(dj, backend=backend, shots=shots).result()
answer = results.get_counts()

plot_histogram(answer)

# Since we measure the state |00> with probability ~0.25, the function is biased as expected

dj = QuantumCircuit(n+1, n)

# Apply H-gates
for qubit in range(n):
    dj.h(qubit)
    
# Put ancilia qubit in state |->
dj.x(n)
dj.h(n)


dj.barrier()

# Add oracle corresponding to the function
dj.append(f1, [*range(n+1)])
dj.append(f2, [*range(n+1)])

dj.barrier()

# Repeat H-gates
for qubit in range(n):
    dj.h(qubit)
dj.barrier()

# Measure
for i in range(n):
    dj.measure(i, i)

# Display circuit
dj.draw('mpl')

# use local simulator
backend = BasicAer.get_backend('qasm_simulator')
shots = 1024
results = execute(dj, backend=backend, shots=shots).result()
answer = results.get_counts()

plot_histogram(answer)


