import cntk as C
import cntkx as Cx
from cntkx.layers import QRNN, SinusoidalPositionalEmbedding, SpatialPyramidPooling, GatedLinearUnit
from cntkx.layers import BertEmbeddings, PositionalEmbedding, SequentialAveragePooling
from cntkx.layers import PreTrainedBertEmbeddings, PositionwiseFeedForward, SequentialMaxPooling
import numpy as np


def test_qrnn():
    HIDDEN_DIM = 20
    i = C.sequence.input_variable(10)
    qrnn = QRNN(hidden_dim=HIDDEN_DIM)(i)

    assert qrnn.shape[0] == HIDDEN_DIM


def test_qrnn1():
    input_dim = 3
    even_length = 10
    odd_length = 15
    i = C.sequence.input_variable(input_dim)
    qrnn = QRNN(hidden_dim=50)(i)

    n1 = np.random.random((odd_length, input_dim)).astype(np.float32)
    n2 = np.random.random((even_length, input_dim)).astype(np.float32)

    qrnn.eval({i: [n1, n2]})


def test_sinusoidal_positional_embedding():
    seq = 50
    dim = 100
    a = C.sequence.input_variable(dim)
    b = SinusoidalPositionalEmbedding(dim)(a)

    assert b.shape == (dim, )

    n = np.random.random((1, seq, dim)).astype(np.float32)
    r = b.eval({a: n})

    dim = 99
    a = C.sequence.input_variable(dim)
    b = SinusoidalPositionalEmbedding(dim)(a)

    assert b.shape == (dim, )

    n = np.random.random((1, seq, dim)).astype(np.float32)
    r = b.eval({a: n})

    assert np.sum(r[0][:, -1]) == 0, "last dimension should be all zeros if dimension is odd"

    # plt.imshow(r[0])  # to view image to confirm correctness
    # plt.show()


def test_spatial_pyramid_pooling():
    # test 1
    n = np.random.random((3, 3, 32, 32)).astype(np.float32)
    a = C.input_variable((3, 32, 32))
    b = SpatialPyramidPooling((1, 2, 4))(a)

    assert b.shape == (3 * (4 * 4 + 2 * 2 + 1),)
    b.eval({a: n})

    # test 2
    n = np.random.random((3, 3, 35, 35)).astype(np.float32)
    a = C.input_variable((3, 35, 35))
    b = SpatialPyramidPooling((1, 2, 4))(a)

    assert b.shape == (3 * (4 * 4 + 2 * 2 + 1),)
    b.eval({a: n})

    # test 3
    n = np.random.random((3, 3, 35, 35)).astype(np.float32)
    a = C.input_variable((3, 35, 35))
    b = SpatialPyramidPooling((1, 3, 5))(a)

    assert b.shape == (3 * (5 * 5 + 3 * 3 + 1), )
    b.eval({a: n})

    # test 3
    n = np.random.random((3, 3, 41, 41)).astype(np.float32)
    a = C.input_variable((3, 41, 41))
    b = SpatialPyramidPooling((1, 3, 5))(a)

    assert b.shape == (3 * (5 * 5 + 3 * 3 + 1),)
    b.eval({a: n})


def test_gated_linear_unit():
    input_dim = 10
    even_length = 10
    odd_length = 15

    a = C.sequence.input_variable(input_dim)
    b = GatedLinearUnit(2, 20)(a)

    assert b.shape == (20, )

    n1 = np.random.random((odd_length, input_dim)).astype(np.float32)
    n2 = np.random.random((even_length, input_dim)).astype(np.float32)

    b.eval({a: [n1, n2]})


def test_positional_embedding():
    max_seq_length = 100
    hidden_dim = 120
    a = C.sequence.input_variable(12)
    b = PositionalEmbedding(max_seq_length, hidden_dim)(a)

    assert b.shape == (hidden_dim, )

    n1 = np.random.random((3, 12)).astype(np.float32)
    n2 = np.random.random((6, 12)).astype(np.float32)
    b.eval({a: [n1, n2]})


