# pylint: disable=no-self-use,invalid-name

import torch

from allennlp.common.testing import ModelTestCase


class CorefTest(ModelTestCase):
    def setUp(self):
        super(CorefTest, self).setUp()
        self.set_up_model(self.FIXTURES_ROOT / 'coref' / 'experiment.json',
                          self.FIXTURES_ROOT / 'coref' / 'coref.gold_conll')

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

    def test_decode(self):

        spans = torch.LongTensor([[1, 2],
                                  [3, 4],
                                  [3, 7],
                                  [5, 6],
                                  [14, 56],
                                  [17, 80]])

        antecedent_indices = torch.LongTensor([[0, 0, 0, 0, 0, 0],
                                               [0, 0, 0, 0, 0, 0],
                                               [1, 0, 0, 0, 0, 0],
                                               [2, 1, 0, 0, 0, 0],
                                               [3, 2, 1, 0, 0, 0],
                                               [4, 3, 2, 1, 0, 0]])

        spans = spans.unsqueeze(0)
        antecedent_indices = antecedent_indices
        # Indices into ``antecedent_indices`` indicating the predicted antecedent
        # index in ``top_spans``.
        predicted_antecedents = torch.LongTensor([-1, 0, -1, -1, 1, 3])
        predicted_antecedents = predicted_antecedents.unsqueeze(0)
        output_dict = {
                "top_spans": spans,
                "antecedent_indices": antecedent_indices,
                "predicted_antecedents": predicted_antecedents
                }
        output = self.model.decode(output_dict)

        clusters = output["clusters"][0]
        gold1 = [(1, 2), (3, 4), (17, 80)]
        gold2 = [(3, 7), (14, 56)]

        assert len(clusters) == 2
        assert gold1 in clusters
        assert gold2 in clusters
