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

from nnaes.base import NeuralAES, SumBumpsLayer, DEVICE, DTYPE
from nnaes.ttables import TTablesNeuralAES
from nnaes.utils import integer_to_bitvector, state_as_ints

def encrypt_list_of_plaintexts(plaintexts, model, dtype=DTYPE):
    state = [integer_to_bitvector(x) for x in plaintexts]
    model = model.to(DEVICE)
    model.eval()
    state_tensor = torch.tensor(state, dtype=dtype).to(DEVICE)
    with torch.inference_mode():
        state_tensor = model(state_tensor)
    return state_as_ints(state_tensor)


def test_aes_testvectors(protected=False, epsilon=None):
    p_fips = [0x3243f6a8885a308d313198a2e0370734]
    k_fips = 0x2b7e151628aed2a6abf7158809cf4f3c
    c_fips = [0x3925841d02dc09fbdc118597196a0b32]

    k_aesavs = 0
    p_aesavs = [
        0x80000000000000000000000000000000,
        0xc0000000000000000000000000000000,
        0xe0000000000000000000000000000000,
        0xf0000000000000000000000000000000,
        0xf8000000000000000000000000000000,
    ]
    c_aesavs = [
        0x3ad78e726c1ec02b7ebfe92b23d9ec34,
        0xaae5939c8efdf2f04e60b9fe7117b2c2,
        0xf031d4d74f5dcbf39daaf8ca3af6e527,
        0x96d9fd5cc4f07441727df0f33e401a36,
        0x30ccdb044646d7e1f3ccea3dca08b8c0,
    ]

    for model_type in [NeuralAES, TTablesNeuralAES]:
        for c_parameter in [1.0, 0.5]:
            aes_enc = model_type(k_fips, c_parameter=c_parameter, protected=protected, epsilon=epsilon)
            ct = encrypt_list_of_plaintexts(p_fips, aes_enc)
            assert ct[0] == c_fips[0]

            aes_enc = model_type(k_aesavs, c_parameter=c_parameter, protected=protected, epsilon=epsilon)
            ct = encrypt_list_of_plaintexts(p_aesavs, aes_enc)
            for i in range(5):
                assert ct[i] == c_aesavs[i]

        for c_parameter in [1.0, 0.5]:
            aes_dec = model_type(
                k_fips,
                direction="Decryption",
                c_parameter=c_parameter,
                protected=protected,
                epsilon=epsilon,
            )
            pt = encrypt_list_of_plaintexts(c_fips, aes_dec)
            assert pt[0] == p_fips[0]

            aes_dec = model_type(
                k_aesavs,
                direction="Decryption",
                c_parameter=c_parameter,
                protected=protected,
                epsilon=epsilon,
            )
            pt = encrypt_list_of_plaintexts(c_aesavs, aes_dec)
            for i in range(5):
                assert pt[i] == p_aesavs[i]


def test_unprotected_aes_testvectors():
    test_aes_testvectors(protected=False)


def test_protected_aes_testvectors():
    test_aes_testvectors(protected=True, epsilon=1 / 4)


def test_protected_aes_functionality(num_samples=10**3, epsilon=1 / 4):
    for model_type in [NeuralAES, TTablesNeuralAES]:
        torch.cuda.empty_cache()
        k = 0x2b7e151628aed2a6abf7158809cf4f3c
        model = model_type(k, c_parameter=0.5, protected=True, epsilon=epsilon).to(DEVICE)
        model.eval()
        p = torch.randint(0, 2, (num_samples, 128), dtype=DTYPE).to(DEVICE)

        with torch.inference_mode():
            c = model(p).detach().cpu().numpy()
        assert np.all(np.logical_or(c == 0, c == 1))

        p1 = p.clone().to(DEVICE)
        p1[:, 3] = torch.abs(p[:, 3] - torch.div(torch.rand(num_samples, dtype=DTYPE), 2).to(DEVICE))
        expected_diff = SumBumpsLayer(epsilon=epsilon).to(DEVICE)(p1)

        with torch.inference_mode():
            c = model(p)
            c0 = F.relu(torch.sub(c, expected_diff[:, None])).detach().cpu().numpy()
            c1 = model(p1).detach().cpu().numpy()
        assert np.all(c1 == c0)