import pytest

import numpy as np

from pycrcnn.functional import average_pool as avg
from pycrcnn.functional import square_layer as sq
from pycrcnn.functional.average_pool import AveragePoolLayer
from pycrcnn.functional.rencryption_layer import RencryptionLayer
from pycrcnn.functional.flatten_layer import FlattenLayer
from pycrcnn.he.HE import MockHE


class TestsAverageLayer:
    @pytest.fixture
    def HE(self):
        HE = MockHE()
        return HE

    @pytest.fixture
    def image(self):
        return np.array([
            [[[1, 1, -1, 0, 0],
              [0, 1, 1, 1, 0],
              [0, 0, 1, 1, 1],
              [0, -2, 1, 1, 0],
              [0, 1, 1, 0, 0]],
             [[1, 0, 1, 1, 0],
              [2, 1, 2, 0, 0],
              [0, 0, -1, 0, 1],
              [0, 1, 1, 1, 0],
              [0, 1, 0, 0, 0]]],
            [[[1, 1, 1, 0, 0],
              [0, 1, 2, 1, 1],
              [0, 0, 1, 1, 1],
              [1, 0, 2, 3, 0],
              [0, 0, 1, 0, 0]],
             [[1, 0, 2, 0, 0],
              [0, 1, 1, 1, 2],
              [1, 0, 0, 1, 1],
              [0, -3, 1, 1, 0],
              [0, 0, 2, 0, 0]]]])

    @pytest.mark.parametrize(
        "kernel_size,stride,padding,expected_result",
        [
            (
                (2, 2), (1, 1), (0, 0),
                np.array(
                    [[[[0.7500, 0.5000, 0.2500, 0.2500],
                       [0.2500, 0.7500, 1.0000, 0.7500],
                       [-0.5000, 0.0000, 1.0000, 0.7500],
                       [-0.2500, 0.2500, 0.7500, 0.2500]],

                      [[1.0000, 1.0000, 1.0000, 0.2500],
                       [0.7500, 0.5000, 0.2500, 0.2500],
                       [0.2500, 0.2500, 0.2500, 0.5000],
                       [0.5000, 0.7500, 0.5000, 0.2500]]],

                     [[[0.7500, 1.2500, 1.0000, 0.5000],
                       [0.2500, 1.0000, 1.2500, 1.0000],
                       [0.2500, 0.7500, 1.7500, 1.2500],
                       [0.2500, 0.7500, 1.5000, 0.7500]],

                      [[0.5000, 1.0000, 1.0000, 0.7500],
                       [0.5000, 0.5000, 0.7500, 1.2500],
                       [-0.5000, -0.5000, 0.7500, 0.7500],
                       [-0.7500, 0.0000, 1.0000, 0.2500]]]])
            ),
            (
                (2, 2), (1, 1), (1, 1),
                np.array(
                    [[[[0.2500, 0.5000, 0.0000, -0.2500, 0.0000, 0.0000],
                       [0.2500, 0.7500, 0.5000, 0.2500, 0.2500, 0.0000],
                       [0.0000, 0.2500, 0.7500, 1.0000, 0.7500, 0.2500],
                       [0.0000, -0.5000, 0.0000, 1.0000, 0.7500, 0.2500],
                       [0.0000, -0.2500, 0.2500, 0.7500, 0.2500, 0.0000],
                       [0.0000, 0.2500, 0.5000, 0.2500, 0.0000, 0.0000]],

                      [[0.2500, 0.2500, 0.2500, 0.5000, 0.2500, 0.0000],
                       [0.7500, 1.0000, 1.0000, 1.0000, 0.2500, 0.0000],
                       [0.5000, 0.7500, 0.5000, 0.2500, 0.2500, 0.2500],
                       [0.0000, 0.2500, 0.2500, 0.2500, 0.5000, 0.2500],
                       [0.0000, 0.5000, 0.7500, 0.5000, 0.2500, 0.0000],
                       [0.0000, 0.2500, 0.2500, 0.0000, 0.0000, 0.0000]]],

                     [[[0.2500, 0.5000, 0.5000, 0.2500, 0.0000, 0.0000],
                       [0.2500, 0.7500, 1.2500, 1.0000, 0.5000, 0.2500],
                       [0.0000, 0.2500, 1.0000, 1.2500, 1.0000, 0.5000],
                       [0.2500, 0.2500, 0.7500, 1.7500, 1.2500, 0.2500],
                       [0.2500, 0.2500, 0.7500, 1.5000, 0.7500, 0.0000],
                       [0.0000, 0.0000, 0.2500, 0.2500, 0.0000, 0.0000]],

                      [[0.2500, 0.2500, 0.5000, 0.5000, 0.0000, 0.0000],
                       [0.2500, 0.5000, 1.0000, 1.0000, 0.7500, 0.5000],
                       [0.2500, 0.5000, 0.5000, 0.7500, 1.2500, 0.7500],
                       [0.2500, -0.5000, -0.5000, 0.7500, 0.7500, 0.2500],
                       [0.0000, -0.7500, 0.0000, 1.0000, 0.2500, 0.0000],
                       [0.0000, 0.0000, 0.5000, 0.5000, 0.0000, 0.0000]]]])
            ),
            (
                (2, 2), (1, 1), (1, 0),
                np.array(
                    [[[[0.5000, 0.0000, -0.2500, 0.0000],
                       [0.7500, 0.5000, 0.2500, 0.2500],
                       [0.2500, 0.7500, 1.0000, 0.7500],
                       [-0.5000, 0.0000, 1.0000, 0.7500],
                       [-0.2500, 0.2500, 0.7500, 0.2500],
                       [0.2500, 0.5000, 0.2500, 0.0000]],

                      [[0.2500, 0.2500, 0.5000, 0.2500],
                       [1.0000, 1.0000, 1.0000, 0.2500],
                       [0.7500, 0.5000, 0.2500, 0.2500],
                       [0.2500, 0.2500, 0.2500, 0.5000],
                       [0.5000, 0.7500, 0.5000, 0.2500],
                       [0.2500, 0.2500, 0.0000, 0.0000]]],

                     [[[0.5000, 0.5000, 0.2500, 0.0000],
                       [0.7500, 1.2500, 1.0000, 0.5000],
                       [0.2500, 1.0000, 1.2500, 1.0000],
                       [0.2500, 0.7500, 1.7500, 1.2500],
                       [0.2500, 0.7500, 1.5000, 0.7500],
                       [0.0000, 0.2500, 0.2500, 0.0000]],

                      [[0.2500, 0.5000, 0.5000, 0.0000],
                       [0.5000, 1.0000, 1.0000, 0.7500],
                       [0.5000, 0.5000, 0.7500, 1.2500],
                       [-0.5000, -0.5000, 0.7500, 0.7500],
                       [-0.7500, 0.0000, 1.0000, 0.2500],
                       [0.0000, 0.5000, 0.5000, 0.0000]]]])
            ),
            (
                (2, 2), (1, 2), (1, 0),
                np.array(
                    [[[[0.5000, -0.2500],
                       [0.7500, 0.2500],
                       [0.2500, 1.0000],
                       [-0.5000, 1.0000],
                       [-0.2500, 0.7500],
                       [0.2500, 0.2500]],

                      [[0.2500, 0.5000],
                       [1.0000, 1.0000],
                       [0.7500, 0.2500],
                       [0.2500, 0.2500],
                       [0.5000, 0.5000],
                       [0.2500, 0.0000]]],

                     [[[0.5000, 0.2500],
                       [0.7500, 1.0000],
                       [0.2500, 1.2500],
                       [0.2500, 1.7500],
                       [0.2500, 1.5000],
                       [0.0000, 0.2500]],

                      [[0.2500, 0.5000],
                       [0.5000, 1.0000],
                       [0.5000, 0.7500],
                       [-0.5000, 0.7500],
                       [-0.7500, 1.0000],
                       [0.0000, 0.5000]]]])
            ),
            (
                (2, 3), (1, 2), (1, 0),
                np.array(
                    [[[[0.1667, -0.1667],
                       [0.5000, 0.1667],
                       [0.5000, 0.8333],
                       [0.0000, 0.8333],
                       [0.1667, 0.5000],
                       [0.3333, 0.1667]],

                      [[0.3333, 0.3333],
                       [1.1667, 0.6667],
                       [0.6667, 0.3333],
                       [0.1667, 0.3333],
                       [0.5000, 0.3333],
                       [0.1667, 0.0000]]],

                     [[[0.5000, 0.1667],
                       [1.0000, 0.8333],
                       [0.6667, 1.1667],
                       [0.6667, 1.3333],
                       [0.6667, 1.0000],
                       [0.1667, 0.1667]],

                      [[0.5000, 0.3333],
                       [0.8333, 1.0000],
                       [0.5000, 1.0000],
                       [-0.1667, 0.6667],
                       [0.0000, 0.6667],
                       [0.3333, 0.3333]]]])
            )
        ]
    )
    def test_avg_pool2d(self, HE, image, kernel_size, stride, padding, expected_result):
        encrypted_image = HE.encrypt_matrix(image)

        avg_layer = AveragePoolLayer(HE, kernel_size=kernel_size, stride=stride, padding=padding)

        result = HE.decrypt_matrix(avg_layer(encrypted_image))

        assert np.allclose(result, expected_result, rtol=1e-3)

    def test__avg(self, HE):
        """ Procedure:
                    1. Create a image in the form
                        [y, x]
                    2. Encrypt the image
                    3. Use __avg
                    ---------
                    Verification:
                    4. Verify the result is the expected
        """
        # Shape of image is [1, 1, 5, 5]. Needed to use encrypt_matrix.
        image = np.array([[[
            [1, -1, 1, 0, 0]
            , [0, 1, -2, 1, 0]
            , [0, 0, 1, 1, 1]
            , [0, -3, 1, 1, 0]
            , [0, 1, 1, 0, 0]
        ]]])

        encrypted_image = HE.encrypt_matrix(image)

        encrypted_result = avg._avg(HE, encrypted_image[0][0], kernel_size=(3, 3), stride=(2, 2))

        result = HE.decrypt_matrix(encrypted_result)

        expected_result = np.array(
            [[0.1111, 0.3333],
            [0.1111, 0.6667]])

        assert np.allclose(result, expected_result, 0.001)


