<reponame>titanous/coreferee
# Copyright 2021 msg systems ag

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#   http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest
from coreferee.test_utils import get_nlps

class GermanSmokeTest(unittest.TestCase):

    def setUp(self):

        self.nlps = get_nlps('de')

    def all_nlps(self, func):
        for nlp in self.nlps:
            func(nlp)

    def compare_annotations(self, doc_text, expected_coref_chains, *, excluded_nlps=[],
        alternative_expected_coref_chains=None):

        def func(nlp):

            if nlp.meta['name'] in excluded_nlps:
                return

            doc = nlp(doc_text)
            chains_representation = str(doc._.coref_chains)
            if alternative_expected_coref_chains is None:
                self.assertEqual(expected_coref_chains,
                    chains_representation, nlp.meta['name'])
            else:
                self.assertTrue(expected_coref_chains == chains_representation or
                    alternative_expected_coref_chains == chains_representation)

        self.all_nlps(func)

    def test_simple(self):
        self.compare_annotations('Ich sah einen Hund und er jagte eine Katze', '[0: [3], [5]]')

    def test_simple_plural(self):
        self.compare_annotations('Ich sah Hunde, und sie jagten eine Katze', '[0: [2], [5]]')

    def test_simple_conjunction_same_word(self):
        self.compare_annotations(
            'Ich sah einen Hund und einen Hund, und sie jagten eine Katze', '[0: [3, 6], [9]]')

    def test_simple_conjunction_different_words(self):
        self.compare_annotations(
            'Ich sah einen Hund und ein Pferd, und sie jagten eine Katze', '[0: [3, 6], [9]]')

    def test_conjunction_different_pronouns(self):
        self.compare_annotations(
            'Ich sah Peter und Jana, und sie und er jagten ein Katze', '[0: [2], [9], 1: [4], [7]]')

    def test_conjunction_involving_pronoun(self):
        self.compare_annotations(
            'Ich sah Peter und Jana. Er und Richard jagten eine Katze', '[0: [2], [6]]')

    def test_different_sentence(self):
        self.compare_annotations(
            'Ich sah Peter. Er jagte eine Katze.', '[0: [2], [4]]')

    def test_proper_noun_coreference(self):
        self.compare_annotations(
            'Ich sah Peter. Peter jagte eine Katze.', '[0: [2], [4]]')

    def test_proper_noun_coreference_multiword(self):
        self.compare_annotations(
            'Ich sah <NAME>. <NAME> jagte eine Katze.', '[0: [3], [6]]')

    def test_proper_noun_coreference_multiword_only_second_repeated(self):
        self.compare_annotations(
            'Ich sah <NAME>. Paul jagte eine Katze.', '[0: [3], [5]]')

    def test_proper_noun_coreference_multiword_only_first_repeated(self):
        self.compare_annotations(
            'Ich sah <NAME>. Peter jagte eine Katze.', '[]')

    def test_proper_noun_coreference_with_gender_difference(self):
        self.compare_annotations(
            'Er arbeitete bei der BMW AG. Das Unternehmen gefiel ihm. Es hatte viel Erfolg auf dem deutschen Markt', '[0: [0], [10], 1: [4], [8], [12]]')

    def test_common_noun_coreference(self):
        self.compare_annotations(
            'Ich sah einen großen Hund. Der Hund jagte eine Katze. Er wedelte mit seinem Schwanz.',
            '[0: [4], [7], [12], [15]]')

    def test_entity_coreference(self):
        self.compare_annotations(
            'Es war Peter, der alles wusste. Alle mochten den freundlichen Mann.',
            '[0: [2], [12]]')

    def test_reflexive_simple(self):
        self.compare_annotations(
            'Der Gepard jagte sich',
            '[0: [1], [3]]')

    def test_reflexive_coordination(self):
        self.compare_annotations(
            'Der Gepard und <NAME> jagten sich',
            '[0: [1, 4], [6]]')

    def test_reflexive_excluded_mix_of_coordination_and_single_member_1(self):
        self.compare_annotations(
            'Peter und Jana kamen rein. Sie sahen ihn.',
            '[0: [0, 2], [6]]')

    def test_reflexive_excluded_mix_of_coordination_and_single_member_2(self):
        self.compare_annotations(
            'Peter und Jana kamen rein. Sie sahen sie.',
            '[0: [0, 2], [6]]')

    def test_reflexive_preceding(self):
        self.compare_annotations(
            'Er wollte, dass sich sein Sohn um die Sache kümmere',
            '[0: [0], [5], 1: [4], [6]]')

    def test_cataphora_simple(self):
        self.compare_annotations(
            'Obwohl er müde war, ging Peter nach Hause',
            '[0: [1], [6]]')

    def test_cataphora_with_coordination(self):
        self.compare_annotations(
            'Obwohl sie ihr Haus erreicht hatten, waren der Mann und die Frau traurig',
            '[0: [1], [12], 1: [2], [9, 12]]',
            alternative_expected_coref_chains='[0: [1], [2], [9, 12]]')

    def test_pronominal_anaphor(self):
        self.compare_annotations(
            'Sie fand ein Messer und aß damit einen Apfel.',
            '[0: [3], [6]]')

    def test_documentation_example(self):
        self.compare_annotations(
            'Weil er mit seiner Arbeit sehr beschäftigt war, hatte Peter genug davon. Er und seine Frau haben entschieden, dass ihnen ein Urlaub gut tun würde. Sie sind nach Spanien gefahren, weil ihnen das Land sehr gefiel.',
            '[0: [1], [3], [10], [14], [16], 1: [4], [12], 2: [14, 17], [22], [29], [36], 3: [32], [38]]', excluded_nlps=['core_news_sm']
        )
