# -*- coding: utf-8 -*-
"""
Defining a pipeline tasks that writes extracted relations into a Neo4j graph database.
"""

# EXT
import luigi
import luigi.format

# PROJECT
from bwg.decorators import time_function
from bwg.helpers import fast_copy
from bwg.neo4j_extensions import Neo4jTarget
from bwg.serializing import deserialize_line
from bwg.tasks.pipeline_run_info_generation import PipelineRunInfoGenerationTask
from bwg.tasks.relation_merging import RelationMergingTask
from bwg.tasks.properties_completion import PropertiesCompletionTask


class RelationsDatabaseWritingTask(luigi.Task):
    """
    Writes relations extracted via (naive) Open Relation Extraction and Participation Extraction into a graph database.
    """
    task_config = luigi.DictParameter()
    pipeline_run_info = None  # Information about current pipeline run, containing processed sentences, timestamp etc.

    def requires(self):
        return RelationMergingTask(task_config=self.task_config), \
               PropertiesCompletionTask(task_config=self.task_config), \
               PipelineRunInfoGenerationTask(task_config=self.task_config)

    def output(self):
        user = self.task_config["NEO4J_USER"]
        password = self.task_config["<PASSWORD>"]
        host = self.task_config["NEO4J_HOST"]
        categories = self.task_config["DATABASE_CATEGORIES"]

        return Neo4jTarget(
            self.pipeline_run_info, user, password,
            node_relevance_function=self.is_relevant_node,
            categorization_function=self.categorize_node,
            host=host,
            categories=categories
        )

    @time_function(is_classmethod=True)
    def run(self):
        with self.input()[0].open("r") as mr_file, self.input()[1].open("r") as pc_file,\
                self.input()[2].open("r") as pri_file:
            self._read_pipeline_run_info(pri_file)
            entity_properties = self._read_properties_file(pc_file)
            with self.output() as database:
                for mr_line in mr_file:
                    self.process_article(mr_line, database, entity_properties)

    @time_function(is_classmethod=True)
    def process_article(self, raw_article, database, entity_properties):
        """
        Process an article generated by previous tasks and process them.

        :param raw_article: Raw unserialized article.
        :type raw_article: str
        :param database: Neo4j luigi target.
        :type database: bwg.db.neo4j.Neo4jTarget
        :param entity_properties: Wikidata properties of all entities as dictionary.
        :type entity_properties: dict
        """
        debug = self.task_config.get("PIPELINE_DEBUG", False)
        encoding = self.task_config["CORPUS_ENCODING"]
        article = deserialize_line(raw_article, encoding)
        article_meta, article_data = article["meta"], article["data"]

        if debug:
            print("{} processing article '{}'...".format(self.__class__.__name__, article["meta"]["title"]))

        for sentence_id, sentence_json in article_data.items():
            if debug:
                print("{} finished sentence #{}.".format(self.__class__.__name__, sentence_id))

            for relation_id, relation_json in sentence_json["data"]["relations"].items():
                database.add_relation(relation_json, sentence_json["data"]["sentence"], entity_properties)

    def _read_properties_file(self, properties_file):
        """
        Read all Wikidata properties from properties file.

        :param properties_file: File with Wikidata properties.
        :type properties_file: luigi.Target.
        :return: Wikidata properties of all entities as dictionary.
        :rtype: dict
        """
        encoding = self.task_config["CORPUS_ENCODING"]
        entity_properties = {}

        for line in properties_file:
            article = deserialize_line(line, encoding)
            article_meta, article_data = article["meta"], article["data"]

            entity_properties.update(self._extract_properties(article_data))

        return entity_properties

    def _extract_properties(self, article_data):
        """
        Extract all Wikidata properties from an article's data.

        :param article_data: Article data.
        :type article_data: dict
        :return: All properties of entities in this article.
        :rtype: dict
        """
        entity_properties = {}

        for sentence_id, sentence_json in article_data.items():
            sentence_meta, sentence_data = sentence_json["meta"], sentence_json["data"]

            for entity_id, entity_json in sentence_data["entities"].items():
                entity_meta, entity_data = entity_json["meta"], entity_json["data"]

                new_entity_data = {
                    "ambiguous": entity_meta["type"] == "Ambiguous Wikidata entity",
                    "senses": [self._convert_entity_sense(entity_sense) for entity_sense in entity_data]
                }

                entity_properties = self._add_entity_data_to_properties_dict(new_entity_data, entity_properties)

        return entity_properties

    @staticmethod
    def _add_entity_data_to_properties_dict(entity_data, entity_properties):
        """
        Add an entity and its aliases to the properties dictionary.

        :param entity_data: Data of target entity.
        :type entity_data: dict
        :param entity_properties: Properties_dictionary that the entity and its aliases are added to.
        :type entity_properties: dict
        :return: Properties dictionary with new entry.
        :rtype: dict
        """
        for sense_data in entity_data["senses"]:
            if "label" not in sense_data:
                continue

            entity_properties[sense_data["label"]] = entity_data

            if "aliases" in sense_data:
                for alias in sense_data["aliases"]:
                    # Adjust label and aliases appropriately
                    alias_sense_data = fast_copy(sense_data)
                    alias_sense_data["label"] = alias
                    alias_sense_data["aliases"].remove(alias)
                    alias_sense_data["aliases"].append(sense_data["label"])
                    alias_entity_data = fast_copy(entity_data)
                    alias_entity_data["senses"][entity_data["senses"].index(sense_data)] = alias_sense_data
                    entity_properties[alias] = alias_entity_data

        return entity_properties

    def _convert_entity_sense(self, entity_data, recursively=True):
        """
        Rename some fields for clarity.

        :param entity_data: Data with fields to be renamed.
        :type entity_data: dict
        :param recursively: Process data of related nodes too.
        :type recursively: bool
        :return: Data with renamed fields.
        :rtype: dict
        """
        new_entity_data = dict(entity_data)
        new_entity_data = self._rename_field("modified", "wikidata_last_modified", new_entity_data)
        new_entity_data = self._rename_field("id", "wikidata_id", new_entity_data)

        # Apply transformation to related Wikidata notes
        if recursively:
            for claim_name, claim_data in new_entity_data.get("claims", {}).items():
                claim_data["target_data"] = {
                    "ambiguous": len(claim_data["target_data"]) > 1,
                    "senses": [
                        self._convert_entity_sense(entity_sense, recursively=False)
                        if claim_data["target_data"] != {} else []
                        for entity_sense in claim_data["target_data"]
                    ]
                }

        return new_entity_data

    @staticmethod
    def _rename_field(field, new_name, dictionary):
        """
        Rename a field in a dictionary.

        :param field: Field to be renamed.
        :type field: str
        :param new_name: New name of field.
        :type new_name: str
        :param dictionary: Dictionary in which the field occurs.
        :type dictionary: dict
        :return: Dictionary with renamed field.
        :rtype: dict
        """
        dictionary[new_name] = dictionary[field]
        del dictionary[field]
        return dictionary

    def _read_pipeline_run_info(self, pri_file):
        """
        Read the current pipeline run info.

        :param pri_file: File with pipeline run info.
        :type: Luigi.target.
        :return: Pipeline run info.
        :rtype: dict
        """
        encoding = self.task_config["CORPUS_ENCODING"]

        for line in pri_file:
            self.pipeline_run_info = deserialize_line(line, encoding)
            break

    def is_relevant_node(self, label, node_data):
        """
        Determine whether a node is relevant and should be written to the database. This function can be overwritten by
        tasks inheriting from this tasks.

        :param label: Node label
        :tyoe label: str
        :param node_data: Node's data.
        :type node_data: dict
        :return: Result of check.
        :rtype: bool
        """
        return True

    def categorize_node(self, label, node_data):
        """
        Assign a node a category out of a pre-defined set of categories. This function can be overwritten by tasks
        inheriting from this tasks.

        :param label: Node label
        :tyoe label: str
        :param node_data: Node's data.
        :type node_data: dict
        :return: Category for node.
        :rtype: str
        """
        if "senses" not in node_data:
            return "miscellaneous"

        # Pick the first sense for now
        sense = node_data["senses"][0]
        ne_tag = sense.get("type", "I-MISC")

        ne_tags_to_model = {
            "I-PER": "Person",
            "I-LOC": "Location",
            "I-ORG": "Organization",
            "DATE": "Date",
            "I-MISC": "Miscellaneous"
        }

        return "Entity" if ne_tag not in ne_tags_to_model else ne_tags_to_model[ne_tag]