class TestsSquareLayer:
    @pytest.fixture
    def HE(self):
        HE = MockHE()
        return HE

    def test_square_layer(self, HE):
        image = np.array([
            [
                [
                    [1, 1, 1, 0, 0]
                    , [0, 1, 1, 1, 0]
                    , [0, 0, 1, 1, 1]
                    , [0, 0, 1, 1, 0]
                    , [0, 1, 1, 0, 0]
                ],
                [
                    [1, 0, 1, 1, 0]
                    , [0, 1, 2, 0, 0]
                    , [0, 0, 1, 0, 1]
                    , [0, 1, 1, 1, 0]
                    , [0, 1, 0, 0, 0]],
            ],
            [
                [
                    [1, 1, 1, 0, 0]
                    , [0, 1, 2, 1, 1]
                    , [0, 0, 1, 1, 1]
                    , [1, 0, 2, 1, 0]
                    , [0, 0, 1, 0, 0]
                ],
                [
                    [-1, 0, 2, 0, 0]
                    , [0, 1, 1, 1, 2]
                    , [1, 0, 0, 1, 1]
                    , [0, 0, 1, 1, 0]
                    , [0, 0, 2, 0, 0]
                ]
            ]
        ])

        expected_result = np.array([
            [
                [
                    [1, 1, 1, 0, 0]
                    , [0, 1, 1, 1, 0]
                    , [0, 0, 1, 1, 1]
                    , [0, 0, 1, 1, 0]
                    , [0, 1, 1, 0, 0]
                ],
                [
                    [1, 0, 1, 1, 0]
                    , [0, 1, 4, 0, 0]
                    , [0, 0, 1, 0, 1]
                    , [0, 1, 1, 1, 0]
                    , [0, 1, 0, 0, 0]],
            ],
            [
                [
                    [1, 1, 1, 0, 0]
                    , [0, 1, 4, 1, 1]
                    , [0, 0, 1, 1, 1]
                    , [1, 0, 4, 1, 0]
                    , [0, 0, 1, 0, 0]
                ],
                [
                    [1, 0, 4, 0, 0]
                    , [0, 1, 1, 1, 4]
                    , [1, 0, 0, 1, 1]
                    , [0, 0, 1, 1, 0]
                    , [0, 0, 4, 0, 0]
                ]
            ]
        ])

        encrypted_image = HE.encrypt_matrix(image)
        encrypted_result = sq.square(HE, encrypted_image)

        result = HE.decrypt_matrix(encrypted_result)

        assert np.allclose(expected_result, result, rtol=1e-3)

    def test_square_layer2D(self, HE):
        image = np.array([
                    [1, -2, 1, 3, 0]
                    , [0, 1, 4, 1, 0]
                    , [0, 2, 1, 1, 1]
                ])

        expected_result = np.array([
                    [1, 4, 1, 9, 0]
                    , [0, 1, 16, 1, 0]
                    , [0, 4, 1, 1, 1]
                ])

        encrypted_image = HE.encrypt_matrix(image)
        encrypted_result = sq.square(HE, encrypted_image)

        result = HE.decrypt_matrix(encrypted_result)

        assert np.allclose(expected_result, result, rtol=1e-3)