def test_bert_embeddings():
    max_seq_length = 512
    hidden_dim = 768

    text_tensor = C.sequence.input_variable(100)
    token_type_tensor = C.sequence.input_variable(2)
    b = BertEmbeddings(max_seq_length, hidden_dim, 0.1)(text_tensor, token_type_tensor)

    assert b.shape == (hidden_dim, )

    n1 = np.random.random((3, 100)).astype(np.float32)
    n2 = np.random.random((6, 100)).astype(np.float32)

    m1 = np.random.random((3, 2)).astype(np.float32)
    m2 = np.random.random((6, 2)).astype(np.float32)
    b.eval({text_tensor: [n1, n2], token_type_tensor: [m1, m2]})


def test_pretrained_bert_embeddings_learnable():
    """ tested to work with 'uncased_L-12_H-768_A-12' """
    text_tensor = C.sequence.input_variable(30522)
    token_type_tensor = C.sequence.input_variable(2)
    filepath_to_tf_bert_model = "../../../pretrained models/BERT/uncased/bert_model.ckpt"

    embeddings = PreTrainedBertEmbeddings(filepath_to_tf_bert_model, 0.1, name='bert')
    b = embeddings(text_tensor, token_type_tensor)

    assert b.shape == (768,)

    n1 = np.random.random((3, 30522)).astype(np.float32)
    n2 = np.random.random((6, 30522)).astype(np.float32)

    m1 = np.random.random((3, 2)).astype(np.float32)
    m2 = np.random.random((6, 2)).astype(np.float32)
    b.eval({text_tensor: [n1, n2], token_type_tensor: [m1, m2]})

    embed = b.bert.word_embeddings
    assert embed.shape == (768,)

    embed = b.bert.position_embeddings
    assert embed.parameters[0].shape == (512, 768)

    embed = b.bert.token_type_embeddings
    assert embed.parameters[0].shape == (2, 768)


def test_positionwise_feedforward():
    model_dim = 20
    inner_dim = 30

    a = C.sequence.input_variable(10)
    b = PositionwiseFeedForward(model_dim, inner_dim, 0.1)(a)

    assert b.shape == (model_dim, )

    n1 = np.random.random((3, 10)).astype(np.float32)
    n2 = np.random.random((6, 10)).astype(np.float32)

    b.eval({a: [n1, n2]})