# -*- coding: utf-8 -*-
"""
Defining a pipeline tasks that writes extracted relations into a Neo4j graph database.
"""

# EXT
import luigi
import luigi.format

# PROJECT
from bwg.decorators import time_function
from bwg.helpers import fast_copy
from bwg.neo4j_extensions import Neo4jTarget
from bwg.serializing import deserialize_line
from bwg.tasks.pipeline_run_info_generation import PipelineRunInfoGenerationTask
from bwg.tasks.relation_merging import RelationMergingTask
from bwg.tasks.properties_completion import PropertiesCompletionTask


class RelationsDatabaseWritingTask(luigi.Task):
    """
    Writes relations extracted via (naive) Open Relation Extraction and Participation Extraction into a graph database.
    """
    task_config = luigi.DictParameter()
    pipeline_run_info = None  # Information about current pipeline run, containing processed sentences, timestamp etc.

    def requires(self):
        return RelationMergingTask(task_config=self.task_config), \
               PropertiesCompletionTask(task_config=self.task_config), \
               PipelineRunInfoGenerationTask(task_config=self.task_config)

    def output(self):
        user = self.task_config["NEO4J_USER"]
        password = self.task_config["<PASSWORD>"]
        host = self.task_config["NEO4J_HOST"]
        categories = self.task_config["DATABASE_CATEGORIES"]

        return Neo4jTarget(
            self.pipeline_run_info, user, password,
            node_relevance_function=self.is_relevant_node,
            categorization_function=self.categorize_node,
            host=host,
            categories=categories
        )

    @time_function(is_classmethod=True)
    def run(self):
        with self.input()[0].open("r") as mr_file, self.input()[1].open("r") as pc_file,\
                self.input()[2].open("r") as pri_file:
            self._read_pipeline_run_info(pri_file)
            entity_properties = self._read_properties_file(pc_file)
            with self.output() as database:
                for mr_line in mr_file:
                    self.process_article(mr_line, database, entity_properties)

    @time_function(is_classmethod=True)
    def process_article(self, raw_article, database, entity_properties):
        """
        Process an article generated by previous tasks and process them.

        :param raw_article: Raw unserialized article.
        :type raw_article: str
        :param database: Neo4j luigi target.
        :type database: bwg.db.neo4j.Neo4jTarget
        :param entity_properties: Wikidata properties of all entities as dictionary.
        :type entity_properties: dict
        """
        debug = self.task_config.get("PIPELINE_DEBUG", False)
        encoding = self.task_config["CORPUS_ENCODING"]
        article = deserialize_line(raw_article, encoding)
        article_meta, article_data = article["meta"], article["data"]

        if debug:
            print("{} processing article '{}'...".format(self.__class__.__name__, article["meta"]["title"]))

        for sentence_id, sentence_json in article_data.items():
            if debug:
                print("{} finished sentence #{}.".format(self.__class__.__name__, sentence_id))

            for relation_id, relation_json in sentence_json["data"]["relations"].items():
                database.add_relation(relation_json, sentence_json["data"]["sentence"], entity_properties)

    def _read_properties_file(self, properties_file):
        """
        Read all Wikidata properties from properties file.

        :param properties_file: File with Wikidata properties.
        :type properties_file: luigi.Target.
        :return: Wikidata properties of all entities as dictionary.
        :rtype: dict
        """
        encoding = self.task_config["CORPUS_ENCODING"]
        entity_properties = {}

        for line in properties_file:
            article = deserialize_line(line, encoding)
            article_meta, article_data = article["meta"], article["data"]

            entity_properties.update(self._extract_properties(article_data))

        return entity_properties

    def _extract_properties(self, article_data):
        """
        Extract all Wikidata properties from an article's data.

        :param article_data: Article data.
        :type article_data: dict
        :return: All properties of entities in this article.
        :rtype: dict
        """
        entity_properties = {}

        for sentence_id, sentence_json in article_data.items():
            sentence_meta, sentence_data = sentence_json["meta"], sentence_json["data"]

            for entity_id, entity_json in sentence_data["entities"].items():
                entity_meta, entity_data = entity_json["meta"], entity_json["data"]

                new_entity_data = {
                    "ambiguous": entity_meta["type"] == "Ambiguous Wikidata entity",
                    "senses": [self._convert_entity_sense(entity_sense) for entity_sense in entity_data]
                }

                entity_properties = self._add_entity_data_to_properties_dict(new_entity_data, entity_properties)

        return entity_properties

    @staticmethod
    def _add_entity_data_to_properties_dict(entity_data, entity_properties):
        """
        Add an entity and its aliases to the properties dictionary.

        :param entity_data: Data of target entity.
        :type entity_data: dict
        :param entity_properties: Properties_dictionary that the entity and its aliases are added to.
        :type entity_properties: dict
        :return: Properties dictionary with new entry.
        :rtype: dict
        """
        for sense_data in entity_data["senses"]:
            if "label" not in sense_data:
                continue

            entity_properties[sense_data["label"]] = entity_data

            if "aliases" in sense_data:
                for alias in sense_data["aliases"]:
                    # Adjust label and aliases appropriately
                    alias_sense_data = fast_copy(sense_data)
                    alias_sense_data["label"] = alias
                    alias_sense_data["aliases"].remove(alias)
                    alias_sense_data["aliases"].append(sense_data["label"])
                    alias_entity_data = fast_copy(entity_data)
                    alias_entity_data["senses"][entity_data["senses"].index(sense_data)] = alias_sense_data
                    entity_properties[alias] = alias_entity_data

        return entity_properties

    def _convert_entity_sense(self, entity_data, recursively=True):
        """
        Rename some fields for clarity.

        :param entity_data: Data with fields to be renamed.
        :type entity_data: dict
        :param recursively: Process data of related nodes too.
        :type recursively: bool
        :return: Data with renamed fields.
        :rtype: dict
        """
        new_entity_data = dict(entity_data)
        new_entity_data = self._rename_field("modified", "wikidata_last_modified", new_entity_data)
        new_entity_data = self._rename_field("id", "wikidata_id", new_entity_data)

        # Apply transformation to related Wikidata notes
        if recursively:
            for claim_name, claim_data in new_entity_data.get("claims", {}).items():
                claim_data["target_data"] = {
                    "ambiguous": len(claim_data["target_data"]) > 1,
                    "senses": [
                        self._convert_entity_sense(entity_sense, recursively=False)
                        if claim_data["target_data"] != {} else []
                        for entity_sense in claim_data["target_data"]
                    ]
                }

        return new_entity_data

    @staticmethod
    def _rename_field(field, new_name, dictionary):
        """
        Rename a field in a dictionary.

        :param field: Field to be renamed.
        :type field: str
        :param new_name: New name of field.
        :type new_name: str
        :param dictionary: Dictionary in which the field occurs.
        :type dictionary: dict
        :return: Dictionary with renamed field.
        :rtype: dict
        """
        dictionary[new_name] = dictionary[field]
        del dictionary[field]
        return dictionary

    def _read_pipeline_run_info(self, pri_file):
        """
        Read the current pipeline run info.

        :param pri_file: File with pipeline run info.
        :type: Luigi.target.
        :return: Pipeline run info.
        :rtype: dict
        """
        encoding = self.task_config["CORPUS_ENCODING"]

        for line in pri_file:
            self.pipeline_run_info = deserialize_line(line, encoding)
            break

    def is_relevant_node(self, label, node_data):
        """
        Determine whether a node is relevant and should be written to the database. This function can be overwritten by
        tasks inheriting from this tasks.

        :param label: Node label
        :tyoe label: str
        :param node_data: Node's data.
        :type node_data: dict
        :return: Result of check.
        :rtype: bool
        """
        return True

    def categorize_node(self, label, node_data):
        """
        Assign a node a category out of a pre-defined set of categories. This function can be overwritten by tasks
        inheriting from this tasks.

        :param label: Node label
        :tyoe label: str
        :param node_data: Node's data.
        :type node_data: dict
        :return: Category for node.
        :rtype: str
        """
        if "senses" not in node_data:
            return "miscellaneous"

        # Pick the first sense for now
        sense = node_data["senses"][0]
        ne_tag = sense.get("type", "I-MISC")

        ne_tags_to_model = {
            "I-PER": "Person",
            "I-LOC": "Location",
            "I-ORG": "Organization",
            "DATE": "Date",
            "I-MISC": "Miscellaneous"
        }

        return "Entity" if ne_tag not in ne_tags_to_model else ne_tags_to_model[ne_tag]

