import numpy
import pytest
import torch

from allennlp.common import Params
from allennlp.common.checks import ConfigurationError
from allennlp.common.testing import AllenNlpTestCase
from allennlp.modules.seq2seq_encoders import IntraSentenceAttentionEncoder
from allennlp.modules.similarity_functions import MultiHeadedSimilarity


class TestIntraSentenceAttentionEncoder(AllenNlpTestCase):
    def test_get_dimension_is_correct(self):
        encoder = IntraSentenceAttentionEncoder(input_dim=5, projection_dim=4, combination="1,2")
        assert encoder.get_input_dim() == 5
        assert encoder.get_output_dim() == 9
        encoder = IntraSentenceAttentionEncoder(input_dim=5, combination="1+2")
        assert encoder.get_input_dim() == 5
        assert encoder.get_output_dim() == 5
        params = Params(
            {
                "input_dim": 8,
                "projection_dim": 8,
                "num_attention_heads": 4,
                "similarity_function": {"type": "multiheaded", "num_heads": 4, "tensor_1_dim": 8},
                "combination": "1,2,1*2",
            }
        )
        encoder = IntraSentenceAttentionEncoder.from_params(params)
        assert encoder.get_input_dim() == 8
        assert encoder.get_output_dim() == 24

    def test_constructor_asserts_multi_head_consistency(self):
        with pytest.raises(ConfigurationError) as exception_info:
            IntraSentenceAttentionEncoder(input_dim=5, num_attention_heads=4)
        assert "Encoder has multiple heads" in exception_info.value.message
        similarity = MultiHeadedSimilarity(3, 6)
        with pytest.raises(ConfigurationError) as exception_info:
            IntraSentenceAttentionEncoder(input_dim=5, similarity_function=similarity)
        assert "Similarity function has multiple heads" in exception_info.value.message
        with pytest.raises(ConfigurationError) as exception_info:
            IntraSentenceAttentionEncoder(
                input_dim=5, num_attention_heads=2, similarity_function=similarity
            )
        assert "Number of heads don't match" in exception_info.value.message

    def test_forward_works_with_simple_attention(self):
        # We're not going to check the output values here, as that's complicated; we'll just make
        # sure the code runs and the shapes are correct.
        encoder = IntraSentenceAttentionEncoder(input_dim=2)
        input_tensor = torch.from_numpy(numpy.random.rand(4, 3, 2))
        encoder_output = encoder(input_tensor, None)
        assert list(encoder_output.size()) == [4, 3, 4]  # default combination is 1,2

    def test_forward_works_with_multi_headed_attention(self):
        # We're not going to check the output values here, as that's complicated; we'll just make
        # sure the code runs and the shapes are correct.
        similarity = MultiHeadedSimilarity(3, 24)
        encoder = IntraSentenceAttentionEncoder(
            input_dim=24,
            projection_dim=24,
            similarity_function=similarity,
            num_attention_heads=3,
            combination="1+2",
        )
        input_tensor = torch.from_numpy(numpy.random.rand(4, 6, 24)).float()
        encoder_output = encoder(input_tensor, None)
        assert list(encoder_output.size()) == [4, 6, 24]
