# Importing standard Qiskit libraries
from qiskit import QuantumCircuit, execute, Aer
from qiskit.visualization import *
from qiskit.quantum_info import state_fidelity
# Magic function to render plots in the notebook after the cell executing the plot command
%matplotlib inline

def run_on_qasm_simulator(quantum_circuit, num_shots):
    """Takes a circuit, the number of shots and a backend and returns the counts for running the circuit
    on the qasm_simulator backend."""
    qasm_simulator = Aer.get_backend('qasm_simulator')
    job = execute(quantum_circuit, backend=qasm_simulator, shots=num_shots)
    result = job.result()
    counts = result.get_counts(quantum_circuit)
    return counts

def oracle_f1():
    "Oracle implementing function f1"
    qc = QuantumCircuit(3)
    qc.cnot(0,2)    
    qc.cnot(1,2)    
    qc.x(2)
    qc.cnot(1,2)    
    qc.cnot(0,2)    
    return qc

def oracle_f2():
    "Oracle implementing function f2"
    qc = QuantumCircuit(3)
    qc.cnot(0,2)
    qc.x(2)
    qc.cnot(1,2)
    return qc

# We visualize the oracle
qc_oracle_f1 = oracle_f1()
qc_oracle_f1.draw('mpl')

# Create a quantum circuit for 3 qubits and 2 classical registers
qc_deutch_josza_oracle1 = QuantumCircuit(3,2)
# Add the Hadamard gate
qc_deutch_josza_oracle1.h(0)
qc_deutch_josza_oracle1.h(1)
# Apply the oracle
qc_deutch_josza_oracle1.compose(qc_oracle_f1, inplace=True)
# Add the z-gate acting on the ancilla
qc_deutch_josza_oracle1.z(2)
# Apply the oracle again
qc_deutch_josza_oracle1.compose(qc_oracle_f1, inplace=True)
# Add the Hadamard gate
qc_deutch_josza_oracle1.h(0)
qc_deutch_josza_oracle1.h(1)
# Add measurement to the first two qubits
qc_deutch_josza_oracle1.barrier()
qc_deutch_josza_oracle1.measure(range(2),range(2))
# Visualize the circuit
qc_deutch_josza_oracle1.draw('mpl')

# The number of shots we use
num_shots = 100
# Now we run the circuit
res_qc_deutch_josza_oracle1 = run_on_qasm_simulator(qc_deutch_josza_oracle1, num_shots)
# Visualize the results in form of a histogram
plot_histogram(res_qc_deutch_josza_oracle1, title='Deutsch-Josza algorithm, oracle 1')

# We visualize the oracle
qc_oracle_f2 = oracle_f2()
qc_oracle_f2.draw('mpl')

# Create a quantum circuit for 3 qubits and 2 classical registers
qc_deutch_josza_oracle2 = QuantumCircuit(3,2)
# Add the Hadamard gate
qc_deutch_josza_oracle2.h(0)
qc_deutch_josza_oracle2.h(1)
# Apply the oracle
qc_deutch_josza_oracle2.compose(qc_oracle_f2, inplace=True)
# Add the z-gate acting on the ancilla
qc_deutch_josza_oracle2.z(2)
# Apply the oracle again
qc_deutch_josza_oracle2.compose(qc_oracle_f2, inplace=True)
# Add the Hadamard gate
qc_deutch_josza_oracle2.h(0)
qc_deutch_josza_oracle2.h(1)
# Add measurement to the first two qubits
qc_deutch_josza_oracle2.barrier()
qc_deutch_josza_oracle2.measure(range(2),range(2))
# Visualize the circuit
qc_deutch_josza_oracle2.draw('mpl')

# The number of shots we use
num_shots = 100
# Now we run the circuit
res_qc_deutch_josza_oracle2 = run_on_qasm_simulator(qc_deutch_josza_oracle2, num_shots)
# Visualize the results in form of a histogram
plot_histogram(res_qc_deutch_josza_oracle2, title='Deutsch-Josza algorithm, oracle 2')

# Generate the circuits that are preparing the basis states
qc_00 = QuantumCircuit(3)

# 00
qc_01 = QuantumCircuit(3)
qc_01.x(1)

# 10
qc_10 = QuantumCircuit(3)
qc_10.x(0)

# 11
qc_11 = QuantumCircuit(3)
qc_11.x(0)
qc_11.x(1)

qcs_basis_states = [qc_00, qc_01, qc_10, qc_11]

# The number of shots we are going to use
num_shots = 100
# Prepare empty lists for the results of the two oracles
res_oracle1 = list()
res_oracle2 = list()

# We run the oracles on the basis states and record the outcomes
for qc_basis_state in qcs_basis_states:
    # The quantum circuit sending a basis state through oracle 1 and measuring the output
    qc_oracle1 = QuantumCircuit(3,1)
    qc_oracle1.compose(qc_basis_state, inplace=True)
    qc_oracle1.compose(oracle_f1(), inplace=True)
    qc_oracle1.measure(2,0)
    
    # The quantum circuit sending a basis state through oracle 2 and measuring the output
    qc_oracle2 = QuantumCircuit(3,1)
    qc_oracle2.compose(qc_basis_state, inplace=True)
    qc_oracle2.compose(oracle_f2(), inplace=True)
    qc_oracle2.measure(2,0)
   
    # We run the circuits on the qasm simulator
    res_oracle1.append(run_on_qasm_simulator(qc_oracle1, num_shots))
    res_oracle2.append(run_on_qasm_simulator(qc_oracle2, num_shots))

# Visualize the results for oracle 1
plot_histogram(res_oracle1, title='Results oracle 1', legend=['input |00>', 'input |01>', 'input |10>', 'input |11>'])

# Visualize the results for oracle 2
plot_histogram(res_oracle2, title='Results oracle 2', legend=['input |00>', 'input |01>', 'input |10>', 'input |11>'])