def test_sequential_max_pooling():
    # ===================================================================
    # sequential max pool across rgb images with width in sequence axis
    # ===================================================================
    a = C.sequence.input_variable((2, 4))
    b = SequentialMaxPooling(filter_shape=(2, 2), strides=(2, 2), pad=False)(a)

    n = np.ascontiguousarray(np.arange(2 * 5 * 4).reshape((1, 5, 2, 4)).astype(np.float32))
    output = b.eval({a: n})

    assert isinstance(output, list) and len(output) == 1
    output = output[0]

    a = C.input_variable((2, 4, 5))
    b = C.layers.MaxPooling(filter_shape=(2, 2), strides=(2, 2), pad=False)(a)

    n = np.arange(2 * 5 * 4).reshape((1, 5, 2, 4)).astype(np.float32)
    n = np.ascontiguousarray(np.moveaxis(n, 1, -1))

    desired = b.eval({a: n})
    desired = np.squeeze(np.moveaxis(desired, -1, 1))

    np.testing.assert_almost_equal(output[:2, ...], desired)

    # BUGBUG: Sequential maxpooling will 'right pad' on sequential axis
    # BUGBUG: once fixed, this assertion should fail
    assert output.shape[0] == 3 and desired.shape[0] == 2
    assert output.shape[0] != desired.shape[0], "Due to bug, sequence length is different between desired and output"

    # ===================================================================
    # sequential max pool across rgb images with width in sequence axis
    # ===================================================================
    a = C.sequence.input_variable((3, 25))
    b = SequentialMaxPooling(filter_shape=(2, 2), strides=(2, 2), pad=False)(a)

    assert b.shape == (3, 12)

    n = np.ascontiguousarray(np.arange(3 * 6 * 25).reshape((1, 6, 3, 25)).astype(np.float32))
    output = b.eval({a: n})

    assert isinstance(output, list) and len(output) == 1
    output = output[0]

    a = C.input_variable((3, 25, 6))
    b = C.layers.MaxPooling(filter_shape=(2, 2), strides=(2, 2), pad=False)(a)

    n = np.arange(3 * 6 * 25).reshape((1, 6, 3, 25)).astype(np.float32)
    n = np.ascontiguousarray(np.moveaxis(n, 1, -1))

    desired = b.eval({a: n})
    desired = np.squeeze(np.moveaxis(desired, -1, 1))

    np.testing.assert_almost_equal(output, desired)

    # ===================================================================
    # sequential max pool across rgb images with width in sequence axis
    # ===================================================================
    a = C.sequence.input_variable((3, 25))
    b = SequentialMaxPooling(filter_shape=(3, 3), strides=(2, 2), pad=False)(a)

    n = np.ascontiguousarray(np.arange(3 * 6 * 25).reshape((1, 6, 3, 25)).astype(np.float32))
    output = b.eval({a: n})

    assert isinstance(output, list) and len(output) == 1
    output = output[0]

    a = C.input_variable((3, 25, 6))
    b = C.layers.MaxPooling(filter_shape=(3, 3), strides=(2, 2), pad=False)(a)

    n = np.arange(3 * 6 * 25).reshape((1, 6, 3, 25)).astype(np.float32)
    n = np.ascontiguousarray(np.moveaxis(n, 1, -1))

    desired = b.eval({a: n})
    desired = np.squeeze(np.moveaxis(desired, -1, 1))

    # BUGBUG: Sequential maxpooling will 'right pad' on sequential axis
    # BUGBUG: once fixed, this assertion should fail
    np.testing.assert_almost_equal(output[:2], desired)

    # ===================================================================
    # sequential max pool across rgb images with width in sequence axis
    # ===================================================================
    a = C.sequence.input_variable((3, 25))
    b = SequentialMaxPooling(filter_shape=(3, 3), strides=(2, 2), pad=True)(a)

    assert b.shape == (3, 13)

    n = np.ascontiguousarray(np.arange(3 * 6 * 25).reshape((1, 6, 3, 25)).astype(np.float32))
    output = b.eval({a: n})

    assert isinstance(output, list) and len(output) == 1
    output = output[0]

    a = C.input_variable((3, 25, 6))
    b = C.layers.MaxPooling(filter_shape=(3, 3), strides=(2, 2), pad=True)(a)

    n = np.arange(3 * 6 * 25).reshape((1, 6, 3, 25)).astype(np.float32)
    n = np.ascontiguousarray(np.moveaxis(n, 1, -1))

    desired = b.eval({a: n})
    desired = np.squeeze(np.moveaxis(desired, -1, 1))

    np.testing.assert_almost_equal(output, desired)

    # ===================================================================
    # sequential max pool across rgb images with width in sequence axis
    # ===================================================================
    a = C.sequence.input_variable((3, 25))
    b = SequentialMaxPooling(filter_shape=(4, 4), strides=(2, 2), pad=True)(a)

    n = np.ascontiguousarray(np.arange(3 * 6 * 25).reshape((1, 6, 3, 25)).astype(np.float32))
    output = b.eval({a: n})

    assert isinstance(output, list) and len(output) == 1
    output = output[0]

    a = C.input_variable((3, 25, 6))
    b = C.layers.MaxPooling(filter_shape=(4, 4), strides=(2, 2), pad=True)(a)

    n = np.arange(3 * 6 * 25).reshape((1, 6, 3, 25)).astype(np.float32)
    n = np.ascontiguousarray(np.moveaxis(n, 1, -1))

    desired = b.eval({a: n})
    desired = np.squeeze(np.moveaxis(desired, -1, 1))

    np.testing.assert_almost_equal(output, desired)

    # ===================================================================
    # sequential max pool across a sequence of vector or B&W images
    # ===================================================================
    a = C.sequence.input_variable((25, ))
    b = SequentialMaxPooling(filter_shape=(4,), strides=(2,), pad=True)(a)

    n = np.ascontiguousarray(np.arange(1 * 6 * 25).reshape((1, 6, 25)).astype(np.float32))
    output = b.eval({a: n})

    assert isinstance(output, list) and len(output) == 1
    output = output[0]

    a = C.input_variable((25, 6))
    b = C.layers.MaxPooling(filter_shape=(4, ), strides=(2, ), pad=True)(a)

    n = np.arange(1 * 6 * 25).reshape((1, 6, 25)).astype(np.float32)
    n = np.ascontiguousarray(np.moveaxis(n, 1, -1))

    desired = b.eval({a: n})
    desired = np.squeeze(np.moveaxis(desired, -1, 1))

    np.testing.assert_almost_equal(output, desired)


