# Source code for nlp_architect.data.cdc_resources.relations.word_embedding_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 math
    from typing import List, Set

    from scipy.spatial.distance import cosine as cos

    from nlp_architect.common.cdc.mention_data import MentionDataLight
    from nlp_architect.data.cdc_resources.embedding.embed_elmo import (
        ElmoEmbedding,
        ElmoEmbeddingOffline,
    )
    from nlp_architect.data.cdc_resources.embedding.embed_glove import (
        GloveEmbedding,
        GloveEmbeddingOffline,
    )
    from nlp_architect.data.cdc_resources.relations.relation_extraction import RelationExtraction
    from nlp_architect.data.cdc_resources.relations.relation_types_enums import (
        EmbeddingMethod,
        RelationType,
    )
    from nlp_architect.utils.string_utils import StringUtils

    logger = logging.getLogger(__name__)

    class WordEmbeddingRelationExtraction(RelationExtraction):
        def __init__(
            self,
            method: EmbeddingMethod = EmbeddingMethod.GLOVE,
            glove_file: str = None,
            elmo_file: str = None,
            cos_accepted_dist: float = 0.7,
        ):
            """
            Extract Relation between two mentions according to Word Embedding cosine distance

            Args:
                method (optional): EmbeddingMethod.{GLOVE/GLOVE_OFFLINE/ELMO/ELMO_OFFLINE}
                    (default = GLOVE)
                glove_file (required on GLOVE/GLOVE_OFFLINE mode): str Location of Glove file
                elmo_file (required on ELMO_OFFLINE mode): str Location of Elmo file
            """
            if method == EmbeddingMethod.GLOVE:
                self.embedding = GloveEmbedding(glove_file)
                self.contextual = False
            elif method == EmbeddingMethod.GLOVE_OFFLINE:
                self.embedding = GloveEmbeddingOffline(glove_file)
                self.contextual = False
            elif method == EmbeddingMethod.ELMO:
                self.embedding = ElmoEmbedding()
                self.contextual = True
            elif method == EmbeddingMethod.ELMO_OFFLINE:
                self.embedding = ElmoEmbeddingOffline(elmo_file)
                self.contextual = True

            self.accepted_dist = cos_accepted_dist
            super(WordEmbeddingRelationExtraction, 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.WORD_EMBEDDING_MATCH)
            )
            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.WORD_EMBEDDING_MATCH:
                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()
            ):
                if not self.contextual:
                    return RelationType.NO_RELATION_FOUND

                if mention_x.mention_context is None or mention_y.mention_context is None:
                    return RelationType.NO_RELATION_FOUND

            if self.is_word_embed_match(mention_x, mention_y):
                return RelationType.WORD_EMBEDDING_MATCH

            return RelationType.NO_RELATION_FOUND

        def is_word_embed_match(self, mention_x: MentionDataLight, mention_y: MentionDataLight):
            """
            Check if input mentions Word Embedding cosine distance below above 0.65

            Args:
                mention_x: MentionDataLight
                mention_y: MentionDataLight

            Returns:
                bool
            """
            match_result = False
            x_embed = self.embedding.get_head_feature_vector(mention_x)
            y_embed = self.embedding.get_head_feature_vector(mention_y)
            # make sure words are not 'unk/None/0'
            if x_embed is not None and y_embed is not None:
                dist = cos(x_embed, y_embed)
                if not math.isnan(dist):
                    sim = 1 - dist
                    if sim >= self.accepted_dist:
                        match_result = True

            return match_result

        @staticmethod
        def get_supported_relations() -> List[RelationType]:
            """
            Return all supported relations by this class

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