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

from .base import NeuralAES

from .utils import get_linear_mapping, MUL2, MUL3, SBOX

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

class TTablesNeuralAES(NeuralAES):
    def __init__(self, secret_key, direction='Encryption', c_parameter=1.0, number_of_rounds=10, protected=False,
                 epsilon=None):
        '''
        Builds a number_of_rounds instance of the AES embedded with the key secret_key, given as an integer. The corresponding NN uses c_parameter in the corner functions.
        If protected is set to True, the protection measure described in the paper is applied, with the specified epsilon value.
        '''
        super().__init__(secret_key, direction, c_parameter, number_of_rounds, protected, epsilon)

        SM2 = np.uint8(MUL2)[np.uint8(SBOX)]
        SM3 = np.uint8(MUL3)[np.uint8(SBOX)]

        if direction == 'Encryption':
            self.m2w_s = torch.tensor(get_linear_mapping(SM2), dtype=DTYPE, requires_grad=False).to(DEVICE)
            self.m3w_s = torch.tensor(get_linear_mapping(SM3), dtype=DTYPE, requires_grad=False).to(DEVICE)
        # TTables is only applicable in the encryption direction, so the performance gain does not port to decryption
        
    def SB_MC_fused(self, state):
        '''
        Fused SubBytes + MixColumns for non-final encryption rounds.
        Expects ShiftRows to have already been applied.
        '''
        all_corners = F.relu(F.linear(state, self.byte_op_w, self.byte_op_b))
        SB = F.linear(all_corners, self.sbw_s)
        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((SB[:, 2:], SB[:, :2]), dim=1),
                         torch.cat((SB[:, 3:], SB[:, :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)
        if not final:
            X = self.SB_MC_fused(X)
        else:
            X = self.SB(X)
        return X