##kode til hellman tables aflevering

from Crypto.Cipher import AES
import  matplotlib.pyplot as plt
from collections import defaultdict
import random
import numpy as np
import math
import time
plt.style.use("seaborn-darkgrid")
start_time = time.time()

def AES128Encrypt(key,plaintext):
    if len(key) < 16:
        missing = 16-len(key)
        for i in range(missing):
            key+= b'0'
    #print("PADDED KEY ",len(key))
    cipher = AES.new(key,AES.MODE_ECB)
    ciphertext = cipher.encrypt(plaintext)
    return ciphertext[0:3] #kun 24 bits, i.e. 3 bytes

NrOfTables = 256 #2^8
plaintext_orig = bytearray(b'0123456789abcdef')

m = [2**11,2**10,2**9,2**8,2**7,2**6,2**5,2**4,2**3,2**2] #m værdier vi vil køre for
mstore = []
domainval = pow(2,24)-1

def PH_hellman(num_tables, l, m, t):
    return 1 - math.exp(-math.sqrt((2 * m * num_tables ** 2) / (2 ** l)) * (math.exp(math.sqrt((2 * m * t ** 2) / (2 ** l))) - 1) / (math.exp(math.sqrt((2 * m * t ** 2) / (2 ** l))) + 1))

#source:
#https://stackoverflow.com/questions/21017698/converting-int-to-bytes-in-python-3
def int_to_bytes(x: int) -> bytes:
    #x = x.item()
    return x.to_bytes((x.bit_length() + 7) // 8, 'big', signed = False)
#source:
#https://stackoverflow.com/questions/21017698/converting-int-to-bytes-in-python-3
def int_from_bytes(xbytes: bytes) -> int:
    return int.from_bytes(xbytes, 'big',signed = False)

def fi(ciphertext,index):
    ciphertext = AES128Encrypt(ciphertext, plaintext_orig)
    temp1 = int_from_bytes(ciphertext)
    f_k_i = (temp1 + index) % pow(2,24)
    #print("RESULT OF FKI ",f_k_i)
    return f_k_i # returns integer

columns = 1000

for mvalue in m:
    currentValues = np.empty((NrOfTables, mvalue), dtype=int)
    storedvalues = set()
    columncoverages = []
    #print("RUNNING FOR VALUE M: ",mvalue)
    for t in range(columns):
        for tableNr in range(0,NrOfTables):
            for rows in range(0,mvalue):
                if t == 0:
                    temprows = int_to_bytes(rows)
                    f_k_i_result = fi(temprows, tableNr)
                    currentValues[tableNr][rows] = f_k_i_result
                else:
                    cipher_temp = int_to_bytes(int(currentValues[tableNr][rows]))
                    f_k_i_result = fi(cipher_temp, tableNr)
                    currentValues[tableNr][rows] = f_k_i_result

                storedvalues.add(f_k_i_result)
        columncoverages.append((len(storedvalues) / domainval)*100)

    mstore.append(columncoverages)

print("--- %s seconds ---" % (time.time() - start_time))


fig1, ax1 = plt.subplots(1)
for idx, columncoveragesin in enumerate(mstore):
    ax1.plot(range(t+1),columncoveragesin, label = 'm = '+str(m[idx]))

plt.legend(loc='upper left')
plt.xlabel('Columns')
plt.ylabel('Coverage [%]')
plt.show()