import pytest

import theano
import theano.tensor as T
import numpy as np

from lasagne.layers import InputLayer, ReshapeLayer, DenseLayer
from lasagne.layers import get_output, get_all_param_values, set_all_param_values
from ntm.layers import NTMLayer
from ntm.heads import WriteHead, ReadHead
from ntm.controllers import DenseController
from ntm.memory import Memory


def model(input_var, batch_size=1):
    l_input = InputLayer((batch_size, None, 8), input_var=input_var)
    _, seqlen, _ = l_input.input_var.shape

    # Neural Turing Machine Layer
    memory = Memory((128, 20), name='memory')
    controller = DenseController(l_input, memory_shape=(128, 20),
        num_units=100, num_reads=1, name='controller')
    heads = [
        WriteHead(controller, num_shifts=3, memory_shape=(128, 20), name='write'),
        ReadHead(controller, num_shifts=3, memory_shape=(128, 20), name='read')
    ]
    l_ntm = NTMLayer(l_input, memory=memory, controller=controller, heads=heads)

    # Output Layer
    l_output_reshape = ReshapeLayer(l_ntm, (-1, 100))
    l_output_dense = DenseLayer(l_output_reshape, num_units=8, name='dense')
    l_output = ReshapeLayer(l_output_dense, (batch_size, seqlen, 8))

    return l_output


def test_batch_size():
    input_var01, input_var16 = T.tensor3s('input01', 'input16')
    l_output01 = model(input_var01, batch_size=1)
    l_output16 = model(input_var16, batch_size=16)

    # Share the parameters for both models
    params01 = get_all_param_values(l_output01)
    set_all_param_values(l_output16, params01)

    posterior_fn01 = theano.function([input_var01], get_output(l_output01))
    posterior_fn16 = theano.function([input_var16], get_output(l_output16))

    example_input = np.random.rand(16, 30, 8)
    example_output16 = posterior_fn16(example_input)
    example_output01 = np.zeros_like(example_output16)

    for i in range(16):
        example_output01[i] = posterior_fn01(example_input[i][np.newaxis, :, :])

    assert example_output16.shape == (16, 30, 8)
    assert np.allclose(example_output16, example_output01, atol=1e-3)

Add test for size of batches
Add test for size of batches