# -*- coding: utf-8 -*- Defining a pipeline tasks that writes extracted relations into a Neo4j graph database. # EXT # PROJECT Writes relations extracted via (naive) Open Relation Extraction and Participation Extraction into a graph database. # Information about current pipeline run, containing processed sentences, timestamp etc. Process an article generated by previous tasks and process them.          :param raw_article: Raw unserialized article.         :type raw_article: str         :param database: Neo4j luigi target.         :type database: bwg.db.neo4j.Neo4jTarget         :param entity_properties: Wikidata properties of all entities as dictionary.         :type entity_properties: dict #{}.".format(self.__class__.__name__, sentence_id)) Read all Wikidata properties from properties file.          :param properties_file: File with Wikidata properties.         :type properties_file: luigi.Target.         :return: Wikidata properties of all entities as dictionary.         :rtype: dict Extract all Wikidata properties from an article's data.          :param article_data: Article data.         :type article_data: dict         :return: All properties of entities in this article.         :rtype: dict Add an entity and its aliases to the properties dictionary.          :param entity_data: Data of target entity.         :type entity_data: dict         :param entity_properties: Properties_dictionary that the entity and its aliases are added to.         :type entity_properties: dict         :return: Properties dictionary with new entry.         :rtype: dict # Adjust label and aliases appropriately Rename some fields for clarity.          :param entity_data: Data with fields to be renamed.         :type entity_data: dict         :param recursively: Process data of related nodes too.         :type recursively: bool         :return: Data with renamed fields.         :rtype: dict # Apply transformation to related Wikidata notes Rename a field in a dictionary.          :param field: Field to be renamed.         :type field: str         :param new_name: New name of field.         :type new_name: str         :param dictionary: Dictionary in which the field occurs.         :type dictionary: dict         :return: Dictionary with renamed field.         :rtype: dict Read the current pipeline run info.          :param pri_file: File with pipeline run info.         :type: Luigi.target.         :return: Pipeline run info.         :rtype: dict Determine whether a node is relevant and should be written to the database. This function can be overwritten by         tasks inheriting from this tasks.          :param label: Node label         :tyoe label: str         :param node_data: Node's data.         :type node_data: dict         :return: Result of check.         :rtype: bool Assign a node a category out of a pre-defined set of categories. This function can be overwritten by tasks         inheriting from this tasks.          :param label: Node label         :tyoe label: str         :param node_data: Node's data.         :type node_data: dict         :return: Category for node.         :rtype: str # Pick the first sense for now