class TestsRencryptionLayer:
    @pytest.fixture
    def HE(self):
        HE = MockHE()
        return HE

    def test_rencryption_layer(self, HE):
        image = np.array([
                      [1, 1, -1, 0, 0],
                      [1, 0, 1, 1, 0]])

        enc_image = HE.encrypt_matrix(image)
        rencryption_layer = RencryptionLayer(HE)

        enc_result = rencryption_layer(enc_image)

        result = HE.decrypt_matrix(enc_result)
        assert np.allclose(image, result)


class TestsReshapeLayer:
    def test_reshape_layer(self):
        image = np.array([
            [
                [
                    [1, 1, -1, 0, 0]
                    , [0, 1, 1, 1, 0]
                    , [0, 0, 1, 1, 1]
                    , [0, -2, 1, 1, 0]
                    , [0, 1, 1, 0, 0]
                ],
                [
                    [1, 0, 1, 1, 0]
                    , [2, 1, 2, 0, 0]
                    , [0, 0, -1, 0, 1]
                    , [0, 1, 1, 1, 0]
                    , [0, 1, 0, 0, 0]],
            ],
            [
                [
                    [1, 1, 1, 0, 0]
                    , [0, 1, 2, 1, 1]
                    , [0, 0, 1, 1, 1]
                    , [1, 0, 2, 3, 0]
                    , [0, 0, 1, 0, 0]
                ],
                [
                    [1, 0, 2, 0, 0]
                    , [0, 1, 1, 1, 2]
                    , [1, 0, 0, 1, 1]
                    , [0, -3, 1, 1, 0]
                    , [0, 0, 2, 0, 0]
                ]
            ]
        ])

        reshape_layer = FlattenLayer()
        result = reshape_layer(image)

        assert result.shape == (2, 50)

