import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from .utils import AES_key_schedule, build_xor_weights_and_biases, get_linear_mapping, \
    build_256_corners, MUL2, MUL3, SBOX, SBOX_INV, MUL14, MUL9, MUL13, MUL11, integer_to_bitvector, state_as_ints, \
    bytes_matrix_to_binary_states, build_xor4_weights_and_biases

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
DTYPE = torch.float16 # torch.float16 for benchmarking

def build_parameters_from_matrices(w, b, ws=None):
    wt = torch.tensor(w, dtype=DTYPE, requires_grad=False).to(DEVICE)
    bt = torch.tensor(b.flatten(), dtype=DTYPE, requires_grad=False).to(DEVICE)
    if ws is None:
        return wt, bt, None
    wst = torch.tensor(ws, dtype=DTYPE, requires_grad=False).to(DEVICE)
    return wt, bt, wst

class ClippingLayer(nn.Module):
    def __init__(self, epsilon=1/4):
        '''
        Implements the step function with parameter epsilon as defined in the paper.
        '''
        assert epsilon == 1/4, "Untested for other values; some values of epsilon, such as 1/3, cause numerical " \
                               "precision issues and result in an incorrect implementation of the cryptographic " \
                               "primitive. If changing epsilon, ensure that ClippingLayer(epsilon).forward(x) == x, " \
                               "for x a tensor containing 0s and 1s only."
        super().__init__()
        self.register_buffer("epsilon", torch.tensor(epsilon))

    def forward(self, x):
        return torch.div(
            torch.sub(
                F.relu(torch.sub(x, self.epsilon)),
                F.relu(torch.sub(x, 1 - self.epsilon))
            )
            , (1 - 2 * self.epsilon))

class SumBumpsLayer(nn.Module):
    def __init__(self, epsilon=1/4):
        '''
        Implements the sumbumps function with parameter epsilon as defined in the paper.
        '''
        assert epsilon == 1/4, "Untested for other values; some values of epsilon, such as 1/3, cause numerical " \
                               "precision issues and result in an incorrect implementation of the cryptographic " \
                               "primitive. If changing epsilon, ensure that SumBumpsLayer(epsilon).forward(x) == 0, " \
                               "for x a tensor containing 0s and 1s only." 
        super().__init__()
        self.register_buffer("epsilon", torch.tensor(epsilon, dtype=DTYPE))

    def forward(self, x):
        return torch.mul(
            torch.add(
                F.relu(torch.sub(x, 1)),
                torch.sub(
                    torch.sub(F.relu(x),
                              F.relu(torch.sub(x, self.epsilon))
                              ),
                    F.relu(torch.sub(x, 1 - self.epsilon))
                )
            )
            , 1 / self.epsilon).sum(-1)