def test_sequential_average_pooling():
    # ===================================================================
    # sequential max pool across rgb images with width in sequence axis
    # ===================================================================
    a = C.sequence.input_variable((2, 4))
    b = SequentialAveragePooling(filter_shape=(2, 2), strides=(2, 2), pad=False)(a)

    n = np.ascontiguousarray(np.arange(2 * 5 * 4).reshape((1, 5, 2, 4)).astype(np.float32))
    output = b.eval({a: n})

    assert isinstance(output, list) and len(output) == 1
    output = output[0]

    a = C.input_variable((2, 4, 5))
    b = C.layers.AveragePooling(filter_shape=(2, 2), strides=(2, 2), pad=False)(a)

    n = np.arange(2 * 5 * 4).reshape((1, 5, 2, 4)).astype(np.float32)
    n = np.ascontiguousarray(np.moveaxis(n, 1, -1))

    desired = b.eval({a: n})
    desired = np.squeeze(np.moveaxis(desired, -1, 1))

    np.testing.assert_almost_equal(output[:2, ...], desired)

    # BUGBUG: Sequential AveragePooling will 'right pad' on sequential axis
    # BUGBUG: once fixed, this assertion should fail
    assert output.shape[0] == 3 and desired.shape[0] == 2
    assert output.shape[0] != desired.shape[0], "Due to bug, sequence length is different between desired and output"

    # ===================================================================
    # sequential max pool across rgb images with width in sequence axis
    # ===================================================================
    a = C.sequence.input_variable((3, 25))
    b = SequentialAveragePooling(filter_shape=(2, 2), strides=(2, 2), pad=False)(a)

    assert b.shape == (3, 12)

    n = np.ascontiguousarray(np.arange(3 * 6 * 25).reshape((1, 6, 3, 25)).astype(np.float32))
    output = b.eval({a: n})

    assert isinstance(output, list) and len(output) == 1
    output = output[0]

    a = C.input_variable((3, 25, 6))
    b = C.layers.AveragePooling(filter_shape=(2, 2), strides=(2, 2), pad=False)(a)

    n = np.arange(3 * 6 * 25).reshape((1, 6, 3, 25)).astype(np.float32)
    n = np.ascontiguousarray(np.moveaxis(n, 1, -1))

    desired = b.eval({a: n})
    desired = np.squeeze(np.moveaxis(desired, -1, 1))

    np.testing.assert_almost_equal(output, desired)

    # ===================================================================
    # sequential max pool across rgb images with width in sequence axis
    # ===================================================================
    a = C.sequence.input_variable((3, 25))
    b = SequentialAveragePooling(filter_shape=(3, 3), strides=(2, 2), pad=False)(a)

    n = np.ascontiguousarray(np.arange(3 * 6 * 25).reshape((1, 6, 3, 25)).astype(np.float32))
    output = b.eval({a: n})

    assert isinstance(output, list) and len(output) == 1
    output = output[0]

    a = C.input_variable((3, 25, 6))
    b = C.layers.AveragePooling(filter_shape=(3, 3), strides=(2, 2), pad=False)(a)

    n = np.arange(3 * 6 * 25).reshape((1, 6, 3, 25)).astype(np.float32)
    n = np.ascontiguousarray(np.moveaxis(n, 1, -1))

    desired = b.eval({a: n})
    desired = np.squeeze(np.moveaxis(desired, -1, 1))

    # BUGBUG: Sequential AveragePooling will 'right pad' on sequential axis
    # BUGBUG: once fixed, this assertion should fail
    np.testing.assert_almost_equal(output[:2], desired)

    # ===================================================================
    # sequential max pool across rgb images with width in sequence axis
    # ===================================================================
    a = C.sequence.input_variable((3, 25))
    b = SequentialAveragePooling(filter_shape=(3, 3), strides=(2, 2), pad=True)(a)

    assert b.shape == (3, 13)

    n = np.ascontiguousarray(np.arange(3 * 6 * 25).reshape((1, 6, 3, 25)).astype(np.float32))
    output = b.eval({a: n})

    assert isinstance(output, list) and len(output) == 1
    output = output[0]

    a = C.input_variable((3, 25, 6))
    b = C.layers.AveragePooling(filter_shape=(3, 3), strides=(2, 2), pad=True)(a)

    n = np.arange(3 * 6 * 25).reshape((1, 6, 3, 25)).astype(np.float32)
    n = np.ascontiguousarray(np.moveaxis(n, 1, -1))

    desired = b.eval({a: n})
    desired = np.squeeze(np.moveaxis(desired, -1, 1))

    # ignore the first seq item because SequentialAveragePooling included padding in average calculation
    np.testing.assert_almost_equal(output[1:], desired[1:])

    # ===================================================================
    # sequential max pool across rgb images with width in sequence axis
    # ===================================================================
    a = C.sequence.input_variable((3, 25))
    b = SequentialAveragePooling(filter_shape=(4, 4), strides=(2, 2), pad=True)(a)

    n = np.ascontiguousarray(np.arange(3 * 6 * 25).reshape((1, 6, 3, 25)).astype(np.float32))
    output = b.eval({a: n})

    assert isinstance(output, list) and len(output) == 1
    output = output[0]

    a = C.input_variable((3, 25, 6))
    b = C.layers.AveragePooling(filter_shape=(4, 4), strides=(2, 2), pad=True)(a)

    n = np.arange(3 * 6 * 25).reshape((1, 6, 3, 25)).astype(np.float32)
    n = np.ascontiguousarray(np.moveaxis(n, 1, -1))

    desired = b.eval({a: n})
    desired = np.squeeze(np.moveaxis(desired, -1, 1))

    # ignore the first and last seq item because SequentialAveragePooling included padding in average calculation
    np.testing.assert_almost_equal(output[1:-1], desired[1:-1])

    # ===================================================================
    # sequential max pool across a sequence of vector or B&W images
    # ===================================================================
    a = C.sequence.input_variable((25, ))
    b = SequentialAveragePooling(filter_shape=(4,), strides=(2,), pad=True)(a)

    n = np.ascontiguousarray(np.arange(1 * 6 * 25).reshape((1, 6, 25)).astype(np.float32))
    output = b.eval({a: n})

    assert isinstance(output, list) and len(output) == 1
    output = output[0]

    a = C.input_variable((25, 6))
    b = C.layers.AveragePooling(filter_shape=(4, ), strides=(2, ), pad=True)(a)

    n = np.arange(1 * 6 * 25).reshape((1, 6, 25)).astype(np.float32)
    n = np.ascontiguousarray(np.moveaxis(n, 1, -1))

    desired = b.eval({a: n})
    desired = np.squeeze(np.moveaxis(desired, -1, 1))

    # ignore the first and last seq item because SequentialAveragePooling included padding in average calculation
    np.testing.assert_almost_equal(output[1:-1], desired[1:-1])


