import matplotlib.pyplot as pyplot
import numpy as np
from scipy.stats import pearsonr
import random
from operator import xor

# Fast implementation of the Hamming weight for 64 bit values
# See book: A Hacker's delight
def popcount(x):
    x -= (x >> 1) & 0x5555555555555555
    x = (x & 0x3333333333333333) + ((x >> 2) & 0x3333333333333333)
    x = (x + (x >> 4)) & 0x0f0f0f0f0f0f0f0f
    return ((x * 0x0101010101010101) & 0xffffffffffffffff ) >> 56

# not sure if the hamming weight model w0rks for Speck

import math

NUM_ROUNDS = 22
BLOCK_SIZE = 32
KEY_SIZE = 64
WORD_SIZE = 16


# SHIFTs for SPECK
ALPHA = 7
BETA = 2

mod_mask = (2 ** WORD_SIZE) -1
mod_mask_sub = (2 ** WORD_SIZE)

def bytesToWords16(b):       
    return [(b >> (x * WORD_SIZE)) & mod_mask for x in
                      range(0, math.ceil(KEY_SIZE // WORD_SIZE))]
        
        

'''
 The 16bit Speck roundfunction
'''
def ER16(x, y, k):

    rs_x = ((x << (16 - ALPHA)) + (x >> ALPHA)) & mod_mask

    add_sxy = (rs_x + y) & mod_mask

    new_x = k ^ add_sxy

    ls_y = ((y >> (16 - BETA)) + (y << BETA)) & mod_mask

    new_y = new_x ^ ls_y

    return new_x, new_y

key_schedule = bytesToWords16(key)

'''

The 16 bit python key schedule

void Speck128256KeySchedule(u64 K[],u64 rk[])
{
    u64 i,D=K[3],C=K[2],B=K[1],A=K[0];
    for(i=0;i<33;){
        rk[i]=A; ER64(B,A,i++);
        rk[i]=A; ER64(C,A,i++);
        rk[i]=A; ER64(D,A,i++);
    }
    rk[i]=A;
}
'''
def key_schedule(k):

    D=k[3]
    C=k[2]
    B=k[1]
    A=k[0]
    out = []
    i = 0
    while i < 21:
        out.append(A)
        B, A = ER16(B, A, i)
        i += 1
        out.append(A)
        C, A = ER16(C, A, i)
        i+= 1
        out.append(A)
        D, A = ER16(D, A, i)
        i+= 1
    out.append(A)
    return out
        

assert key_schedule(l_schedule) == [0x708,0xf32, 0x2bf1,0x8035,0xa48e,0x8482, 0x74ee, 0xf589, 0xb396, 0xb231, 0xdab2, 0x57bc, 0x704e,0x9947,0xe2d2, 0xea6a, 0x4ebe, 0xdd24, 0x6b64, 0x3ab1, 0x1c57, 0x7bde]