import pytest

from keras.utils.test_utils import layer_test
from keras.utils.test_utils import keras_test
from keras.layers import local


@keras_test
def test_locallyconnected_1d():
    num_samples = 2
    num_steps = 8
    input_dim = 5
    filter_length = 3
    filters = 4

    for padding in ['valid']:
        for strides in [1]:
            if padding == 'same' and strides != 1:
                continue
            layer_test(local.LocallyConnected1D,
                       kwargs={'filters': filters,
                               'kernel_size': filter_length,
                               'padding': padding,
                               'strides': strides},
                       input_shape=(num_samples, num_steps, input_dim))

            layer_test(local.LocallyConnected1D,
                       kwargs={'filters': filters,
                               'kernel_size': filter_length,
                               'padding': padding,
                               'kernel_regularizer': 'l2',
                               'bias_regularizer': 'l2',
                               'activity_regularizer': 'l2',
                               'strides': strides},
                       input_shape=(num_samples, num_steps, input_dim))


@keras_test
def test_locallyconnected_2d():
    num_samples = 8
    filters = 3
    stack_size = 4
    num_row = 6
    num_col = 10

    for padding in ['valid']:
        for strides in [(1, 1), (2, 2)]:
            if padding == 'same' and strides != (1, 1):
                continue

            layer_test(local.LocallyConnected2D,
                       kwargs={'filters': filters,
                               'kernel_size': 3,
                               'padding': padding,
                               'kernel_regularizer': 'l2',
                               'bias_regularizer': 'l2',
                               'activity_regularizer': 'l2',
                               'strides': strides,
                               'data_format': 'channels_last'},
                       input_shape=(num_samples, num_row, num_col, stack_size))

            layer_test(local.LocallyConnected2D,
                       kwargs={'filters': filters,
                               'kernel_size': (3, 3),
                               'padding': padding,
                               'kernel_regularizer': 'l2',
                               'bias_regularizer': 'l2',
                               'activity_regularizer': 'l2',
                               'strides': strides,
                               'data_format': 'channels_first'},
                       input_shape=(num_samples, stack_size, num_row, num_col))


if __name__ == '__main__':
    pytest.main([__file__])