def test_convolution_2d():
    # --------------------------------------------------------------------------
    # Normal use case
    # --------------------------------------------------------------------------
    image_shape = (3, 32, 32)
    a = C.input_variable(image_shape)
    b = C.layers.Convolution2D(filter_shape=3, num_filters=16, reduction_rank=1, pad=True)(a)
    c = Cx.layers.Convolution2D(filter_shape=3, num_filters=16, reduction_rank=1, pad=True,
                                init=b.W.value, init_bias=b.b.value)(a)

    assert b.shape == c.shape

    n = np.random.random((1, ) + image_shape).astype(np.float32)

    desired = b.eval({a: n})
    actual = c.eval({a: n})

    np.testing.assert_equal(actual, desired)

    image_shape = (3, 32, 32)
    a = C.input_variable(image_shape)
    b = C.layers.Convolution2D(filter_shape=3, num_filters=16, reduction_rank=1, pad=(True, True))(a)
    c = Cx.layers.Convolution2D(filter_shape=3, num_filters=16, reduction_rank=1, pad=(True, True),
                                init=b.W.value, init_bias=b.b.value)(a)

    assert b.shape == c.shape

    n = np.random.random((1,) + image_shape).astype(np.float32)

    desired = b.eval({a: n})
    actual = c.eval({a: n})

    np.testing.assert_equal(actual, desired)

    # --------------------------------------------------------------------------
    # Normal use case, no padding with stride 2
    # --------------------------------------------------------------------------
    image_shape = (3, 32, 32)
    a = C.input_variable(image_shape)
    b = C.layers.Convolution2D(filter_shape=3, num_filters=16, reduction_rank=1, pad=False, strides=2)(a)
    c = Cx.layers.Convolution2D(filter_shape=3, num_filters=16, reduction_rank=1, pad=False, strides=2,
                                init=b.W.value, init_bias=b.b.value)(a)

    assert b.shape == c.shape

    n = np.random.random((1,) + image_shape).astype(np.float32)

    desired = b.eval({a: n})
    actual = c.eval({a: n})

    np.testing.assert_equal(actual, desired)

    # --------------------------------------------------------------------------
    # Normal use case, reduction rank = 0
    # --------------------------------------------------------------------------
    image_shape = (32, 32)
    a = C.input_variable(image_shape)
    b = C.layers.Convolution2D(filter_shape=3, num_filters=16, reduction_rank=0, pad=False, strides=2)(a)
    c = Cx.layers.Convolution2D(filter_shape=3, num_filters=16, reduction_rank=0, pad=False, strides=2,
                                init=b.W.value, init_bias=b.b.value)(a)

    assert b.shape == c.shape

    n = np.random.random((1,) + image_shape).astype(np.float32)

    desired = b.eval({a: n})
    actual = c.eval({a: n})

    np.testing.assert_equal(actual, desired)

    # --------------------------------------------------------------------------
    # Normal use case, reduction rank = 0 with padding
    # --------------------------------------------------------------------------
    image_shape = (32, 32)
    a = C.input_variable(image_shape)
    b = C.layers.Convolution2D(filter_shape=3, num_filters=16, reduction_rank=0, pad=True, strides=2)(a)
    c = Cx.layers.Convolution2D(filter_shape=3, num_filters=16, reduction_rank=0, pad=True, strides=2,
                                init=b.W.value, init_bias=b.b.value)(a)

    assert b.shape == c.shape

    n = np.random.random((1,) + image_shape).astype(np.float32)

    desired = b.eval({a: n})
    actual = c.eval({a: n})

    np.testing.assert_equal(actual, desired)

    # --------------------------------------------------------------------------
    # Normal use case, reduction rank = 0 with pad & stride on one dim
    # --------------------------------------------------------------------------
    image_shape = (32, 32)
    a = C.input_variable(image_shape)
    b = C.layers.Convolution2D(filter_shape=3, num_filters=16, reduction_rank=0, pad=(True, False), strides=(2, 1))(a)
    c = Cx.layers.Convolution2D(filter_shape=3, num_filters=16, reduction_rank=0, pad=(True, False), strides=(2, 1),
                                init=b.W.value, init_bias=b.b.value)(a)

    assert b.shape == c.shape

    n = np.random.random((1,) + image_shape).astype(np.float32)

    desired = b.eval({a: n})
    actual = c.eval({a: n})

    np.testing.assert_equal(actual, desired)

    # --------------------------------------------------------------------------
    # Normal use case, reduction rank = 0 with pad & stride on one dim and no output depth
    # --------------------------------------------------------------------------
    image_shape = (32, 32)
    a = C.input_variable(image_shape)
    b = C.layers.Convolution2D(filter_shape=3, num_filters=None, reduction_rank=0, pad=(True, False), strides=(2, 1))(a)
    c = Cx.layers.Convolution2D(filter_shape=3, num_filters=None, reduction_rank=0, pad=(True, False), strides=(2, 1),
                                init=b.W.value, init_bias=b.b.value)(a)

    assert b.shape == c.shape

    n = np.random.random((1,) + image_shape).astype(np.float32)

    desired = b.eval({a: n})
    actual = c.eval({a: n})

    np.testing.assert_equal(actual, desired)

    # --------------------------------------------------------------------------
    # Image is same size as kernel
    # --------------------------------------------------------------------------
    image_shape = (3, 3, 3)
    a = C.input_variable(image_shape)
    b = C.layers.Convolution2D(filter_shape=3, num_filters=16, reduction_rank=1, pad=True)(a)
    c = Cx.layers.Convolution2D(filter_shape=3, num_filters=16, reduction_rank=1, pad=True,
                                init=b.W.value, init_bias=b.b.value)(a)

    assert b.shape == c.shape

    n = np.random.random((1, ) + image_shape).astype(np.float32)

    desired = b.eval({a: n})
    actual = c.eval({a: n})

    np.testing.assert_equal(actual, desired)

    # --------------------------------------------------------------------------
    # Image is same size as kernel, no padding
    # --------------------------------------------------------------------------
    image_shape = (3, 3, 3)
    a = C.input_variable(image_shape)
    b = C.layers.Convolution2D(filter_shape=3, num_filters=16, reduction_rank=1, pad=False)(a)
    c = Cx.layers.Convolution2D(filter_shape=3, num_filters=16, reduction_rank=1, pad=False,
                                init=b.W.value, init_bias=b.b.value)(a)

    assert b.shape == c.shape

    n = np.random.random((1,) + image_shape).astype(np.float32)

    desired = b.eval({a: n})
    actual = c.eval({a: n})

    np.testing.assert_equal(actual, desired)

    # --------------------------------------------------------------------------
    # Check kernel initialised same shape
    # --------------------------------------------------------------------------
    image_shape = (32, 32)
    a = C.input_variable(image_shape)
    b = C.layers.Convolution2D(filter_shape=(3, 2), num_filters=None, reduction_rank=0, pad=(True, False), strides=(2, 1))(a)
    c = Cx.layers.Convolution2D(filter_shape=(3, 2), num_filters=None, reduction_rank=0, pad=(True, False), strides=(2, 1))(a)

    assert b.shape == c.shape
    assert b.W.shape == c.W.shape
    assert b.b.shape == c.b.shape

    image_shape = (3, 32, 32)
    a = C.input_variable(image_shape)
    b = C.layers.Convolution2D(filter_shape=(3, 4), num_filters=None, reduction_rank=1, pad=(True, False), strides=(2, 1))(a)
    c = Cx.layers.Convolution2D(filter_shape=(3, 4), num_filters=None, reduction_rank=1, pad=(True, False), strides=(2, 1))(a)

    assert b.shape == c.shape
    assert b.W.shape == c.W.shape
    assert b.b.shape == c.b.shape


