Demonstrator BB84.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d2ee0436",
   "metadata": {},
   "source": [
    "Author: Jakub Gnyp; contact: gnyp.jakub@gmail.com, LinkedIn: https://www.linkedin.com/in/gnypit/\n",
    "Functions and constants for hashing based on FIPS PUB 180-4 from:\n",
    "http://dx.doi.org/10.6028/NIST.FIPS.180-4\n",
    "\n",
    "For a summary of applied knowledge I refer to \"Applied Quantum Cryptography\" (Lecture notes in physics 797) published by Springer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "295a5654",
   "metadata": {},
   "source": [
    "import ast\n",
    "import random\n",
    "import math\n",
    "import numpy as np\n",
    "from scipy.stats import binom\n",
    "\n",
    "random.seed(a=261135)  # useful for debugging/presentations, should be omitted for practical uses"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "a51d26c1",
   "metadata": {},
   "source": [
    "# Defining functions for privacy amplification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "cd8e5c3b",
   "metadata": {},
   "source": [
    "def constant_sha1(index):\n",
    "    if 0 <= index <= 19:\n",
    "        return 0x5a827999\n",
    "    elif 20 <= index <= 39:\n",
    "        return 0x6eD9eba1\n",
    "    elif 40 <= index <= 59:\n",
    "        return 0x8f1bbcdc\n",
    "    elif 60 <= index <= 79:\n",
    "        return 0xca62c1d6\n",
    "    else:\n",
    "        return 'Error: index exceeds allowed range.'\n",
    "\n",
    "\n",
    "def ch(x, y, z):  # one of internal hashing functions for SHA-1\n",
    "    return (x & y) ^ (~x & z)\n",
    "\n",
    "\n",
    "def parity(x, y, z):  # one of internal hashing functions for SHA-1\n",
    "    return x ^ y ^ z\n",
    "\n",
    "\n",
    "def maj(x, y, z):  # one of internal hashing functions for SHA-1\n",
    "    return (x & y) ^ (x & z) ^ (y & z)\n",
    "\n",
    "\n",
    "def hashing_function_sha1(x, y, z, j):\n",
    "    if 0 <= j <= 19:\n",
    "        return ch(x=x, y=y, z=z)\n",
    "    elif 20 <= j <= 39 or 60 <= j <= 79:\n",
    "        return parity(x=x, y=y, z=z)\n",
    "    elif 40 <= j <= 59:\n",
    "        return maj(x=x, y=y, z=z)\n",
    "    else:\n",
    "        return 'Error: j exceeds possible range.'\n",
    "\n",
    "\n",
    "def rotl(n, n_bits, word):  # word will be given as an int, so we will return it as an int too\n",
    "    \"\"\"python's left shift doesn't discard any bits, just adds padding - we will discard bits manually\"\"\"\n",
    "    first = bin(word << n)[2 + n:n_bits + n + 2]  # bin changes type to str\n",
    "    first = int('0b' + first, base=2)\n",
    "\n",
    "    \"\"\"python's right shift only discards the bits - we will add the padding manually\"\"\"\n",
    "    second = bin(word >> n_bits - n)[2:n + 2]  # bin changes type to str\n",
    "    second = '0' * (n_bits - n) + second\n",
    "    second = int('0b' + second, base=2)\n",
    "\n",
    "    \"\"\"Result is a bitwise OR on first and second\"\"\"\n",
    "    result = first | second  # result = (word << n) | (word >> (n_bits - n)) with proper discarding and padding\n",
    "    return result\n",
    "\n",
    "\n",
    "def sha1(message):\n",
    "    # message = str(input('Please give a message:'))  # should be a series of 0's and 1's\n",
    "\n",
    "    \"\"\"SHA-1: firstly we parse the message into equally long substrings - blocks:\"\"\"\n",
    "    message_length, blocks_length = len(message), 512\n",
    "    blocks = [message[i:i + blocks_length] for i in range(0, message_length, blocks_length)]\n",
    "    if len(blocks[-1]) < blocks_length:  # in case the last block is too short, we add padding with 0's\n",
    "        padding_length = blocks_length - len(blocks[-1])\n",
    "        padding = '0' * padding_length\n",
    "        blocks[-1] = padding + blocks[-1]\n",
    "\n",
    "    \"\"\"Now for the actual hash computation. We will use:\n",
    "    words -> groups 32 bits\n",
    "    hashing function -> as defined above for SHA-1\n",
    "    ROTL function -> as defined above\n",
    "    ...to create a 'message schedule' for every block of parsed message from the list 'blocks':\n",
    "    \"\"\"\n",
    "\n",
    "    \"\"\"SHA-1: setting initial hash value\"\"\"\n",
    "    hash_value = [\n",
    "        0x67452301,\n",
    "        0xefcdab89,\n",
    "        0x98badcfe,\n",
    "        0x10325476,\n",
    "        0xc3d2e1f0\n",
    "    ]\n",
    "\n",
    "    for block in blocks:\n",
    "        \"\"\"Firstly we prepare the message schedule:\"\"\"\n",
    "        for t in range(0, 80, 1):\n",
    "            \"\"\"First 16 words of the message schedule are words of the block of parsed message which we are working on\n",
    "            in this iteration of the loop.\n",
    "            \"\"\"\n",
    "            message_schedule = [block[i:i + 32] for i in range(0, 512, 32)]\n",
    "\n",
    "            \"\"\"Last 64 words of the message schedule are created based on a recurrent formula using ROTL function:\"\"\"\n",
    "            for i in range(16, 80, 1):\n",
    "                \"\"\"We locally represent selected previous words as binary values\"\"\"\n",
    "                schedule_i3 = int('0b' + message_schedule[i - 3], base=2)\n",
    "                schedule_i8 = int('0b' + message_schedule[i - 8], base=2)\n",
    "                schedule_i14 = int('0b' + message_schedule[i - 14], base=2)\n",
    "                schedule_i16 = int('0b' + message_schedule[i - 16], base=2)\n",
    "\n",
    "                \"\"\"Next word of the message schedule is value of ROTL on above words:\"\"\"\n",
    "                next_word = bin(rotl(\n",
    "                    n=1,\n",
    "                    n_bits=32,\n",
    "                    word=schedule_i3 ^ schedule_i8 ^ schedule_i14 ^ schedule_i16\n",
    "                ))[2:32]\n",
    "                message_schedule.append(next_word)\n",
    "\n",
    "            \"\"\"We initialize the five working variables with last (possibly initial) hash value\"\"\"\n",
    "            a = hash_value[0]\n",
    "            b = hash_value[1]\n",
    "            c = hash_value[2]\n",
    "            d = hash_value[3]\n",
    "            e = hash_value[4]\n",
    "\n",
    "            \"\"\"We process the working variables...\"\"\"\n",
    "            for i in range(0, 80, 1):\n",
    "                t = rotl(n=5, n_bits=32, word=a) + hashing_function_sha1(x=b, y=c, z=d, j=i) + e + constant_sha1(index=i) \\\n",
    "                    + int('0b' + message_schedule[i], base=2)\n",
    "                e = d\n",
    "                d = c\n",
    "                c = rotl(n=30, n_bits=32, word=b)\n",
    "                b = a\n",
    "                a = t\n",
    "\n",
    "            \"\"\"...and compute the next hash value.\n",
    "            As we always need just the last hash value, we won't store all of them in memory; instead we will\n",
    "            save each next one in the list hash_value, which decomposes such a value\n",
    "            into its' five words - the list elements. \n",
    "            \"\"\"\n",
    "            hash_value[0] = hash_value[0] + a\n",
    "            hash_value[1] = hash_value[1] + b\n",
    "            hash_value[2] = hash_value[2] + c\n",
    "            hash_value[3] = hash_value[3] + d\n",
    "            hash_value[4] = hash_value[4] + e\n",
    "\n",
    "    \"\"\"After the loop the resulting 160-bit message digest of the original message is:\"\"\"\n",
    "    digest = ''\n",
    "    for value in hash_value:\n",
    "        word = bin(value)\n",
    "        digest += word[2:len(word)]\n",
    "\n",
    "    return digest"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "df2b55fa",
   "metadata": {},
   "source": [
    "# Defining function for BINARY:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "c588c6bf",
   "metadata": {},
   "source": [
    "def binary(sender_block, receiver_block, indexes, receiver_name='Bob'):\n",
    "    is_binary = True\n",
    "    sender_current_block = sender_block\n",
    "    receiver_current_block = receiver_block\n",
    "    indexes = list(indexes)\n",
    "\n",
    "    \"\"\"As in real-life applications, in this simulation of the BINARY algorithm Alice (sender) and Bob (receiver) \n",
    "    exchange messages. Additionally, we count how many bits are exchanged between them for the future\n",
    "    privacy amplification. Afterwards we return this value together with the bit to be changed as a result of the\n",
    "    algorithm.\n",
    "    \"\"\"\n",
    "    bit_counter = 0\n",
    "\n",
    "    while is_binary:\n",
    "        \"\"\"Sender starts by sending to the Receiver parity of the first half of her string\"\"\"\n",
    "        half_index = len(sender_current_block.values()) // 2  # same as Bob's\n",
    "        first_half_indexes = indexes[0:half_index:1]  # same as Bob's\n",
    "        sender_first_half_list = []\n",
    "\n",
    "        for index in first_half_indexes:\n",
    "            sender_first_half_list.append(int(sender_current_block.get(index)))\n",
    "\n",
    "        sender_first_half_parity = sum(sender_first_half_list) % 2\n",
    "        print(\"[Alice] My string's first half has a parity: {}\".format(sender_first_half_parity))\n",
    "        bit_counter += 1  # At this point sender informs receiver about their 1st half's parity\n",
    "\n",
    "        \"\"\"Now Receiver determines whether an odd number of errors occurred in the first or in the\n",
    "        second half by testing the parity of his string and comparing it to the parity sent\n",
    "        by Sender\n",
    "        \"\"\"\n",
    "\n",
    "        receiver_first_half_list = []\n",
    "\n",
    "        for index in first_half_indexes:\n",
    "            receiver_first_half_list.append(int(receiver_current_block.get(index)))\n",
    "\n",
    "        receiver_first_half_parity = sum(receiver_first_half_list) % 2\n",
    "\n",
    "        \"\"\"Single (at least) error is in the 'half' of a different parity; we change current strings\n",
    "        that are analysed into halves of different parities until one bit is left - the error\n",
    "        \"\"\"\n",
    "\n",
    "        if receiver_first_half_parity != sender_first_half_parity:\n",
    "            print('[{}] I have an odd number of errors in my first half.'.format(\n",
    "                receiver_name\n",
    "            ))\n",
    "            bit_counter += 1\n",
    "\n",
    "            sender_subscription_block = {}\n",
    "            receiver_subscription_block = {}\n",
    "\n",
    "            for index in first_half_indexes:\n",
    "                receiver_subscription_block[index] = receiver_current_block.get(index)\n",
    "                sender_subscription_block[index] = sender_current_block.get(index)\n",
    "\n",
    "            sender_current_block = sender_subscription_block\n",
    "            receiver_current_block = receiver_subscription_block\n",
    "\n",
    "            indexes = list(sender_current_block.keys())  # same as Bob's\n",
    "            \n",
    "            print('Alice and Bob updated their strings.')\n",
    "        else:\n",
    "            print('[{}] I have an odd number of errors in my second half.'.format(\n",
    "                receiver_name\n",
    "            ))\n",
    "            bit_counter += 1\n",
    "\n",
    "            \"\"\"We have to repeat the whole procedure for the second halves\"\"\"\n",
    "            second_half_indexes = indexes[half_index::1]\n",
    "            sender_subscription_block = {}\n",
    "            receiver_subscription_block = {}\n",
    "\n",
    "            for index in second_half_indexes:\n",
    "                receiver_subscription_block[index] = receiver_current_block.get(index)\n",
    "                sender_subscription_block[index] = sender_current_block.get(index)\n",
    "\n",
    "            sender_current_block = sender_subscription_block\n",
    "            receiver_current_block = receiver_subscription_block\n",
    "\n",
    "            indexes = list(sender_current_block.keys())  # same as Bob's\n",
    "\n",
    "            print('Alice and Bob updated their strings.')\n",
    "\n",
    "        if len(receiver_current_block) == 1:  # at some point this clause will be true\n",
    "            print(\"[{}] I have one bit left, I'm changing it.\".format(\n",
    "                receiver_name\n",
    "            ))\n",
    "            bit_counter += 1  # At this point receiver would send a message (?) about one bit left and changing it\n",
    "\n",
    "            \"\"\"We return the correct value of the bit which was erroneous, together with it's index in the original Bob's\n",
    "            bits string & the number of bits exchanged during BINARY.\n",
    "            \"\"\"\n",
    "            if receiver_current_block[indexes[0]] == '0':\n",
    "                return ['1', indexes[0], bit_counter]\n",
    "            else:\n",
    "                return ['0', indexes[0], bit_counter]"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "9607ef77",
   "metadata": {},
   "source": [
    "# Performing BB84"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b38b503",
   "metadata": {},
   "source": [
    "Let's set up the quantum channel (BB84):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "3aa02bf8",
   "metadata": {},
   "source": [
    "basis_mapping = {'rectilinear': 0, 'diagonal': 1}\n",
    "states_mapping = {'0': 0, '1': 1, '+': 0, '-': 1}\n",
    "quantum_channel = {\n",
    "    'rectilinear': {\n",
    "        'basis_vectors': {'first_state': '0', 'second_state': '1'}\n",
    "    },\n",
    "    'diagonal': {\n",
    "        'basis_vectors': {'first_state': '+', 'second_state': '-'}\n",
    "    }\n",
    "}"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "9689e409",
   "metadata": {},
   "source": [
    "## Let's define some useful functions for the protocol to run:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "25226388",
   "metadata": {},
   "source": [
    "def qc_gain(mean_photon_number=1., fiber_loss=1., detection_efficiency=1., k_dead=1.,\n",
    "            additional_loss=1.):  # quantum channel gain -> do lamusa\n",
    "    g = mean_photon_number * fiber_loss * detection_efficiency * k_dead * additional_loss\n",
    "    return g\n",
    "\n",
    "\n",
    "def received_key_material(quantum_channel_gain, sender_data_rate):\n",
    "    receiver = quantum_channel_gain * sender_data_rate\n",
    "    return receiver\n",
    "\n",
    "\n",
    "def random_choice(length, p=0.5):  # function for random choosing of basis for each photon\n",
    "    \"\"\"p -> probability of selecting rectilinear basis\"\"\"\n",
    "    chosen_basis = ''\n",
    "    for index in range(int(np.floor(length))):\n",
    "        basis = random.uniform(0, 1)\n",
    "        if basis <= p:\n",
    "            chosen_basis += str(0)\n",
    "        else:\n",
    "            chosen_basis += str(1)\n",
    "\n",
    "    return chosen_basis\n",
    "\n",
    "\n",
    "def measurement(state, basis):  # meant for classical encoding\n",
    "\n",
    "    if basis == '1':  # meaning diagonal basis\n",
    "        basis = 'diagonal'\n",
    "\n",
    "        if state == '+':\n",
    "            final_state = '+'\n",
    "        elif state == '-':\n",
    "            final_state = '-'\n",
    "        elif state == '0' or state == '1':  # in this case there's a 50% chance of getting either polarization\n",
    "            if random.randint(0, 1) == 0:\n",
    "                final_state = quantum_channel.get(basis).get('basis_vectors').get('first_state')\n",
    "            else:\n",
    "                final_state = quantum_channel.get(basis).get('basis_vectors').get('second_state')\n",
    "        else:\n",
    "            return 'U'  # U for unknown\n",
    "\n",
    "    elif basis == '0':  # meaning rectilinear basis\n",
    "\n",
    "        if state == '0':\n",
    "            final_state = '0'\n",
    "        elif state == '1':\n",
    "            final_state = '1'\n",
    "        elif state == '+' or state == '-':\n",
    "            final_state = str(random.randint(0, 1))  # since '0' and '1' are states, there's no need for if...else\n",
    "        else:\n",
    "            return 'U'  # U for unknown\n",
    "\n",
    "    elif basis == 'L':\n",
    "        final_state = 'L'  # L for loss, as basis L reflects unperformed measurement due to quantum channel loss\n",
    "    else:\n",
    "        return 'U'  # U for unknown\n",
    "\n",
    "    return final_state\n",
    "\n",
    "\n",
    "def numerical_error_prob(n_errors, pass_size, qber):  # probability that 2*n_errors remain\n",
    "    prob = binom.pmf(2 * n_errors, pass_size, qber) + binom.pmf(2 * n_errors + 1, pass_size, qber)\n",
    "    return prob\n",
    "\n",
    "\n",
    "def cascade_blocks_sizes(quantum_bit_error_rate, key_length, n_passes=1):\n",
    "    print('--------------------\\nCOMPUTING SIZES OF BLOCKS FOR CASCADE')\n",
    "    max_expected_value = -1 * math.log(0.5, math.e)\n",
    "    # best_expected_value = max_expected_value\n",
    "    best_size = 0\n",
    "    print(\"Currently the greatest size of block fulfilling both criteria is {}.\".format(best_size))\n",
    "\n",
    "    for size in range(key_length // 2):  # we need at lest 2 blocks to begin with\n",
    "        # Firstly we check condition for expected values\n",
    "        expected_value = 0\n",
    "\n",
    "        for j in range(size // 2):\n",
    "            expected_value += 2 * (j + 1) * numerical_error_prob(n_errors=(j + 1), pass_size=size,\n",
    "                                                                 qber=quantum_bit_error_rate)\n",
    "\n",
    "        if expected_value <= max_expected_value:\n",
    "            first_condition = True\n",
    "        else:\n",
    "            first_condition = False\n",
    "\n",
    "        # Secondly we check condition for probabilities per se\n",
    "        second_condition = False\n",
    "        for j in range(size // 2):\n",
    "            prob_sum = 0\n",
    "            for k in list(np.arange(j + 1, size // 2 + 1, 1)):\n",
    "                prob_sum += numerical_error_prob(n_errors=k, pass_size=size, qber=quantum_bit_error_rate)\n",
    "\n",
    "            if prob_sum <= numerical_error_prob(n_errors=j, pass_size=size, qber=quantum_bit_error_rate) / 4:\n",
    "                second_condition = True\n",
    "            else:\n",
    "                second_condition = False\n",
    "\n",
    "        if first_condition and second_condition:\n",
    "            if size > best_size:\n",
    "                # best_expected_value = expected_value\n",
    "                best_size = size\n",
    "                print(\"Currently the greatest size of block fulfilling both criteria is {}.\".format(best_size))\n",
    "\n",
    "    sizes = [best_size]\n",
    "\n",
    "    for j in range(n_passes - 1):  # corrected interpretation of number of passes\n",
    "        next_size = 2 * sizes[-1]\n",
    "        if next_size <= key_length:\n",
    "            sizes.append(next_size)\n",
    "        else:\n",
    "            break\n",
    "\n",
    "    print(\"Since the best initial size of blocks equals {} and we are ment to perform {} passes of the algorithm,\"\n",
    "          \"sizes of blocks for consecutive passes are: {}\".format(best_size, n_passes, sizes))\n",
    "\n",
    "    return sizes\n",
    "\n",
    "\n",
    "def cascade_blocks_generator(string_length, blocks_size):\n",
    "    string_index = list(np.arange(0, string_length, 1))  # I create a list of all indexes (list of ndarray)\n",
    "    blocks = random.sample(population=string_index, k=string_length)  # I shuffle the list randomly\n",
    "\n",
    "    for j in range(0, string_length, blocks_size):  # I generate equally long chunks of shuffled indexes\n",
    "        yield blocks[j:j + blocks_size]\n",
    "\n",
    "\n",
    "def naive_error(alice_key, bob_key, publication_prob_rect):\n",
    "    \"\"\"\n",
    "    Originally I've been testing algorithm on short strings, i.e. up to 20 bits in original Alice's message.\n",
    "    Sifted keys were therefore too short for publication of any parts of them for error estimation.\n",
    "\n",
    "    A refined estimation is based on two subsets: one of the sifted key and the other on the bits originating from\n",
    "    measurements with different basis (Alice's and Bob's). First we present the 'naive' estimation.\n",
    "\n",
    "    Let's randomly publish subsets of bits of matching positions in the sifted keys strings. Then we count how many\n",
    "    bits differ between these two strings, divide this by their length and get a naive estimator of the total error\n",
    "    rate (between error correction algorithm, i.e. CASCADE).\n",
    "\n",
    "    As a result of this function we return the naive estimator and Alice's & Bob's keys without published bits.\n",
    "    \"\"\"\n",
    "    print('----------------------------------------\\nPERFORMING NAIVE ERROR ESTIMATION\\n')\n",
    "    alice_published_bits = ''\n",
    "    alice_sifted_key_after_publication = ''\n",
    "    bob_published_bits = ''\n",
    "    bob_sifted_key_after_publication = ''\n",
    "\n",
    "    naive_error_estimate = 0\n",
    "\n",
    "    for index in range(len(alice_key)):  # could be bob_key as well\n",
    "        if random.uniform(0, 1) <= publication_prob_rect:\n",
    "            alice_bit = alice_key[index]\n",
    "            bob_bit = bob_key[index]\n",
    "\n",
    "            \"\"\"First we add those bits to strings meant for publication\"\"\"\n",
    "            alice_published_bits += alice_bit\n",
    "            bob_published_bits += bob_bit\n",
    "\n",
    "            \"\"\"Now for the estimation of the error:\"\"\"\n",
    "            if alice_bit != bob_bit:\n",
    "                naive_error_estimate += 1\n",
    "        else:  # if a bit wasn't published, we reuse it in the sifted key\n",
    "            alice_sifted_key_after_publication += alice_key[index]\n",
    "            bob_sifted_key_after_publication += bob_key[index]\n",
    "\n",
    "    print('[Alice] My subset of bits for error estimation is: {}'.format(alice_published_bits))\n",
    "    print('[Bob] My subset of bits for error estimation is: {}'.format(bob_published_bits))\n",
    "\n",
    "    try:\n",
    "        naive_error_estimate = naive_error_estimate / len(alice_published_bits)\n",
    "    except ZeroDivisionError:\n",
    "        naive_error_estimate = 0  # this will obviously be false, but easy to notice and work on in genalqkd.py\n",
    "\n",
    "    \"\"\"At this point we count the number of bits exchanged between Alice and Bob via the public channel,\n",
    "    for future estimation of the computational cost.\"\"\"\n",
    "\n",
    "    no_published_bits = len(alice_published_bits) + len(bob_published_bits)\n",
    "    print('The probability of bit publishing was {}. In actuality {} bits were published out of {} bits succesfully '\n",
    "          'sent & received.'.format(publication_prob_rect, no_published_bits, len(alice_key)))\n",
    "\n",
    "    results = {\n",
    "        'error estimator': naive_error_estimate,\n",
    "        'alice key': alice_sifted_key_after_publication,\n",
    "        'bob key': bob_sifted_key_after_publication,\n",
    "        'number of published bits': no_published_bits\n",
    "    }\n",
    "\n",
    "    return results\n",
    "\n",
    "\n",
    "def refined_average_error(rect_prob, rect_pub_prob, diag_pub_prob, alice_key, bob_key, alice_bases, bob_bases):\n",
    "    \"\"\"In the refined error analysis for simplicity we DO NOT divide raw keys into two separate strings (by the basis).\n",
    "    Instead, we create two empty strings - alice_key_after_error_estimation & bob_key_after_error_estimation - into\n",
    "    which we shall rewrite bits unused for error estimation. As for the others, chosen with probability\n",
    "    rect_pub_prob & diag_pub_prob, respectively, we count them as 'published' and additionally count an error,\n",
    "    if they differ from each other. Those counter are:\n",
    "    rect_pub_counter & diag_pub_counter, rect_error & diag_error.\n",
    "\n",
    "    The last two will be divided at the end by the first two, respectively, to obtain estimations as ratios.\n",
    "    \"\"\"\n",
    "    print('----------------------------------------\\nPERFORMING REFINED ERROR ESTIMATION\\n')\n",
    "    alice_key_after_error_estimation = ''\n",
    "    bob_key_after_error_estimation = ''\n",
    "    rect_error = 0\n",
    "    rect_pub_counter = 0\n",
    "    diag_error = 0\n",
    "    diag_pub_counter = 0\n",
    "    alice_published_bits_rect = ''\n",
    "    bob_published_bits_rect = ''\n",
    "    alice_published_bits_diag = ''\n",
    "    bob_published_bits_diag = ''\n",
    "\n",
    "    \"\"\"By the way, we don't have to worry with the conditional probability, because the final formula for the error\n",
    "    estimation takes that into consideration.\n",
    "    \"\"\"\n",
    "    for index in range(len(alice_key)):\n",
    "        if alice_bases[index] == bob_bases[index] == '0':  # rectilinear basis\n",
    "            if random.uniform(0, 1) >= rect_pub_prob:\n",
    "                alice_key_after_error_estimation += alice_key[index]\n",
    "                bob_key_after_error_estimation += bob_key[index]\n",
    "            else:\n",
    "                rect_pub_counter += 1\n",
    "                if alice_key[index] != bob_key[index]:\n",
    "                    rect_error += 1\n",
    "                alice_published_bits_rect += alice_key[index]\n",
    "                bob_published_bits_rect += bob_key[index]\n",
    "        else:\n",
    "            if random.uniform(0, 1) >= diag_pub_prob:\n",
    "                alice_key_after_error_estimation += alice_key[index]\n",
    "                bob_key_after_error_estimation += bob_key[index]\n",
    "            else:\n",
    "                diag_pub_counter += 1\n",
    "                if alice_key[index] != bob_key[index]:\n",
    "                    diag_error += 1\n",
    "                alice_published_bits_diag += alice_key[index]\n",
    "                bob_published_bits_diag += bob_key[index]\n",
    "\n",
    "    \"\"\"As it is possible to get a VERY small probability of publication, we check for possible divisions by zero:\"\"\"\n",
    "    try:\n",
    "        rect_error = float(rect_error) / float(rect_pub_counter)\n",
    "    except ZeroDivisionError:\n",
    "        rect_error = 0.0\n",
    "\n",
    "    try:\n",
    "        diag_error = float(diag_error) / float(diag_pub_counter)\n",
    "    except ZeroDivisionError:\n",
    "        diag_error = 0.0\n",
    "\n",
    "    print('[Alice] My subset of bits measured in the rectilinear basis is: {}'.format(alice_published_bits_rect))\n",
    "    print('[Bob] My subset of bits measured in the rectilinear basis is: {}'.format(bob_published_bits_rect))\n",
    "    print('The error rate estimate for the rectilinear basis is {}.\\n'.format(rect_error))\n",
    "\n",
    "    print('[Alice] My subset of bits measured in the diagonal basis is: {}'.format(alice_published_bits_diag))\n",
    "    print('[Bob] My subset of bits measured in the diagonal basis is: {}'.format(bob_published_bits_diag))\n",
    "    print('The error rate estimate for the diagonal basis is {}.\\n'.format(diag_error))\n",
    "\n",
    "    \"\"\"Now, given that measurements in the rectilinear basis were not necessarily with the same probability \n",
    "    as those in the diagonal basis, we need a more complicated formula for the 'average error estimate' \n",
    "    (Lo, Chau, Ardehali, 2004).\n",
    "    \"\"\"\n",
    "    p = rect_prob  # just a reminder that it's the probability of choosing rect. basis for measurements\n",
    "    e1 = rect_error\n",
    "    e2 = diag_error\n",
    "\n",
    "    average_error = (p ** 2 * e1 + (1 - p) ** 2 * e2) / (p ** 2 + (1 - p) ** 2)\n",
    "    print('The average error estimate is {}.'.format(average_error))\n",
    "    e = (e1 + e2) / 2\n",
    "    print('With the refined error analysis, the actual QBER is {}'.format(e))\n",
    "\n",
    "    results = {\n",
    "        'error estimator': e,\n",
    "        'alice key': alice_key_after_error_estimation,\n",
    "        'bob key': bob_key_after_error_estimation,\n",
    "        'number of published bits': rect_pub_counter + diag_pub_counter\n",
    "    }\n",
    "\n",
    "    return results"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "da8c7c75",
   "metadata": {},
   "source": [
    "## Start of the protocol (input section)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b4fcd7a",
   "metadata": {},
   "source": [
    "At the begining we define the inputs, namely: quantum gain of a quantum channel used by Alice and Bob, Alice's basis length, etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "434b9b4f",
   "metadata": {},
   "source": [
    "\"\"\"Please enter a numerical value between 0 and 1, e.g. 0.95:\"\"\"\n",
    "gain = 0.95\n",
    "\n",
    "\"\"\"Alice will choose her basis randomly. How long should the string with her choices be?\"\"\"\n",
    "alice_basis_length = 512\n",
    "\n",
    "\"\"\"Please set the probability of choosing a rectilinear basis:\"\"\"\n",
    "rectilinear_basis_prob = 0.75\n",
    "\n",
    "\"\"\"There's a probability that due to disturbances in quantum channel or eavesdropping some bits change while being sent\"\"\"\n",
    "change_probability = 0.05\n",
    "\n",
    "\"\"\"Please set the number of CASCADE passes to perform:\"\"\"\n",
    "cascade_passes = 4"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "e8437c3b",
   "metadata": {},
   "source": [
    "At this point the quantum channel is set up and we know how Alice prepares the measurements.\n",
    "It's time to perform sending states by Alice. She will choose tem randomly, with unifrom probability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "692a8a07",
   "metadata": {},
   "source": [
    "if alice_basis_length != 0:  # firstly it was 0, if now it's !=0, then we are supposed to choose randomly her basis\n",
    "    alice_basis = random_choice(alice_basis_length, p=rectilinear_basis_prob)\n",
    "    print('Random choices of basis for Alice are: {}'.format(alice_basis))\n",
    "\n",
    "alice_bits_length = 0\n",
    "if alice_bits_length == 0:  # if it's 0, then we didn't get the bits form the user\n",
    "    alice_bits_length = alice_basis_length\n",
    "    alice_bits = random_choice(alice_bits_length, p=0.5)\n",
    "    print('Random choices of bits for Alice are: {}'.format(alice_bits))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "d6da0dc6",
   "metadata": {},
   "source": [
    "At this point it is impractical to encode states as bits because Bob's measurements results depend on both basis\n",
    "and bit choice of Alice, but he shouldn't know the first one. Because of that we will now translate Alice's bits\n",
    "to proper states, changing 0 and 1 into + and - for the diagonal basis, respectively.\n",
    "\n",
    "Please remember that we do that in order to demonstrate what exactly is going on - in practice a measurement is \n",
    "performed, results are remembered and post-processed and they are not being published, in contrary to us printing\n",
    "all of the information ;)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "2ac3872c",
   "metadata": {},
   "source": [
    "i = 0\n",
    "alice_states_list = list(alice_bits)\n",
    "for bit in alice_bits:\n",
    "    if alice_basis[i] == '1':\n",
    "        if bit == '0':\n",
    "            alice_states_list[i] = '+'\n",
    "        elif bit == '1':\n",
    "            alice_states_list[i] = '-'\n",
    "        else:\n",
    "            alice_states_list[i] = 'U'  # U for unknown\n",
    "    i += 1\n",
    "\n",
    "alice_states = ''.join(alice_states_list)\n",
    "print(\"Alice's states are: {}, where U stands for unknown states due to incorrect encoding of bits.\".format(\n",
    "    alice_states\n",
    "))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "839edef9",
   "metadata": {},
   "source": [
    "Now that we have Alice's data rate, quantum channel's gain and Alice's states,\n",
    "we can randomly choose m (Alice's basis choices number) bases for Bob. While he performs his measurements,\n",
    "a portion of Alice's photons do not reach him due to loss on quantum channel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "ec903fde",
   "metadata": {},
   "source": [
    "bob_basis_list = random_choice(alice_basis_length, p=rectilinear_basis_prob)\n",
    "bob_basis = ''\n",
    "for basis in bob_basis_list:\n",
    "    if random.uniform(0, 1) <= gain:  # for small numbers of bits quantum gain won't change anything in here\n",
    "        bob_basis += str(int(basis))\n",
    "    else:\n",
    "        bob_basis += 'L'  # L for loss\n",
    "\n",
    "print(\"Bob's basis choices are: {}, where L stands for measurement that cannot be done\".format(bob_basis),\n",
    "      \"\\ndue to quantum channel loss.\")"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "d96c3b6d",
   "metadata": {},
   "source": [
    "Now we can perform disturbances in the quantum channel:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "0ec60dec",
   "metadata": {},
   "source": [
    "received_states = ''\n",
    "change_states = {'0': '0', '1': '1', '2': '+', '3': '-'}  # dictionary for randomizing changed states\n",
    "for state in alice_states:  # to be updated following concrete attack strategies\n",
    "    if random.uniform(0, 1) <= change_probability:\n",
    "        change_indicator = str(random.randint(0, 3))\n",
    "        while state == change_states.get(change_indicator):  # we repeat as long as it takes to actually change state\n",
    "            change_indicator = str(random.randint(0, 3))\n",
    "        received_states += change_states.get(change_indicator)\n",
    "    else:\n",
    "        received_states += state\n",
    "\n",
    "print(\"States received by Bob before measurement are: {}\".format(received_states))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "7dea2408",
   "metadata": {},
   "source": [
    "After Bob chose a basis and a photon has reached him, he performs his measurement. If for this particular photon \n",
    "he chooses the same basis as Alice chose before, his measurement result will be the same - that's because\n",
    "in reality Bob is choosing polarisators for photons to go through. If photon's polarization is the same as\n",
    "polarisators, then there's 100% probability of preserving this polarization. This is reflected by the same bit\n",
    "as before the measurement. Otherwise polarization after measurement is random. This mechanism is implemented in\n",
    "\"measurement\" function.\n",
    "\n",
    "Results of Bob's measurements are states - either 0 for |0> or 1 for |1> or + for |+> or - for |->. If he couldn't\n",
    "perform the measurement (encoded by basis L) then the result will be encoded in the string by L as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "66bd6dcf",
   "metadata": {},
   "source": [
    "bob_states = ''\n",
    "for i in range(alice_basis_length):  # it's the same length as of Bob's basis choices\n",
    "    bob_states += measurement(state=received_states[i], basis=bob_basis[i])\n",
    "\n",
    "print(\"Bob's measurement results are: {}\".format(bob_states))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "4ee1c62e",
   "metadata": {},
   "source": [
    "Now, having Bob's measurement results, we can translate states into bits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "37667266",
   "metadata": {},
   "source": [
    "bob_bits = ''\n",
    "for state in bob_states:\n",
    "    try:\n",
    "        bob_bits += str(states_mapping[state])  # we want to use indexing to raise errors for unsuccessful measur.\n",
    "    except KeyError:\n",
    "        bob_bits += random_choice(length=1)\n",
    "        continue\n",
    "\n",
    "print(\"Bob's bits are: {}\".format(bob_bits))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "087d9b06",
   "metadata": {},
   "source": [
    "## Sifting"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "375ad557",
   "metadata": {},
   "source": [
    "Alice and Bob each have a string of bits, which will shortly become a key for cipher.\n",
    "At this point Alice and Bob can switch to communicating on a public channel. Their first step is to \n",
    "perform sifting - decide which bits to keep in their key.\n",
    "    \n",
    "Bob begins by telling Alice, which photons he measured. He then tells her which basis were used in each measurement.\n",
    "Then it's Alice's turn to send Bob her basis and to cancel out bits (representing states!) from her string \n",
    "that do not match both successful Bob's measurement and his usage of the same basis in each case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "007da9e6",
   "metadata": {},
   "source": [
    "bob_measurement_indicators = ''\n",
    "for bit in bob_bits:  # is it possible to optimise length of such an indicator?\n",
    "    if bit == '0' or bit == '1':\n",
    "        bob_measurement_indicators += '1'\n",
    "    else:\n",
    "        bob_measurement_indicators += '0'\n",
    "\n",
    "print(\"[Bob] My measurement indicator: {}\".format(bob_measurement_indicators))\n",
    "\n",
    "bob_indicated_basis = ''\n",
    "bob_indicated_bits = ''\n",
    "index = 0\n",
    "\n",
    "for indicator in bob_measurement_indicators:  # in optimised approach send only bases for successful measurements\n",
    "    if indicator == '1':\n",
    "        bob_indicated_basis += bob_basis[index]\n",
    "        bob_indicated_bits += bob_bits[index]\n",
    "    index += 1\n",
    "\n",
    "print(\"[Bob] I used bases: {}\".format(bob_indicated_basis))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "bc012195",
   "metadata": {},
   "source": [
    "Now it's Alice's turn to send Bob her bases and to cancel out bits (representing states!) from her string that do not\n",
    "match both successful Bob's measurement and his usage of the same basis in each case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "f45da563",
   "metadata": {},
   "source": [
    "alice_indicated_bits = ''\n",
    "alice_indicated_basis = ''\n",
    "alice_sifted_key = ''\n",
    "\n",
    "index = 0\n",
    "for indicator in bob_measurement_indicators:  # this loop ic practically copied - room for optimisation\n",
    "    if indicator == '1':\n",
    "        alice_indicated_bits += alice_bits[index]\n",
    "        alice_indicated_basis += alice_basis[index]\n",
    "    index += 1\n",
    "\n",
    "print(\"[Alice] I used bases: {}\".format(alice_indicated_basis))\n",
    "\n",
    "index = 0\n",
    "for basis in alice_indicated_basis:\n",
    "    if basis == bob_indicated_basis[index]:\n",
    "        alice_sifted_key += alice_indicated_bits[index]\n",
    "    index += 1"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "7e175f00",
   "metadata": {},
   "source": [
    "Now Bob gets info from Alice about her choices of bases, so that he can omit bits resulting from measurements\n",
    "when he used different basis than Alice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "75c3610c",
   "metadata": {},
   "source": [
    "bob_sifted_key = ''\n",
    "\n",
    "index = 0\n",
    "for basis in bob_indicated_basis:\n",
    "    if basis == alice_indicated_basis[index]:\n",
    "        bob_sifted_key += bob_indicated_bits[index]\n",
    "    index += 1\n",
    "\n",
    "print(\"Alice's sifted key is {},\\n& Bob's sifted key is {}.\".format(\n",
    "    alice_sifted_key, bob_sifted_key\n",
    "))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "1e6487a4",
   "metadata": {},
   "source": [
    "## Error estimation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "928cf748",
   "metadata": {},
   "source": [
    "Sifted keys generally differ from each other due to changes between states sent by Alice and received by Bob.\n",
    "In order to estimate empirical probability of error occurrence in the sifted keys we can publish parts\n",
    "of keys, compare them and calculate numbers of errors. Then published parts of key should be deleted, as\n",
    "they have just been exchanged via the public channel.\n",
    "\n",
    "Originally I've been using 'normal' error estimation; naive one. In this version of the demonstrator we can choose between naive and refined one (Lo, Chau, Ardehali, 2004). For this purpose I defined above two functions: naive_error & refined_average_error.\n",
    "Therefore, we must choose the method of error estimation (at the begining)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "14335bc1",
   "metadata": {},
   "source": [
    "alice_sifted_key_after_error_estimation = ''\n",
    "bob_sifted_key_after_error_estimation = ''\n",
    "qber = None\n",
    "pub_prob_rect = 0.5  # default probability of publication of bits measured in rectilinear basis for refined error est.\n",
    "pub_prob_diag = 0.5  # default probability of publication of bits measured in diagonal basis for refined error est.\n",
    "while True:\n",
    "    print('Please choose a method of error estimation (naive/refined):')\n",
    "    answer6 = str(input())\n",
    "    if answer6 not in ('naive', 'refined'):\n",
    "        print(\"Sorry, I didn't understand that.\")\n",
    "        continue\n",
    "    else:\n",
    "        break\n",
    "\n",
    "if answer6 == 'naive':\n",
    "    published_key_length_ratio = 0.2\n",
    "    print(\"Default ratio of length of published part of the key to it's full length is 0.2.\",\n",
    "          \"\\nDo you want to change it ('yes'/'no')?\")\n",
    "    while True:\n",
    "        answer7 = str(input())\n",
    "        if answer7 not in ('yes', 'no'):\n",
    "            print(\"Sorry, I didn't understand that.\")\n",
    "            continue\n",
    "        else:\n",
    "            break\n",
    "\n",
    "    if answer7 == 'yes':\n",
    "        print(\"Please set ratio of number of published bits of the key to it's full length\",\n",
    "              \"\\n as a numerical value between 0 and 1, e.g. 0.15:\")\n",
    "        while True:\n",
    "            try:\n",
    "                published_key_length_ratio = float(input())\n",
    "                if 0 <= published_key_length_ratio <= 1:\n",
    "                    break\n",
    "                else:\n",
    "                    print(\"Please enter a numerical value between 0 and 1\")\n",
    "                    continue\n",
    "            except TypeError:\n",
    "                print(\"Please enter a numerical value between 0 and 1\")\n",
    "                continue\n",
    "    naive_results = naive_error(\n",
    "        alice_key=alice_sifted_key,\n",
    "        bob_key=bob_sifted_key,\n",
    "        publication_prob_rect=published_key_length_ratio)\n",
    "    qber = naive_results.get('error estimation')\n",
    "    alice_sifted_key_after_error_estimation = naive_results.get('alice key')\n",
    "    bob_sifted_key_after_error_estimation = naive_results.get('bob key')\n",
    "else:  # refined error estimation\n",
    "    while True:\n",
    "        try:\n",
    "            print(\"Please set the probability of bits' publication in the rectilinear basis:\")\n",
    "            pub_prob_rect = float(input())\n",
    "        except TypeError:\n",
    "            print(\"Please enter a numerical value between 0 and 1, e.g. 0.57.\")\n",
    "            continue\n",
    "        else:\n",
    "            break\n",
    "    while True:\n",
    "        try:\n",
    "            print(\"Please set the probability of bits' publication in the diagonal basis:\")\n",
    "            pub_prob_diag = float(input())\n",
    "        except TypeError:\n",
    "            print(\"Please enter a numerical value between 0 and 1, e.g. 0.57.\")\n",
    "            continue\n",
    "        else:\n",
    "            break\n",
    "    \"\"\"Now that we have the inputs for the refined error estimation, we perform it with a function defined above.\"\"\"\n",
    "    refined_results = refined_average_error(\n",
    "        rect_prob=rectilinear_basis_prob,\n",
    "        rect_pub_prob=pub_prob_rect,\n",
    "        diag_pub_prob=pub_prob_diag,\n",
    "        alice_key=alice_sifted_key,\n",
    "        bob_key=bob_sifted_key,\n",
    "        alice_bases=alice_basis,\n",
    "        bob_bases=bob_basis\n",
    "    )\n",
    "    qber = refined_results.get('error estimator')\n",
    "    alice_sifted_key_after_error_estimation = refined_results.get('alice key')\n",
    "    bob_sifted_key_after_error_estimation = refined_results.get('bob key')"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "ffd0475d",
   "metadata": {},
   "source": [
    "## CASCADE"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8015247",
   "metadata": {},
   "source": [
    "Assuming the estimate of empirical probability of error is reasonably small we can continue\n",
    "with the error correction. Naturally for BB84 we assume it's Bob's key that's flawed.\n",
    "\n",
    "We begin by checking the parity of Alice's and Bob's sifted keys, \n",
    "shortened by the subsets used for error estimation.\n",
    "\n",
    "CASCADE: Firstly, we need to assign bits to their indexes in original strings. Therefore, we create dictionaries\n",
    "for Alice and for Bob."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "05374991",
   "metadata": {},
   "source": [
    "n = len(alice_sifted_key_after_error_estimation)\n",
    "alice_cascade = {}\n",
    "bob_cascade = {}\n",
    "\n",
    "for i in range(n):  # I dynamically create dictionaries with indexes as keys and bits as values\n",
    "    alice_cascade[str(i)] = alice_sifted_key_after_error_estimation[i]\n",
    "    bob_cascade[str(i)] = bob_sifted_key_after_error_estimation[i]"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "29a51240",
   "metadata": {},
   "source": [
    "Secondly, we need to set up CASCADE itself: sizes of blocks in each pass, numeration of passes and a dictionary\n",
    "for corrected bits with their indexes from original Bob's string as keys and correct bits as values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "5f641ee2",
   "metadata": {},
   "source": [
    "blocks_sizes = cascade_blocks_sizes(quantum_bit_error_rate=qber, key_length=n, n_passes=cascade_passes)\n",
    "bob_corrected_bits = {}"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "861088ee",
   "metadata": {},
   "source": [
    "Thirdly, in order to return to blocks from earlier passes of CASCADE we need a history of blocks with indexes and bits,\n",
    "so implemented by dictionaries as list elements per pass, nested in general history list. Moreover, we create lists\n",
    "and variables to remember error rates of bits after each consecutive pass of CASCADE and to count the bits exchanged."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "63bede9e",
   "metadata": {},
   "source": [
    "history = []\n",
    "error_rates = []\n",
    "pass_number = 0\n",
    "exchanged_bits_counter = 0\n",
    "\n",
    "for size in blocks_sizes:\n",
    "    print(\"--------------------\\nThis is CASCADE pass number {}\".format(pass_number + 1))\n",
    "    try:\n",
    "        pass_number_of_blocks = int(\n",
    "            -1 * np.floor(-1 * n // size))  # I calculate how many blocks are in total in this pass\n",
    "    except ZeroDivisionError:\n",
    "        print(\"Initial block size equal to 0, please check why.\")\n",
    "        quit()\n",
    "\n",
    "    print(\"With {} bits in a single block and {} bits in total we have {} blocks in this CASCADE pass\".format(\n",
    "        size, n, pass_number_of_blocks\n",
    "    ))\n",
    "\n",
    "    alice_pass_parity_list = []\n",
    "    bob_pass_parity_list = []\n",
    "    alice_blocks = []\n",
    "    bob_blocks = []\n",
    "\n",
    "    for block_index in cascade_blocks_generator(string_length=n, blocks_size=size):\n",
    "        \"\"\"We sample bits from the raw key (after the error estimation phase) into blocks, with the generator\n",
    "        defined above.\n",
    "        \"\"\"\n",
    "        alice_block = {}  # a dictionary for a single block for Alice\n",
    "        bob_block = {}  # a dictionary for a single block for Bob\n",
    "\n",
    "        for index in block_index:  # I add proper bits to these dictionaries\n",
    "            alice_block[str(index)] = alice_cascade[str(index)]\n",
    "            bob_block[str(index)] = bob_cascade[str(index)]\n",
    "\n",
    "        \"\"\"I append single blocks created for given indexes to lists of block for this particular CASCADE's pass\"\"\"\n",
    "        alice_blocks.append(alice_block)\n",
    "        bob_blocks.append(bob_block)\n",
    "\n",
    "    for i in range(pass_number_of_blocks):\n",
    "        print(\">>>\\nThis is block number {} out of {} in CASCADE pass number {}.\\n\".format(\n",
    "            i+1, pass_number_of_blocks, pass_number))\n",
    "\n",
    "        current_indexes = list(alice_blocks[i].keys())  # same as Bob's\n",
    "        print(\"Indexes of bits from the raw key which are sampled for the current block in CASCADE are: {}\".format(\n",
    "            current_indexes\n",
    "        ))\n",
    "\n",
    "        alice_current_bits = list(alice_blocks[i].values())\n",
    "        print(\"Alice's bits with these indexes are: {}\".format(alice_current_bits))\n",
    "\n",
    "        bob_current_bits = list(bob_blocks[i].values())\n",
    "        print(\"Bob's bits with these indexes are: {}\".format(bob_current_bits))\n",
    "\n",
    "        alice_bit_values = []\n",
    "        bob_bit_values = []\n",
    "\n",
    "        for j in range(len(current_indexes)):\n",
    "            alice_bit_values.append(int(alice_current_bits[j]))\n",
    "            bob_bit_values.append(int(bob_current_bits[j]))\n",
    "\n",
    "        alice_pass_parity_list.append(sum(alice_bit_values) % 2)\n",
    "        print(\"[Alice] Bits in my block have parity {}\".format(alice_pass_parity_list[-1]))\n",
    "\n",
    "        bob_pass_parity_list.append(sum(bob_bit_values) % 2)\n",
    "        print(\"[Bob] Bits in my block have parity {}\".format(bob_pass_parity_list[-1]))\n",
    "\n",
    "        if alice_pass_parity_list[i] != bob_pass_parity_list[i]:  # we check if we should perform BINARY\n",
    "            print(\"[Alice] We have different parities. Let's perform binary search for the erroneous bit.\")\n",
    "\n",
    "            binary_results = binary(\n",
    "                sender_block=alice_blocks[i],\n",
    "                receiver_block=bob_blocks[i],\n",
    "                indexes=current_indexes\n",
    "            )\n",
    "\n",
    "            \"\"\"Firstly we add the number of exchanged bits during this BINARY performance to the general number\n",
    "            of bits exchanged via the public channel.\n",
    "            \"\"\"\n",
    "            exchanged_bits_counter += binary_results[2]\n",
    "\n",
    "            \"\"\"Secondly we change main dictionary with final results and current blocks for history\"\"\"\n",
    "            bob_cascade[binary_results[1]] = binary_results[0]\n",
    "            bob_blocks[i][binary_results[1]] = binary_results[0]\n",
    "\n",
    "            \"\"\"Thirdly we change the error bit in blocks' history\n",
    "            We need to perform BINARY on all blocks which we correct in history list\n",
    "            history[number of pass][owner][number of block]\n",
    "            \"\"\"\n",
    "            if pass_number > 0:  # in the first pass of CASCADE there are no previous blocks\n",
    "                print(\"[Alice] Since this CASCADE pass number {} we should find blocks \"\n",
    "                      \"with the corrected bit from previous passes\".format(pass_number))\n",
    "                for n_pass in range(pass_number):  # we check all previous passes\n",
    "                    no_blocks = len(history[n_pass - 1][1])\n",
    "                    for n_block in range(no_blocks):  # we check all Bob's blocks in each previous pass\n",
    "                        correct_bit_index = binary_results[1]\n",
    "                        if correct_bit_index in history[n_pass - 1][1][n_block]:\n",
    "                            \"\"\"After we found the correct bit's index in one of Bob's blocks from previous pass of the\n",
    "                            CASCADE, we can correct it and once again perform BINARY on this block (from Alice & Bob).\n",
    "                            That's because there has to be another erroneous bit, that together with the 'old' error\n",
    "                            we have just corrected made the parity the same between the blocks.\n",
    "                            \"\"\"\n",
    "                            history[n_pass - 1][1][n_block][correct_bit_index] = binary_results[0]\n",
    "\n",
    "                            print(\"[Bob] I found a previous block. Let's go!\")\n",
    "                            print(\"[Alice] Mee too.\")\n",
    "\n",
    "                            try:\n",
    "                                if type(history[n_pass - 1][1][n_block]) == str:\n",
    "                                    indexes = ast.literal_eval(history[n_pass - 1][1][n_block])\n",
    "                                    binary_previous = binary(\n",
    "                                        sender_block=history[n_pass - 1][0][n_block],\n",
    "                                        receiver_block=history[n_pass - 1][1][n_block],\n",
    "                                        indexes=indexes.keys()\n",
    "                                    )\n",
    "                                elif type(history[n_pass - 1][1][n_block]) == dict:\n",
    "                                    binary_previous = binary(\n",
    "                                        sender_block=history[n_pass - 1][0][n_block],\n",
    "                                        receiver_block=history[n_pass - 1][1][n_block],\n",
    "                                        indexes=history[n_pass - 1][1][n_block].keys()\n",
    "                                    )\n",
    "                            except AttributeError:\n",
    "                                print(\"AttributeError for binary_previous\")\n",
    "\n",
    "                                file = open(\"error.txt\", \"w\")\n",
    "                                file.write('\\n' + 'type of history: ' + str(type(history)) + '\\n' + 'type of '\n",
    "                                                                                                    'history['\n",
    "                                                                                                    'n_pass]: ' +\n",
    "                                           str(type(history[n_pass])) + '\\n' + 'type of history[n_pass][1]: ' +\n",
    "                                           str(type(history[n_pass][1])) + '\\n' + 'type of history[n_pass][1]['\n",
    "                                                                                  'n_block]: ' +\n",
    "                                           str(type(history[n_pass][1][n_block])) + '\\n' + str(history) + '\\n')\n",
    "                                file.close()\n",
    "                                exit()\n",
    "\n",
    "                            exchanged_bits_counter += binary_previous[2]\n",
    "                            bob_cascade[binary_previous[1]] = binary_previous[0]\n",
    "                            bob_blocks[i][binary_previous[1]] = binary_previous[0]\n",
    "\n",
    "    history.append([alice_blocks, bob_blocks])\n",
    "    pass_number += 1\n",
    "\n",
    "    \"\"\"For the purposes of demonstration we check the error rate after each pass:\"\"\"\n",
    "    alice_key_error_check = ''.join(list(alice_cascade.values()))\n",
    "    bob_key_error_check = ''.join(list(bob_cascade.values()))\n",
    "\n",
    "    key_error_rate = 0\n",
    "    index = 0\n",
    "    for bit in alice_key_error_check:\n",
    "        if bit != bob_key_error_check[index]:\n",
    "            key_error_rate += 1\n",
    "        index += 1\n",
    "    try:\n",
    "        key_error_rate = key_error_rate / len(alice_key_error_check)\n",
    "        error_rates.append(key_error_rate)  # its length is equivalent to no. CASCADE passes performed\n",
    "    except ZeroDivisionError:\n",
    "        print(\"ZeroDivisionError in key error rate calculation.\")\n",
    "        exit()\n",
    "    print(\"After this CASCADE pass the ACTUAL error rate is equal to {}.\".format(key_error_rate))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "9de30976",
   "metadata": {},
   "source": [
    "Time to create strings from cascade dictionaries into corrected keys:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "617036b6",
   "metadata": {},
   "source": [
    "alice_correct_key = ''.join(list(alice_cascade.values()))\n",
    "bob_correct_key = ''.join(list(bob_cascade.values()))\n",
    "\n",
    "print(\"Alice's correct key:\", \"\\n{}\".format(alice_correct_key))\n",
    "print(\"Bob's key after performing CASCADE error correction:\", \"\\n{}\".format(bob_correct_key))\n",
    "print(\"Number of bits exchanged during error correction: {}\".format(exchanged_bits_counter))\n",
    "\n",
    "print(\"History:\", \"\\n{}\".format(history))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "4f07814e",
   "metadata": {},
   "source": [
    "All that remains is to randomly choose number of bits for deletion, equal to number of exchanged bits\n",
    "during error correction phase. It's a form of a rudimentary privacy amplification. Let's say Alice randomly deletes \n",
    "bits and informs Bob which indexes they were on, so that the computational cost would be equal \n",
    "to the number of deleted bits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "651ea21b",
   "metadata": {},
   "source": [
    "deleted_bits_counter = 0\n",
    "try:\n",
    "    deletion_prob = exchanged_bits_counter / len(alice_correct_key)\n",
    "except ZeroDivisionError:\n",
    "    print('Error in bit deletion.')\n",
    "    deletion_prob = 0  # no idea how to set it better in such a case\n",
    "\n",
    "index = 0\n",
    "while deleted_bits_counter < exchanged_bits_counter:\n",
    "    if index == len(alice_correct_key):  # just in case we won't delete enough bits in the first 'run'\n",
    "        index = 0\n",
    "    if random.uniform(0, 1) <= deletion_prob:  # we \"increase\" the prob. by < OR =\n",
    "        alice_correct_key = alice_correct_key[0: index:] + alice_correct_key[index + 1::]\n",
    "        bob_correct_key = bob_correct_key[0: index] + bob_correct_key[index + 1::]\n",
    "        deleted_bits_counter += 1\n",
    "\n",
    "    index += 1"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "daf26c88",
   "metadata": {},
   "source": [
    "Now we finally have the proper keys."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "f660773b",
   "metadata": {},
   "source": [
    "print(\"Alice's correct key:\", \"\\n{}\".format(alice_correct_key))\n",
    "print(\"Bob's key after performing CASCADE error correction:\", \"\\n{}\".format(bob_correct_key))\n",
    "print(\"Number of bits exchanged during error correction: {}\".format(exchanged_bits_counter))\n",
    "\n",
    "print(\"History:\", \"\\n{}\".format(history))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "1d39856c",
   "metadata": {},
   "source": [
    "## Privacy Amplification"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3097cf0",
   "metadata": {},
   "source": [
    "Finally we perform privacy amplification, using hashing functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "746f804d",
   "metadata": {},
   "source": [
    "alice_digest = sha1(alice_correct_key)\n",
    "print(\"Alice's correct key's digest is: {}\".format(alice_digest))\n",
    "\n",
    "bob_digest = sha1(bob_correct_key)\n",
    "print(\"Bob's correct key's digest is: {}\".format(bob_digest))"
   ],
   "outputs": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}