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

from allennlp.common import Params
from allennlp.common.testing import ModelTestCase
from allennlp.data import Vocabulary
from allennlp.data.iterators import EpochTrackingBucketIterator
from allennlp.models import Model
from allennlp.models.archival import load_archive


class NlvrCoverageSemanticParserTest(ModelTestCase):
    def setUp(self):
        super(NlvrCoverageSemanticParserTest, self).setUp()
        self.set_up_model("tests/fixtures/semantic_parsing/nlvr_coverage_semantic_parser/experiment.json",
                          "tests/fixtures/data/nlvr/sample_grouped_data.jsonl")

    def test_model_can_train_save_and_load(self):
        self.ensure_model_can_train_save_and_load(self.param_file)

    def test_ungrouped_model_can_train_save_and_load(self):
        # pylint: disable=line-too-long
        self.ensure_model_can_train_save_and_load("tests/fixtures/semantic_parsing/nlvr_coverage_semantic_parser/ungrouped_experiment.json")

    def test_mml_initialized_model_can_train_save_and_load(self):
        # pylint: disable=line-too-long
        self.ensure_model_can_train_save_and_load("tests/fixtures/semantic_parsing/nlvr_coverage_semantic_parser/mml_init_experiment.json")

    def test_get_checklist_info(self):
        # Creating a fake all_actions field where actions 0, 2 and 4 are terminal productions.
        all_actions = [('<o,o> -> top', True, None),
                       ('fake_action', True, None),
                       ('c -> color_black', True, None),
                       ('fake_action2', True, None),
                       ('e -> 6', True, None)]
        # Of the actions above, those at indices 0 and 4 are on the agenda, and there are padding
        # indices at the end.
        test_agenda = Variable(torch.Tensor([[0], [4], [-1], [-1]]))
        checklist_info = self.model._get_checklist_info(test_agenda, all_actions)
        target_checklist, terminal_actions, checklist_mask = checklist_info
        assert_almost_equal(target_checklist.data.numpy(), [[1], [0], [1]])
        assert_almost_equal(terminal_actions.data.numpy(), [[0], [2], [4]])
        assert_almost_equal(checklist_mask.data.numpy(), [[1], [1], [1]])

    def test_forward_with_epoch_num_changes_cost_weight(self):
        # Redefining model. We do not want this to change the state of ``self.model``.
        params = Params.from_file(self.param_file)
        model = Model.from_params(self.vocab, params['model'])
        # Initial cost weight, before forward is called.
        assert model._checklist_cost_weight == 0.8
        iterator = EpochTrackingBucketIterator(sorting_keys=[['sentence', 'num_tokens']])
        cost_weights = []
        for epoch_data in iterator(self.dataset, num_epochs=4):
            model.forward(**epoch_data)
            cost_weights.append(model._checklist_cost_weight)
        # The config file has ``wait_num_epochs`` set to 0, so the model starts decreasing the cost
        # weight at epoch 0 itself.
        assert_almost_equal(cost_weights, [0.72, 0.648, 0.5832, 0.52488])

    def test_initialize_weights_from_archive(self):
        original_model_parameters = self.model.named_parameters()
        original_model_weights = {name: parameter.data.clone().numpy()
                                  for name, parameter in original_model_parameters}
        # pylint: disable=line-too-long
        mml_model_archive_file = "tests/fixtures/semantic_parsing/nlvr_direct_semantic_parser/serialization/model.tar.gz"
        archive = load_archive(mml_model_archive_file)
        archived_model_parameters = archive.model.named_parameters()
        self.model._initialize_weights_from_archive(archive)
        changed_model_parameters = dict(self.model.named_parameters())
        for name, archived_parameter in archived_model_parameters:
            archived_weight = archived_parameter.data.numpy()
            original_weight = original_model_weights[name]
            changed_weight = changed_model_parameters[name].data.numpy()
            # We want to make sure that the weights in the original model have indeed been changed
            # after a call to ``_initialize_weights_from_archive``.
            with self.assertRaises(AssertionError, msg=f"{name} has not changed"):
                assert_almost_equal(original_weight, changed_weight)
            # This also includes the sentence token embedder. Those weights will be the same
            # because the two models have the same vocabulary.
            assert_almost_equal(archived_weight, changed_weight)

    def test_get_vocab_index_mapping(self):
        # pylint: disable=line-too-long
        mml_model_archive_file = "tests/fixtures/semantic_parsing/nlvr_direct_semantic_parser/serialization/model.tar.gz"
        archive = load_archive(mml_model_archive_file)
        mapping = self.model._get_vocab_index_mapping(archive.model.vocab)
        expected_mapping = [(i, i) for i in range(16)]
        assert mapping == expected_mapping

        new_vocab = Vocabulary()
        def copy_token_at_index(i):
            token = self.vocab.get_token_from_index(i, "tokens")
            new_vocab.add_token_to_namespace(token, "tokens")
        copy_token_at_index(5)
        copy_token_at_index(7)
        copy_token_at_index(10)
        mapping = self.model._get_vocab_index_mapping(new_vocab)
        # Mapping of indices from model vocabulary to new vocabulary. 0 and 1 are padding and unk
        # tokens.
        assert mapping == [(0, 0), (1, 1), (5, 2), (7, 3), (10, 4)]