def test_sequential_convolution():
    # --------------------------------------------------------------------------
    # Normal use case - image check kernel initialisation shape
    # --------------------------------------------------------------------------
    image_shape = (3, 32)  # rgb image of variable width
    a = C.sequence.input_variable(image_shape)
    b = C.layers.SequentialConvolution(filter_shape=(3, 3), num_filters=16, reduction_rank=1, pad=True)(a)
    c = Cx.layers.SequentialConvolution(filter_shape=(3, 3), num_filters=16, reduction_rank=1, pad=True)(a)

    assert b.shape == c.shape
    assert b.W.shape == c.W.shape
    assert b.b.shape == c.b.shape

    image_shape = (32, )  # black and white image of variable width
    a = C.sequence.input_variable(image_shape)
    b = C.layers.SequentialConvolution(filter_shape=(2, 2), num_filters=16, reduction_rank=0, pad=True)(a)
    c = Cx.layers.SequentialConvolution(filter_shape=(2, 2), num_filters=16, reduction_rank=0, pad=True)(a)

    assert b.shape == c.shape
    assert b.W.shape == c.W.shape
    assert b.b.shape == c.b.shape

    image_shape = (32,)  # text vector of variable sequence length
    a = C.sequence.input_variable(image_shape)
    b = C.layers.SequentialConvolution(filter_shape=(2, 2), num_filters=16, reduction_rank=1, pad=True)(a)
    c = Cx.layers.SequentialConvolution(filter_shape=(2, 2), num_filters=16, reduction_rank=1, pad=True)(a)

    assert b.shape == c.shape
    assert b.W.shape == c.W.shape
    assert b.b.shape == c.b.shape

    # --------------------------------------------------------------------------
    # Normal use case - image
    # --------------------------------------------------------------------------
    kernel_init = C.glorot_normal(seed=5)
    image_shape = (3, 32)
    width = 40
    a = C.sequence.input_variable(image_shape)
    b = C.layers.SequentialConvolution(filter_shape=(3, 3), num_filters=16, reduction_rank=1, pad=True, init=kernel_init)(a)
    c = Cx.layers.SequentialConvolution(filter_shape=(3, 3), num_filters=16, reduction_rank=1, pad=True, init=kernel_init)(a)

    assert b.shape == c.shape

    n = np.random.random((1, width) + image_shape).astype(np.float32)

    desired = b.eval({a: n})
    actual = c.eval({a: n})

    if isinstance(desired, list) and len(desired) == 1:
        desired = desired[0]

    if isinstance(actual, list) and len(actual) == 1:
        actual = actual[0]

    np.testing.assert_equal(actual, desired)

    # --------------------------------------------------------------------------
    # Normal use case - image, mix padding across axis
    # --------------------------------------------------------------------------
    kernel_init = C.glorot_normal(seed=5)
    image_shape = (3, 32)
    width = 40
    a = C.sequence.input_variable(image_shape)
    b = C.layers.SequentialConvolution(filter_shape=(3, 3), num_filters=16, reduction_rank=1, pad=(False, True), init=kernel_init)(a)
    c = Cx.layers.SequentialConvolution(filter_shape=(3, 3), num_filters=16, reduction_rank=1, pad=(False, True), init=kernel_init)(a)

    assert b.shape == c.shape

    n = np.random.random((1, width) + image_shape).astype(np.float32)

    desired = b.eval({a: n})
    actual = c.eval({a: n})

    if isinstance(desired, list) and len(desired) == 1:
        desired = desired[0]

    if isinstance(actual, list) and len(actual) == 1:
        actual = actual[0]

    np.testing.assert_equal(actual, desired)

    # --------------------------------------------------------------------------
    # Normal use case - vector data for qrnn implementation
    # --------------------------------------------------------------------------
    kernel_init = C.glorot_normal(seed=5)
    image_shape = (25,)
    width = 40
    a = C.sequence.input_variable(image_shape)
    b = C.layers.SequentialConvolution(filter_shape=(2, ), num_filters=16, reduction_rank=1, pad=True, init=kernel_init)(a)
    c = Cx.layers.SequentialConvolution(filter_shape=(2,), num_filters=16, reduction_rank=1, pad=True, init=kernel_init)(a)

    assert b.shape == c.shape

    n = np.random.random((1, width) + image_shape).astype(np.float32)

    desired = b.eval({a: n})
    actual = c.eval({a: n})

    if isinstance(desired, list) and len(desired) == 1:
        desired = desired[0]

    if isinstance(actual, list) and len(actual) == 1:
        actual = actual[0]

    np.testing.assert_equal(actual, desired)

    # --------------------------------------------------------------------------
    # Normal use case - B&W image
    # --------------------------------------------------------------------------
    kernel_init = C.glorot_normal(seed=5)
    image_shape = (32, )
    width = 40
    a = C.sequence.input_variable(image_shape)
    b = C.layers.SequentialConvolution(filter_shape=(3, 3), num_filters=16, reduction_rank=0, pad=True, init=kernel_init)(a)
    c = Cx.layers.SequentialConvolution(filter_shape=(3, 3), num_filters=16, reduction_rank=0, pad=True, init=kernel_init)(a)

    assert b.shape == c.shape

    n = np.random.random((1, width) + image_shape).astype(np.float32)

    desired = b.eval({a: n})
    actual = c.eval({a: n})

    if isinstance(desired, list) and len(desired) == 1:
        desired = desired[0]

    if isinstance(actual, list) and len(actual) == 1:
        actual = actual[0]

    np.testing.assert_equal(actual, desired)

    image_shape = (32,)
    width = 40
    a = C.sequence.input_variable(image_shape)
    b = C.layers.SequentialConvolution(filter_shape=(3, 3), num_filters=16, reduction_rank=0, pad=True)(a)
    c = Cx.layers.SequentialConvolution(filter_shape=(3, 3), num_filters=16, reduction_rank=0, pad=True, init=b.W.value)(a)

    assert b.shape == c.shape

    n = np.random.random((1, width) + image_shape).astype(np.float32)

    desired = b.eval({a: n})
    actual = c.eval({a: n})

    if isinstance(desired, list) and len(desired) == 1:
        desired = desired[0]

    if isinstance(actual, list) and len(actual) == 1:
        actual = actual[0]

    np.testing.assert_equal(actual, desired)
