# %pip install numpy==1.19 qiskit==0.20 pylatexenc # Please uncomment this line if you are running on Google Colab or uncomment and run once if you are running locally.
%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 *
import qiskit
from packaging.version import parse as parse_version
assert parse_version(np.__version__) >= parse_version('1.19'),  "Please install the correct version of numpy using the command 'pip install --upgrade numpy==1.19'"
assert parse_version(qiskit.__qiskit_version__['qiskit-terra']) >= parse_version('0.15'), "Please make sure you have the correct version of Qiskit installed or run this on IBM Quantum Experience" 
assert parse_version(qiskit.__qiskit_version__['qiskit-aer']) >= parse_version('0.6'), "Please make sure you have the correct version of Qiskit installed or run this on IBM Quantum Experience" 
assert parse_version(qiskit.__qiskit_version__['qiskit']) >= parse_version('0.20'),"Please make sure you have the correct version of Qiskit installed or run this on IBM Quantum Experience" 
from cryptography.fernet import Fernet
import base64
basis_gates = ['id', 'x', 'y', 'z', 's', 't', 'sdg', 'tdg', 'h', 'p', 'sx' ,'r', 'rx', 'ry', 'rz', 'u', 'u1', 'u2', 'u3', 'cx', 'barrier', 'measure']
secret_message = b'gAAAAABfevgMDRKfpM75bCBMUfAvaUW_Fjs2PxFYkYOSCldJTUnl8oLKVZRaiPitXqwQwbMTx4YwSCf_n0HQ-RIBvLa58AN4Pi7Fp9hFxGtjwzIpWUXIUr-BGE_9SLvjUGgsQCyrhK9ZJ5Yy9R5F6w4Me0Csr19UU3IqQQIP3ffhInE5o68_CI_URCjHXpBUnztJoDmlBnZz3Ka5NykfUN22iulaFvXOyw=='
print(f"The secret message is {secret_message.decode()}")

num_qubits = 64

rng = np.random.default_rng(seed=10)

alice_state = rng.integers(0, 2, size=num_qubits)
bob_basis = rng.integers(0, 2, size=num_qubits)

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

def make_b92_circ(enc_state, meas_basis):
    '''
    A function that makes a B92 QKD protocol simulation circuit
    
    enc_state: array of 0s and 1s denoting the state to be encoded using the following scheme:
                0 -> |0>
                1 -> |+>
    
    meas_basis: array of 0s and 1s denoting the basis to be used for measurement
                0 -> Hadamard Basis
                1 -> Computational Basis
    
    Note that both enc_state and meas_basis are arrays of integers, so if you are using them in 
    if statements, compare them to integer values like 0 and 1 (without quotes).
    
    Since this is a function, you only have access to the variables enc_state and meas_basis.
    You may define other local variables. One such variable, num_qubits has been defined for you. 
    This is the number of qubits in the B92 simulation QuantumCircuit()
    '''
    num_qubits = len(enc_state)

    b92 = QuantumCircuit(num_qubits)

    # Sender prepares qubits
    # Add code below to encode the state in qubits

    
    b92.barrier()  
    # Receiver measures the received qubits
    # Add code below to change basis for measurements. DO NOT add a measure() or measure_all()

    # Do not change below this line
    b92.measure_all()
    return b92

try:
    b92_circ = make_b92_circ(alice_state, bob_basis)
    assert list(b92_circ.count_ops()) != [], "Circuit cannot be empty"
    assert set(b92_circ.count_ops().keys()).difference(basis_gates) == set(), f"Only the following basic gates are allowed: {basis_gates}"
    assert all([type(gate[0]) == qiskit.circuit.measure.Measure for gate in b92_circ.data[-b92_circ.num_qubits:len(b92_circ.data)]]), "Measurement must be the last operation in a circuit."
    assert b92_circ.count_ops()['measure'] == b92_circ.num_qubits, "Please do not add or remove measurements."
    temp_key = execute(
        b92_circ.reverse_bits(),
        backend=QasmSimulator(),
        shots=1, 
        seed_simulator=10
    ).result().get_counts().most_frequent()
    assert temp_key == bin(16228741048440553634)[2:], "Your circuit did not perform as expected. Please check the gates again."
    print(f"Bob's results:\t{temp_key}\nYour answer is correct.")
except AssertionError as e:
    print(f'Your code has an error:  {e.args[0]}')
except Exception as e:
    print(f'This error occured: {e.args[0]}')

def b92_sifting(enc_state, meas_basis, meas_result):
    '''
    The function that implements key sifting for the B92 QKD protocol.
    
    enc_state:  array of 0s and 1s denoting the state to be encoded.
                (Array of integers)
    
    meas_basis: array of 0s and 1s denoting the basis to be used for measurement.
                (Array of integers)
    
    meas_result: A string of characters representing the results of measurement after the 
                B92 QKD protocol. Note that this is a string and its elements are characters, 
                so while using any if statements, compare the elements to '0' and '1' (with quotes)
                
    Since this is a function, you only have access to the variables enc_state, meas_basis and meas_result.
    You may define other local variables. num_qubits has been defined for you. 
    This is the number of qubits in the B92 simulation QuantumCircuit.
    
    sender_key and receiver_key are initialised as two empty strings. You may append bits using the += 
    operation as shown in the BB84 notebook. Note that you can only add characters. To change from other 
    data types to characters, you may use str(). Check the BB84 notebook for examples. 
    '''
    
    num_qubits = len(enc_state)

    sender_key = ''
    receiver_key = ''
    
    # Loop over all bits in the meas_result string and add the necessary bits to both sender_key and receiver_key
    # Add your code below

    # Do not change bolow this line. 
    return (sender_key, receiver_key)

try: 
    alice_key, bob_key = b92_sifting(alice_state, bob_basis, temp_key)
    assert ''.join([str(x ^ y) for x, y in zip(alice_key.encode(), bob_key.encode())]) != '1'*len(alice_key), "Please check your measurement convention"
    assert alice_key == bob_key, "They keys are different for Alice and Bob."
    assert alice_key == bob_key == bin(49522)[2:], "They keys is incorrect. Please check your solutions."
    print(f"Alice's Key: \t{alice_key}\nBob's Key: \t{bob_key}\nYour answer is correct.")
    g = Fernet(base64.b64encode(bob_key.encode()*2))
    print(f"The secret message is: {g.decrypt(secret_message).decode()}")
except AssertionError as e:
    print(f'Your code has an error:  {e.args[0]}')
except Exception as e:
    print(f'This error occured: {e.args[0]}')    