# Source code for nlp_architect.data.cdc_resources.relations.referent_dict_relation_extraction

    # ******************************************************************************
    # Copyright 2017-2018 Intel Corporation
    #
    # 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 logging
    import os
    from typing import Dict, List, Set

    from nlp_architect.common.cdc.mention_data import MentionDataLight
    from nlp_architect.data.cdc_resources.relations.relation_extraction import RelationExtraction
    from nlp_architect.data.cdc_resources.relations.relation_types_enums import (
        OnlineOROfflineMethod,
        RelationType,
    )
    from nlp_architect.utils.io import load_json_file
    from nlp_architect.utils.string_utils import StringUtils

    logger = logging.getLogger(__name__)

    class ReferentDictRelationExtraction(RelationExtraction):
        def __init__(
            self, method: OnlineOROfflineMethod = OnlineOROfflineMethod.ONLINE, ref_dict: str = None
        ):
            """
            Extract Relation between two mentions according to Referent Dictionary knowledge

            Args:
                method (optional): OnlineOROfflineMethod.{ONLINE/OFFLINE} run against full referent
                    dictionary or a sub-set of (default = ONLINE)
                ref_dict (required): str Location of referent dictionary file to work with
            """
            logger.info("Loading ReferentDict module")
            if ref_dict is not None and os.path.isfile(ref_dict):
                if method == OnlineOROfflineMethod.OFFLINE:
                    self.ref_dict = load_json_file(ref_dict)
                elif method == OnlineOROfflineMethod.ONLINE:
                    self.ref_dict = self.load_reference_dict(ref_dict)
                logger.info("ReferentDict module lead successfully")
            else:
                raise FileNotFoundError("Referent Dict file not found or not in path:" + ref_dict)

            super(ReferentDictRelationExtraction, self).__init__()

        def extract_all_relations(
            self, mention_x: MentionDataLight, mention_y: MentionDataLight
        ) -> Set[RelationType]:
            ret_ = set()
            ret_.add(self.extract_sub_relations(mention_x, mention_y, RelationType.REFERENT_DICT))
            return ret_

        def extract_sub_relations(
            self, mention_x: MentionDataLight, mention_y: MentionDataLight, relation: RelationType
        ) -> RelationType:
            """
            Check if input mentions has the given relation between them

            Args:
                mention_x: MentionDataLight
                mention_y: MentionDataLight
                relation: RelationType

            Returns:
                RelationType: relation in case mentions has given relation or
                    RelationType.NO_RELATION_FOUND otherwise
            """
            if relation is not RelationType.REFERENT_DICT:
                return RelationType.NO_RELATION_FOUND

            mention_x_str = mention_x.tokens_str
            mention_y_str = mention_y.tokens_str
            if StringUtils.is_pronoun(mention_x_str.lower()) or StringUtils.is_pronoun(
                mention_y_str.lower()
            ):
                return RelationType.NO_RELATION_FOUND

            if self.is_referent_dict(mention_x, mention_y):
                return RelationType.REFERENT_DICT

            return RelationType.NO_RELATION_FOUND

        def is_referent_dict(self, mention_x: MentionDataLight, mention_y: MentionDataLight) -> bool:
            """
            Check if input mentions has referent dictionary relation between them

            Args:
                mention_x: MentionDataLight
                mention_y: MentionDataLight

            Returns:
                bool
            """
            match_result = False
            x_head = mention_x.mention_head
            y_head = mention_y.mention_head
            x_head_lemma = mention_x.mention_head_lemma
            y_head_lemma = mention_y.mention_head_lemma

            if (x_head in self.ref_dict and y_head in self.ref_dict[x_head]) or (
                y_head in self.ref_dict and x_head in self.ref_dict[y_head]
            ):
                match_result = True
            if (x_head_lemma in self.ref_dict and y_head_lemma in self.ref_dict[x_head_lemma]) or (
                y_head_lemma in self.ref_dict and x_head_lemma in self.ref_dict[y_head_lemma]
            ):
                match_result = True
            if (x_head_lemma in self.ref_dict and y_head in self.ref_dict[x_head_lemma]) or (
                y_head_lemma in self.ref_dict and x_head in self.ref_dict[y_head_lemma]
            ):
                match_result = True
            if (y_head in self.ref_dict and x_head_lemma in self.ref_dict[y_head]) or (
                x_head in self.ref_dict and y_head_lemma in self.ref_dict[x_head]
            ):
                match_result = True

            return match_result

        @staticmethod
        def get_supported_relations():
            """
            Return all supported relations by this class

            Returns:
                List[RelationType]
            """
            return [RelationType.REFERENT_DICT]

        @staticmethod
        def load_reference_dict(dict_fname: str) -> Dict[str, List[str]]:
            """
            Method to load referent dictionary to memory

            Returns:
                List[RelationType]
            """
            word_dict = {}
            first = True
            with open(dict_fname, "r", encoding="utf-8") as f:
                for line in f:
                    if first:
                        first = False
                        continue
                    word1, word2, _, npmi = line.strip().split("\t")
                    npmi = float(npmi)
                    if npmi >= 0.2:
                        if word1 not in word_dict:
                            word_dict[word1] = []
                        word_dict[word1].append(word2)
            return word_dict