import numpy as np
from qiskit import QuantumCircuit, execute, Aer, IBMQ, QuantumRegister, ClassicalRegister
from qiskit import IBMQ, BasicAer
from IPython.core.display import Image, display
simulator = Aer.get_backend('qasm_simulator')


# Alice prepares the qubits
input_register = QuantumRegister(1, "x")
output_register = QuantumRegister(1, "f(x)")
classical_register = ClassicalRegister(1, "c")

circuit_Deutsch = QuantumCircuit(input_register, output_register, classical_register)

# Prepare the qubit in the output register in the |1> state
circuit_Deutsch.x(output_register[0])

# Hadamard gates applied on input and output registers
circuit_Deutsch.h(input_register[0])
circuit_Deutsch.h(output_register[0])

# Add a barrier
circuit_Deutsch.barrier()

# Draw the circuit
circuit_Deutsch.draw(output="mpl")

# Bob's four functions/circuits

# The first option: constant f(0) = f(1) = 0
circuit_Bob1 = QuantumCircuit(input_register, output_register)
circuit_Bob1.barrier()

# The second option: constant f(0) = f(1) = 1
circuit_Bob2 = QuantumCircuit(input_register, output_register)
circuit_Bob2.cx(input_register[0], output_register)
circuit_Bob2.x(input_register[0])
circuit_Bob2.cx(input_register[0], output_register)
circuit_Bob2.x(input_register[0])
circuit_Bob2.barrier()

# The third option: balanced f(0) = 0 & f(1) = 1
circuit_Bob3 = QuantumCircuit(input_register, output_register)
circuit_Bob3.cx(input_register[0], output_register)
circuit_Bob3.barrier()

# The fourth option: balanced f(0) = 1 & f(1) = 0
circuit_Bob4 = QuantumCircuit(input_register, output_register)
circuit_Bob4.x(input_register[0])
circuit_Bob4.cx(input_register[0], output_register)
circuit_Bob4.x(input_register[0])
circuit_Bob4.barrier()

import random
list_options_Bob = [circuit_Bob2,circuit_Bob2, circuit_Bob3, circuit_Bob4]
circuit_Bob_choice = random.choice(list_options_Bob)

# Add the chosen circuit to the main circuit
circuit_Deutsch += circuit_Bob_choice

# Draw the circuit 
print("The circuit with Bob's chosen function")
circuit_Deutsch.draw(output="mpl")

# Alice's final operations
circuit_Deutsch.h(input_register[0])
circuit_Deutsch.measure(input_register[0], classical_register[0])

# Draw the final version of the circuit 
print("The final version of the circuit")
circuit_Deutsch.draw(output="mpl")

# The execution of the circuit
counts = execute(circuit_Deutsch, simulator, shots=1).result().get_counts()

# Finding the only key/measurement outcome
measurement_result = list(counts.keys())[0]
print("The results of the measurements: {}".format(counts))
print("The final result is: {}".format(measurement_result))

# From the final measurement result, Alice understands if 
# the Bob's chosen function balanced or constant
if measurement_result == '0':
    print("Bob's chosen function was constant")
elif measurement_result == '1':
    print("Bob's chosen function was balanced")

# Alice: qubit preparation
input_register = QuantumRegister(3, "x")
output_register = QuantumRegister(1, "f(x)")
classical_register = ClassicalRegister(3, "c")

circuit_Deutsch_Jozsa = QuantumCircuit(input_register, output_register, classical_register)

# Prepare the qubit in the output register in the |1> state
circuit_Deutsch_Jozsa.x(output_register[0])

# Hadamard gates on both input and output registers
circuit_Deutsch_Jozsa.h(input_register) # Hadamard gate is applied to all qubits in the input_register
circuit_Deutsch_Jozsa.h(output_register[0])

# Add a barrier
circuit_Deutsch_Jozsa.barrier()
circuit_Deutsch_Jozsa.draw('mpl')


# Circuit for the constant function
circtuit_Deutsch_Jozsa_constant  = QuantumCircuit(input_register, output_register, classical_register)
circtuit_Deutsch_Jozsa_constant += circuit_Deutsch_Jozsa

# Implementing the f(x) = 1 constant function
circtuit_Deutsch_Jozsa_constant.x(output_register[0])

# Add a barrier
circtuit_Deutsch_Jozsa_constant.barrier()

# Final Hadamard gates applied on all qubits in the input register
circtuit_Deutsch_Jozsa_constant.h(input_register)

# Measurements executed for all qubits in the input register
circtuit_Deutsch_Jozsa_constant.measure(input_register, classical_register)

# Draw the circuit
circtuit_Deutsch_Jozsa_constant.draw(output="mpl")

# A function that will help Alice to define if
# the chosen function was balanced or constant 
def is_balanced_or_constant(circuit, bakend):
    # The execution of the circuit
    counts = execute(circuit, bakend, shots=1).result().get_counts()

    # Finding the only key/measurement outcome
    measurement_result = list(counts.keys())[0]
    print("The results of the measurements: {}".format(counts))
    print("The final result is: {}".format(measurement_result))

    # Alice checks if Bob's function was constant or balanced
    if '000' in counts:
        print("Bob's chosen function was constant")
    else:
        print("Bob's chosen function was balanced")

# Alice uses is_balanced_or_constant() function
is_balanced_or_constant(circtuit_Deutsch_Jozsa_constant, simulator)

# Circuit for the balanced case
circtuit_Deutsch_Jozsa_balanced  = QuantumCircuit(input_register, output_register, classical_register)
circtuit_Deutsch_Jozsa_balanced += circuit_Deutsch_Jozsa

# implementing the balanced function
circtuit_Deutsch_Jozsa_balanced.cx(input_register[0], output_register[0])
circtuit_Deutsch_Jozsa_balanced.cx(input_register[1], output_register[0])
circtuit_Deutsch_Jozsa_balanced.cx(input_register[2], output_register[0])

# add a barrier
circtuit_Deutsch_Jozsa_balanced.barrier()

# final Hadamard gates applied on all qubits in the input register
circtuit_Deutsch_Jozsa_balanced.h(input_register)


# measurements executed for all qubits in the input register
circtuit_Deutsch_Jozsa_balanced.measure(input_register, classical_register)

# draw the circuit
circtuit_Deutsch_Jozsa_balanced.draw(output="mpl")

# Alice uses is_balanced_or_constant() function
simulator = Aer.get_backend('qasm_simulator')
is_balanced_or_constant(circtuit_Deutsch_Jozsa_balanced, simulator)