# pylint: disable=no-self-use,invalid-name
from numpy.testing import assert_almost_equal
import pytest
import torch
from torch.autograd import Variable

from allennlp.common import Params
from allennlp.common.checks import ConfigurationError
from allennlp.modules import Maxout
from allennlp.nn import InitializerApplicator
from allennlp.common.testing import AllenNlpTestCase


class TestMaxout(AllenNlpTestCase):
    def test_init_checks_output_dims_consistency(self):
        with pytest.raises(ConfigurationError):
            Maxout(input_dim=2,
                   num_layers=2,
                   output_dims=[5, 4, 3],
                   pool_sizes=4,
                   dropout=0.0)

    def test_init_checks_pool_sizes_consistency(self):
        with pytest.raises(ConfigurationError):
            Maxout(input_dim=2,
                   num_layers=2,
                   output_dims=5,
                   pool_sizes=[4, 5, 2],
                   dropout=0.0)

    def test_init_checks_dropout_consistency(self):
        with pytest.raises(ConfigurationError):
            Maxout(input_dim=2,
                   num_layers=3,
                   output_dims=5,
                   pool_sizes=4,
                   dropout=[0.2, 0.3])

    def test_forward_gives_correct_output(self):
        params = Params({
                'input_dim': 2,
                'output_dims': 3,
                'pool_sizes': 4,
                'dropout': 0.0,
                'num_layers': 2
                })
        maxout = Maxout.from_params(params)

        constant_init = lambda tensor: torch.nn.init.constant(tensor, 1.)
        initializer = InitializerApplicator([(".*", constant_init)])
        initializer(maxout)

        input_tensor = Variable(torch.FloatTensor([[-3, 1]]))
        output = maxout(input_tensor).data.numpy()
        assert output.shape == (1, 3)
        # This output was checked by hand
        # The output of the first maxout layer is [-1, -1, -1], since the
        # matrix multiply gives us [-2]*12. Reshaping and maxing
        # produces [-2, -2, -2] and the bias increments these values.
        # The second layer output is [-2, -2, -2], since the matrix
        # matrix multiply gives us [-3]*12. Reshaping and maxing
        # produces [-3, -3, -3] and the bias increments these values.
        assert_almost_equal(output, [[-2, -2, -2]])
