%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
# Importing standard Qiskit libraries
from qiskit import QuantumCircuit, execute
from qiskit.providers.aer import QasmSimulator
from qiskit.visualization import *

num_qubits = 32

alice_basis = np.random.randint(2, size=num_qubits)
alice_state = np.random.randint(2, size=num_qubits)
bob_basis = np.random.randint(2, size=num_qubits)
oscar_basis = np.random.randint(2, size=num_qubits)

print(f"Alice's State:\t {np.array2string(alice_state, separator='')}")
print(f"Alice's Bases:\t {np.array2string(alice_basis, separator='')}")
print(f"Oscar's Bases:\t {np.array2string(oscar_basis, separator='')}")
print(f"Bob's Bases:\t {np.array2string(bob_basis, separator='')}")

def make_bb84_circ(enc_state, enc_basis, meas_basis):
    '''
    enc_state: array of 0s and 1s denoting the state to be encoded
    enc_basis: array of 0s and 1s denoting the basis to be used for encoding
                0 -> Computational Basis
                1 -> Hadamard Basis
    meas_basis: array of 0s and 1s denoting the basis to be used for measurement
                0 -> Computational Basis
                1 -> Hadamard Basis
    '''
    num_qubits = len(enc_state)
    
    bb84_circ = QuantumCircuit(num_qubits)

    # Sender prepares qubits
    for index in range(len(enc_basis)):
        if enc_state[index] == 1:
            bb84_circ.x(index)
        if enc_basis[index] == 1:
            bb84_circ.h(index)
    bb84_circ.barrier()  

    # Receiver measures the received qubits
    for index in range(len(meas_basis)):
        if meas_basis[index] == 1:
            bb84_circ.h(index)
     
    bb84_circ.measure_all()
    
    return bb84_circ


bb84_AO = make_bb84_circ(alice_state, alice_basis, oscar_basis)
oscar_result = execute(bb84_AO.reverse_bits(),
            backend=QasmSimulator(),
            shots=1).result().get_counts().most_frequent()
print(f"Oscar's results:\t {oscar_result}")
# Converting string to array
oscar_state = np.array(list(oscar_result), dtype=int)
print(f"Oscar's State:\t\t{np.array2string(oscar_state, separator='')}")

bb84_OB = make_bb84_circ(oscar_state, oscar_basis, bob_basis)
temp_key = execute(bb84_OB.reverse_bits(),
                   backend=QasmSimulator(),
                   shots=1).result().get_counts().most_frequent()
print(f"Bob's results:\t\t {temp_key}")

alice_key = ''
bob_key = ''
oscar_key = ''
for i in range(num_qubits):
    if alice_basis[i] == bob_basis[i]: # Only choose bits where Alice and Bob chose the same basis
        alice_key += str(alice_state[i])
        bob_key += str(temp_key[i])
        oscar_key += str(oscar_result[i])
print(f"The length of the key is {len(bob_key)}")
print(f"Alice's key contains\t {(alice_key).count('0')} zeroes and {(alice_key).count('1')} ones")
print(f"Bob's key contains\t {(bob_key).count('0')} zeroes and {(bob_key).count('1')} ones")
print(f"Oscar's key contains\t {(oscar_key).count('0')} zeroes and {(oscar_key).count('1')} ones")
print(f"Alice's Key:\t {alice_key}")
print(f"Bob's Key:\t {bob_key}")
print(f"Oscar's Key:\t {oscar_key}")