# pylint: disable=no-self-use,invalid-name
import numpy as np
from numpy.testing import assert_almost_equal
import torch
from allennlp.common import Params
from allennlp.data import Vocabulary
from allennlp.modules.token_embedders import BagOfWordCountsTokenEmbedder
from allennlp.common.testing import AllenNlpTestCase


class TestBagOfWordCountsTokenEmbedder(AllenNlpTestCase):
    def setUp(self):
        super(TestBagOfWordCountsTokenEmbedder, self).setUp()
        self.vocab = Vocabulary()
        self.vocab.add_token_to_namespace("1")
        self.vocab.add_token_to_namespace("2")
        self.vocab.add_token_to_namespace("3")
        self.vocab.add_token_to_namespace("4")

    def test_forward_calculates_bow_properly(self):
        params = Params({})
        embedder = BagOfWordCountsTokenEmbedder.from_params(self.vocab, params=params)
        numpy_tensor = np.array([[2, 0], [3, 0], [4, 4]])
        inputs = torch.from_numpy(numpy_tensor).unsqueeze(1)
        embedder_output = embedder(inputs)
        numpy_tensor = np.array([[1, 0, 1, 0, 0, 0], [1, 0, 0, 1, 0, 0], [0, 0, 0, 0, 2, 0]])
        manual_output = torch.from_numpy(numpy_tensor).float()
        assert_almost_equal(embedder_output.data.numpy(), manual_output.data.numpy())

    def test_projects_properly(self):
        params = Params({"projection_dim": 50})
        embedder = BagOfWordCountsTokenEmbedder.from_params(self.vocab, params=params)
        numpy_tensor = np.array([self.vocab.get_token_index(x) for x in ["1", "2", "3"]])
        inputs = torch.from_numpy(numpy_tensor).unsqueeze(1)
        embedder_output = embedder(inputs)
        assert embedder_output.shape[1] == 50

[((14, 21, 14, 33), 'allennlp.data.Vocabulary', 'Vocabulary', ({}, {}), '()', False, 'from allennlp.data import Vocabulary\n'), ((21, 17, 21, 27), 'allennlp.common.Params', 'Params', ({(21, 24, 21, 26): '{}'}, {}), '({})', False, 'from allennlp.common import Params\n'), ((22, 19, 22, 86), 'allennlp.modules.token_embedders.BagOfWordCountsTokenEmbedder.from_params', 'BagOfWordCountsTokenEmbedder.from_params', (), '', False, 'from allennlp.modules.token_embedders import BagOfWordCountsTokenEmbedder\n'), ((23, 23, 23, 57), 'numpy.array', 'np.array', ({(23, 32, 23, 56): '[[2, 0], [3, 0], [4, 4]]'}, {}), '([[2, 0], [3, 0], [4, 4]])', True, 'import numpy as np\n'), ((26, 23, 26, 93), 'numpy.array', 'np.array', ({(26, 32, 26, 92): '[[1, 0, 1, 0, 0, 0], [1, 0, 0, 1, 0, 0], [0, 0, 0, 0, 2, 0]]'}, {}), '([[1, 0, 1, 0, 0, 0], [1, 0, 0, 1, 0, 0], [0, 0, 0, 0, 2, 0]])', True, 'import numpy as np\n'), ((31, 17, 31, 47), 'allennlp.common.Params', 'Params', ({(31, 24, 31, 46): "{'projection_dim': 50}"}, {}), "({'projection_dim': 50})", False, 'from allennlp.common import Params\n'), ((32, 19, 32, 86), 'allennlp.modules.token_embedders.BagOfWordCountsTokenEmbedder.from_params', 'BagOfWordCountsTokenEmbedder.from_params', (), '', False, 'from allennlp.modules.token_embedders import BagOfWordCountsTokenEmbedder\n'), ((24, 17, 24, 47), 'torch.from_numpy', 'torch.from_numpy', ({(24, 34, 24, 46): 'numpy_tensor'}, {}), '(numpy_tensor)', False, 'import torch\n'), ((27, 24, 27, 54), 'torch.from_numpy', 'torch.from_numpy', ({(27, 41, 27, 53): 'numpy_tensor'}, {}), '(numpy_tensor)', False, 'import torch\n'), ((34, 17, 34, 47), 'torch.from_numpy', 'torch.from_numpy', ({(34, 34, 34, 46): 'numpy_tensor'}, {}), '(numpy_tensor)', False, 'import torch\n')]