class NeuralAES(nn.Module):
    def __init__(self,
                 secret_key,
                 direction='Encryption',
                 c_parameter=1.0,
                 number_of_rounds=10,
                 protected=False,
                 epsilon=None):


        super().__init__()

        self.number_of_rounds = number_of_rounds
        self.c_parameter = c_parameter

        # Init ARK Parameters
        self.round_keys = AES_key_schedule(secret_key, self.number_of_rounds)  # TODO: add number of rounds
        self.round_key_tensors = torch.tensor(bytes_matrix_to_binary_states(self.round_keys), dtype=DTYPE)
        self.register_buffer("buff_round_keys", self.round_key_tensors[:, None, :, :])

        self.xorw, self.xorb, _ = build_parameters_from_matrices(*build_xor_weights_and_biases(c_parameter))
        self.xor4w, self.xor4b, _ = build_parameters_from_matrices(*build_xor4_weights_and_biases(c_parameter))
        if protected:
            self.sum_bumps = SumBumpsLayer(epsilon)
            self.step = ClippingLayer(epsilon)

        self.byte_op_w, self.byte_op_b, _ = build_parameters_from_matrices(*build_256_corners(self.c_parameter))

        if direction == 'Encryption':
            self.fw_function = self.encrypt_protected if protected else self.encrypt

            self.sbw_s = torch.tensor(get_linear_mapping(SBOX), dtype=DTYPE, requires_grad=False).to(DEVICE)
            self.m2w_s = torch.tensor(get_linear_mapping(MUL2), dtype=DTYPE, requires_grad=False).to(DEVICE)
            self.m3w_s = torch.tensor(get_linear_mapping(MUL3), dtype=DTYPE, requires_grad=False).to(DEVICE)

        elif direction == 'Decryption':
            self.fw_function = self.decrypt_protected if protected else self.decrypt

            self.sbiw_s = torch.tensor(get_linear_mapping(SBOX_INV), dtype=DTYPE, requires_grad=False).to(DEVICE)
            self.m14w_s = torch.tensor(get_linear_mapping(MUL14), dtype=DTYPE, requires_grad=False).to(DEVICE)
            self.m9w_s = torch.tensor(get_linear_mapping(MUL9), dtype=DTYPE, requires_grad=False).to(DEVICE)
            self.m13w_s = torch.tensor(get_linear_mapping(MUL13), dtype=DTYPE, requires_grad=False).to(DEVICE)
            self.m11w_s = torch.tensor(get_linear_mapping(MUL11), dtype=DTYPE, requires_grad=False).to(DEVICE)

    def apply_op(self, x, w0, b0, w1=None):
        h0 = F.relu(F.linear(x, w0, b0))
        if w1 is None:
            return h0.sum(dim=-1)
        else:
            return F.linear(h0, w1)

    def SR(self, state):
        return torch.cat((
            state[:, 0:1],
            state[:, 1:2, [1, 2, 3, 0]],
            state[:, 2:3, [2, 3, 0, 1]],
            state[:, 3:,  [3, 0, 1, 2]])
            , dim=1
        )

    def SR_inv(self, state):
        return torch.cat((
            state[:, 0:1],
            state[:, 1:2, [3, 0, 1, 2]],
            state[:, 2:3, [2, 3, 0, 1]],
            state[:, 3:,  [1, 2, 3, 0]])
            , dim=1
        )

    def SB(self, state):
        return self.apply_op(state, self.byte_op_w, self.byte_op_b, self.sbw_s)

    def SB_inv(self, state):
        state = state.contiguous()
        return self.apply_op(state, self.byte_op_w, self.byte_op_b, self.sbiw_s)

    def MC(self, state):
        all_corners = F.relu(F.linear(state, self.byte_op_w, self.byte_op_b))
        M2 = F.linear(all_corners, self.m2w_s)
        M3 = F.linear(all_corners, self.m3w_s)
        X = torch.stack((
            M2,
            torch.cat((M3[:, 1:], M3[:, :1]), dim=1),
            torch.cat((state[:, 2:], state[:, :2]), dim=1),
            torch.cat((state[:, 3:], state[:, :3]), dim=1)
            ), dim=-1)

        return self.apply_op(X, self.xor4w, self.xor4b)

    def MC_inv(self, state):
        all_corners = F.relu(F.linear(state, self.byte_op_w, self.byte_op_b))
        M9 = F.linear(all_corners, self.m9w_s)
        M11 = F.linear(all_corners, self.m11w_s)
        M13 = F.linear(all_corners, self.m13w_s)
        M14 = F.linear(all_corners, self.m14w_s)
        X = torch.stack((
            M14,
            torch.cat((M11[:, 1:], M11[:, :1]), dim=1),
            torch.cat((M13[:, 2:], M13[:, :2]), dim=1),
            torch.cat((M9[:, 3:], M9[:, :3]), dim=1)
            ), dim=-1)

        return self.apply_op(X, self.xor4w, self.xor4b)

    def AES_round(self, state, final=False):
        X = self.SR(state)
        X = self.SB(X)
        if not final:
            X = self.MC(X)
        return X

    def AES_round_inverse(self, X, final=False):
        if not final:
            X = self.MC_inv(X)
        X = self.SB_inv(X)
        X = self.SR_inv(X)
        return X

    def ARK(self, state, round_key):
        # Possible optimisation: precompute self.xorw[0, 1] * round_key + self.xorb[0] and self.xorw[1, 1] * round_key + self.xorb[1]
        corner_0 = self.xorw[0, 0] * state + self.xorw[0, 1] * round_key + self.xorb[0]
        corner_1 = self.xorw[1, 0] * state + self.xorw[1, 1] * round_key + self.xorb[1]
        return F.relu(corner_0) + F.relu(corner_1)

    def encrypt(self, state_tensor):
        '''
        Applies NN_AES to a binary tensor of plaintexts with dimensions batch_size x 128.
        '''
        state_tensor = state_tensor.view(-1, 4, 4, 8).transpose(1, 2)
        state_tensor = self.ARK(state_tensor, self.buff_round_keys[0])
        for r in range(1, self.number_of_rounds):
            state_tensor = self.AES_round(state_tensor)
            state_tensor = self.ARK(state_tensor, self.buff_round_keys[r])
        state_tensor = self.AES_round(state_tensor, final=True)
        state_tensor = self.ARK(state_tensor, self.buff_round_keys[self.number_of_rounds])
        return state_tensor.transpose(1, 2).reshape(-1, 128)

    def decrypt(self, state_tensor):
        '''
        Applies NN_AES to a binary tensor of ciphertexts with dimensions batch_size x 128.
        '''
        state_tensor = state_tensor.view(-1, 4, 4, 8).transpose(1, 2)
        state_tensor = self.ARK(state_tensor, self.buff_round_keys[self.number_of_rounds])
        state_tensor = self.AES_round_inverse(state_tensor, final=True)
        for r in range(1, self.number_of_rounds):
            state_tensor = self.ARK(state_tensor, self.buff_round_keys[self.number_of_rounds - r])
            state_tensor = self.AES_round_inverse(state_tensor)
        state_tensor = self.ARK(state_tensor, self.buff_round_keys[0])
        return state_tensor.transpose(1, 2).reshape(-1, 128)

    def encrypt_protected(self, state_tensor):
        '''
        Applies NN_AES to a binary tensor of plaintexts with dimensions batch_size x 128.
        '''
        sum_bumps = self.sum_bumps(state_tensor)
        state_tensor = self.step(state_tensor)
        state_tensor = self.encrypt(state_tensor)
        state_tensor = F.relu(torch.sub(state_tensor, sum_bumps[:, None]))
        return state_tensor

    def decrypt_protected(self, state_tensor):
        '''
        Applies NN_AES to a binary tensor of plaintexts with dimensions batch_size x 128.
        '''
        sum_bumps = self.sum_bumps(state_tensor)
        state_tensor = self.step(state_tensor)
        state_tensor = self.decrypt(state_tensor)
        state_tensor = F.relu(torch.sub(state_tensor, sum_bumps[:, None]))
        return state_tensor

    def forward(self, state_tensor):
        '''
        Applies NN_AES to a binary tensor of plaintexts with dimensions batch_size x 128.
        '''
        return self.fw_function(state_tensor)
