"""
-----------------------------------------
@Author: xhj
@Email: 20212010075@fudan.edu.cn
@Created: 2022/4/13
------------------------------------------
@Modify: 2022/4/13
------------------------------------------
@Description:
"""
from pathlib import Path
import csv
import pickle
import math
import string
import random

from py2neo import Graph
from py2neo.cypher import Cursor
from py2neo.data import Node
import en_core_web_sm
from nltk.corpus import stopwords

from definitions import DATA_DIR
from sckg.constant_util import ConstantUtil


class SoftwareKGSearcher:
    def __init__(self, uri, user, password):
        self.constant = ConstantUtil()
        self.graph = Graph(uri, auth=(user, password))
        self.relation_weight = self.__construct_relation_weight()
        self.concepts = self.__construct_concepts()
        self.nlp = en_core_web_sm.load()
        self.stopwords = set(stopwords.words('english'))

    def dumps_entity2csv(self, label, properties: list, limit_size=None):
        # TODO: dump pagerank_score next
        cypher = f'"MATCH (n:{label}) RETURN labels(n) as labels, %s"' % ','.join([f'n.{_} as {_}' for _ in properties])
        if limit_size:
            cypher = f'{cypher[:-1]} LIMIT {limit_size}"'
        cypher = 'WITH ' + cypher + ' AS query\n'
        cypher += 'CALL apoc.export.csv.query(query, "%s.csv", {})\n' % label
        cypher += 'YIELD file, source, format, nodes, properties, time, rows, batchSize, batches, done\n'
        cypher += 'RETURN file, source, format, nodes, properties, time, rows, batchSize, batches, done;'
        self.graph.run(cypher)

    def dumps_entity2json(self, label, properties: list, limit_size=None):
        cypher = f'"MATCH (n:{label}) RETURN labels(n) as labels, %s"' % ','.join([f'n.{_} as {_}' for _ in properties])
        if limit_size:
            cypher = f'{cypher[:-1]} LIMIT {limit_size}"'
        cypher = 'WITH ' + cypher + ' AS query\n'
        cypher += 'CALL apoc.export.json.query(query, "%s.json", {})\n' % label
        cypher += 'YIELD file, nodes, relationships, properties, data\n'
        cypher += 'RETURN file, nodes, relationships, properties, data;'
        self.graph.run(cypher)

    @staticmethod
    def __construct_concepts():
        if (Path(DATA_DIR) / 'concept.bin').exists():
            with open(Path(DATA_DIR) / 'concept.bin', 'rb') as rf:
                return pickle.load(rf)
        concepts = {}
        with open(Path(DATA_DIR) / 'concept.csv', 'r', encoding='utf-8') as rf:
            reader = csv.reader(rf)
            next(reader)
            for row in reader:
                concepts[row[2]] = (row[1], row[0])
        with open(Path(DATA_DIR) / 'concept.bin', 'wb') as wf:
            pickle.dump(concepts, wf)
        return concepts

    @staticmethod
    def __construct_relation_weight():
        relation_weight = {}
        with open(Path(DATA_DIR) / 'new_relation_weight_10', 'r', encoding='utf-8') as rf:
            for line in rf:
                data_line = line.split(':')
                relation_weight[data_line[0]] = data_line[1]
        return relation_weight

    def get_node_num(self):
        return self.graph.run('MATCH (n) RETURN COUNT(*)').evaluate()

    def get_relation_num(self):
        return self.graph.run('MATCH P=()-->() RETURN COUNT(*)').evaluate()

    def get_node_info_by_id(self, concept_id):
        query = f'MATCH (n:entity{{id:{concept_id}}}) RETURN n.id as id, labels(n) as labels, properties(n) as properties'
        data = self.graph.run(query).data()
        if not data:
            return None
        return data[0]

    def get_concept_by_id(self, concept_id):
        query = f'MATCH (n:entity{{id:{concept_id}}}) RETURN n.concept_name'
        cursor: Cursor = self.graph.run(query)
        return cursor.evaluate()

    def get_node_by_concept(self, concept):
        query = f'MATCH (n:entity{{concept_name:{repr(concept)}}}) RETURN n.id as id, labels(n) as labels, properties(n) as properties'
        data = self.graph.run(query).data()
        if not data:
            return None
        return data[0]

    def get_id_by_concept(self, concept):
        query = f'MATCH (n:entity{{concept_name:{repr(concept)}}}) RETURN n.id'
        cursor: Cursor = self.graph.run(query)
        return cursor.evaluate()

    def get_id_by_concept_local(self, concept):
        if not self.is_exist_concept_local(concept):
            return None
        return self.concepts[concept][0]

    def is_exist_concept(self, concept):
        query = f'MATCH (n:entity{{concept_name:{repr(concept)}}}) RETURN COUNT(*)'
        return True if self.graph.run(query).evaluate() > 0 else False

    def is_exist_concept_local(self, concept):
        return self.concepts.__contains__(concept)

    def is_exist_concept_by_id(self, concept_id):
        query = f'MATCH (n:entity{{id:{concept_id}}}) RETURN COUNT(*)'
        return True if self.graph.run(query).evaluate() > 0 else False

    def get_concept_score(self, concept):
        query = f'MATCH (n:entity{{concept_name:{repr(concept)}}}) RETURN n.score'
        return self.graph.run(query).evaluate()

    def get_concept_pagerank_score(self, concept):
        query = f'MATCH (n:entity{{concept_name:{repr(concept)}}}) RETURN n.pagerank_score'
        return self.graph.run(query).evaluate()

    def get_concept_labels(self, concept):
        query = f'MATCH (n:concept{{concept_name:{repr(concept)}}}) RETURN labels(n)'
        return self.graph.run(query).evaluate()

    def get_concept_labels_local(self, concept):
        if not self.is_exist_concept_local(concept):
            return None
        return self.concepts[concept][1]

    def is_action(self, word):
        labels = self.get_concept_labels(word)
        if not labels:
            return False
        return True if 'action' in self.get_concept_labels(word) else False

    def is_action_local(self, word):
        labels = self.get_concept_labels_local(word)
        if not labels:
            return False
        return True if 'action' in labels else False

    def is_characteristic(self, word):
        labels = self.get_concept_labels(word)
        if not labels:
            return False
        return True if 'characteristic' in labels else False

    def is_characteristic_local(self, word):
        labels = self.get_concept_labels_local(word)
        if not labels:
            return False
        return True if 'characteristic' in labels else False

    def is_concept(self, word):
        labels = self.get_concept_labels(word)
        if not labels:
            return False
        return False if 'action' in labels or 'characteristic' in labels else True

    def is_concept_local(self, word):
        labels = self.get_concept_labels_local(word)
        if not labels:
            return False
        return False if 'action' in labels or 'characteristic' in labels else True

    def get_out_relation_by_concept(self, concept):
        query = f'MATCH (n:entity{{concept_name:{repr(concept)}}})-[r]->(x) RETURN n.concept_name, type(r), x.concept_name'
        cursor: Cursor = self.graph.run(query)
        return [tuple(record) for record in cursor]

    def get_out_relation_info_by_concept(self, concept):
        """
        返回结点的完整信息
        """
        query = f'MATCH (n:entity{{concept_name:{repr(concept)}}})-[r]->(x) RETURN n.concept_name, type(r), x'
        cursor: Cursor = self.graph.run(query)
        # return [(record[0], record[1],
        #          {
        #              'id': record[2]['id'],
        #              'labels': list(record[2].labels),
        #              'properties': dict(record[2])
        #          }) for record in cursor]
        result = []
        for record in cursor:
            node: Node = record[2]
            result.append((
                record[0],
                record[1],
                {
                    'id': node['id'],
                    'labels': list(node.labels),
                    'properties': dict(node)
                }
            ))
        return result

    def get_in_relation_by_concept(self, concept):
        query = f'MATCH (n:entity{{concept_name:{repr(concept)}}})<-[r]-(x) RETURN x.concept_name, type(r), n.concept_name'
        cursor: Cursor = self.graph.run(query)
        return [tuple(record) for record in cursor]

    def get_in_relation_info_by_concept(self, concept):
        """
        返回结点的完整信息
        """
        query = f'MATCH (n:entity{{concept_name:{repr(concept)}}})<-[r]-(x) RETURN x, type(r), n.concept_name'
        cursor: Cursor = self.graph.run(query)
        # return [({
        #              'id': record[0]['id'],
        #              'labels': list(record[0].labels),
        #              'properties': dict(record[0])
        #          }, record[1], record[2]) for record in cursor]
        result = []
        for record in cursor:
            node: Node = record[0]
            result.append((
                {
                    'id': node['id'],
                    'labels': list(node.labels),
                    'properties': dict(node)
                },
                record[1],
                record[2]
            ))
        return result

    def find_all_concept_from_sentence(self, sentence):
        """
        给一段文本，返回在KG里面出现的概念
        """
        concepts = []
        sentence = sentence.lower().translate(str.maketrans('', '', string.punctuation)).strip("\n")
        doc = self.nlp(sentence)
        words = []
        for token in doc:
            words.append(token.lemma_)
        # words = sentence.strip("\n").lower().split()
        for i, word in enumerate(words):
            concept = word
            if i == len(words) - 1:
                break
            for j in range(i + 1, i + 5):
                if j < len(words):
                    concept += " " + words[j]
                if self.is_exist_concept_local(concept):
                    concepts.append(concept)
        for word in words:
            if not self.is_concept_local(word):
                continue
            concepts.append(word)
        return set(concepts)

    def is_exit_concept_in_sentence(self, concept, sentence):
        return True if concept in self.find_all_concept_from_sentence(sentence) else False

    def find_longest_valid_concept_from_sentence(self, sentence):
        """
        给一段文本，返回在KG里面出现的有效最长概念
        """
        concepts = []
        sentence = sentence.lower().translate(str.maketrans('', '', string.punctuation)).strip("\n")
        doc = self.nlp(sentence)
        words = []
        for token in doc:
            words.append(token.lemma_)
        # words = sentence.strip("\n").lower().split()
        index = 0
        while index < len(words):
            concept = words[index]
            longest_concept = concept
            for j in range(index + 1, index + 6):
                if j < len(words):
                    concept += " " + words[j]
                    if self.is_concept_local(concept):
                        longest_concept = concept
                        index = j
                    else:
                        break
            if self.is_concept_local(longest_concept):
                concepts.append(longest_concept)
            index = index + 1
        return concepts

    def is_stop_word(self, word):
        return True if word in self.stopwords else False

    def t_find_longest_valid_concept_from_sentence(self, sentence):
        """
        给一段文本，返回在KG里面出现的有效最长概念
        """
        concepts = []
        sentence = sentence.lower().translate(str.maketrans('', '', string.punctuation)).strip("\n")
        doc = self.nlp(sentence)
        words = []
        for token in doc:
            words.append(token.lemma_)
        # words = sentence.strip("\n").lower().split()
        index = 0
        words = ['data' if i == 'datum' else i for i in words]
        while index < len(words):
            concept = words[index]
            longest_concept = concept
            for j in range(index + 1, index + 6):
                if j < len(words):
                    concept += " " + words[j]
                    if self.is_concept(concept):
                        longest_concept = concept
                        index = j
                    else:
                        break
            if self.is_concept(longest_concept):
                concepts.append(longest_concept)
            index = index + 1
        return concepts

    def find_longest_valid_concept_non_stop_words_from_sentence(self, sentence):
        """
        给一段文本，返回在KG里面出现的有效最长概念
        """
        concepts = []
        sentence = sentence.lower().translate(str.maketrans('', '', string.punctuation)).strip("\n")
        doc = self.nlp(sentence)
        words = []
        for token in doc:
            words.append(token.lemma_)
        index = 0
        while index < len(words):
            concept = words[index]
            longest_concept = concept
            for j in range(index + 1, index + 6):
                if j < len(words):
                    concept += ' ' + words[j]
                    if self.is_concept_local(concept):
                        longest_concept = concept
                        index = j
                    else:
                        break
            if self.is_concept_local(longest_concept) and longest_concept not in self.stopwords:
                concepts.append(longest_concept)
            index = index + 1
        return concepts

    def find_longest_concept_from_sentence(self, sentence):
        """
        给一段文本，返回在KG里面出现的最长概念
        """
        concepts = []
        sentence = sentence.lower().translate(str.maketrans('', '', string.punctuation)).strip("\n")
        doc = self.nlp(sentence)
        words = []
        for token in doc:
            words.append(token.lemma_)
        # words = sentence.strip("\n").lower().split()
        index = 0
        while index < len(words):
            concept = words[index]
            longest_concept = concept
            for j in range(index + 1, index + 6):
                if j < len(words):
                    concept += " " + words[j]
                    if self.is_exist_concept_local(concept):
                        longest_concept = concept
                        index = j
                    else:
                        break
            if self.is_concept_local(longest_concept):
                concepts.append((longest_concept, 'valid'))
            else:
                concepts.append((longest_concept, 'invalid'))
            index = index + 1
        return concepts

    def get_random_node(self):
        query = 'MATCH (n:entity) RETURN n.id as id, labels(n) as labels, properties(n) as properties LIMIT 50'
        cursor: Cursor = self.graph.run(query)
        return random.sample(cursor.data(), 5)

    def get_random_lastest_node(self):
        result = []
        max_num = self.get_node_num()
        count, i = 0, 0
        while count < 5:
            node_info = self.get_node_info_by_id(max_num - i)
            i += 1
            if not node_info:
                continue
            result.append(node_info)
            count += 1
        return result

    def get_hot_concepts(self):
        # result = sorted(self.concepts, key=lambda e: e['pagerank'], reverse=True)
        # return result[:10]
        # TODO cache 过期时间
        if (Path(DATA_DIR) / 'hot_concepts_cache.bin').exists():
            with open(Path(DATA_DIR) / 'hot_concepts_cache.bin', 'rb') as rf:
                return pickle.load(rf)
        score_query = 'MATCH (n:entity) RETURN max(n.pagerank_score)'
        max_pagerank_score = self.graph.run(score_query).evaluate()
        if not max_pagerank_score:
            return None
        max_pagerank_score -= 0.01
        query = f'MATCH (n:entity) WHERE n.pagerank_score>{max_pagerank_score} RETURN n.id as id, labels(n) as labels, properties(n) as properties'
        cursor: Cursor = self.graph.run(query)
        concepts = [{
            'id': record[0],
            'pagerank': record[2]['pagerank_score'],
            'labels': record[1],
            'properties': record[2]
        } for record in cursor]
        concepts = sorted(concepts, key=lambda x: x['pagerank'], reverse=True)[:10]
        with open(Path(DATA_DIR) / 'hot_concepts_cache.bin', 'wb') as wf:
            pickle.dump(concepts, wf)
        # for concept in concepts:
        #     print(concept[2]['pagerank_score'])
        # print(concepts)
        return concepts

    # def return_hot_concepts(self):
    #     """
    #     提前运行get_hot_concept保存结果
    #     """
    #     if (Path(DATA_DIR) / 'hot_concepts_cache.bin').exists():
    #         with open(Path(DATA_DIR) / 'hot_concepts_cache.bin', 'rb') as rf:
    #             return pickle.load(rf)
    #     return self.get_hot_concepts()

    def __cul_related_node(self, relations, orig_relations) -> list:
        nodes = {}
        try:
            for index, relation_tuple in enumerate(relations):
                node = relation_tuple[0]
                relation = relation_tuple[1]
                if relation not in self.constant.relation_score:
                    # TODO relation not in relation weight
                    relation_score = 0.6 * 0.6 + float(self.relation_weight[relation])
                else:
                    relation_score = self.constant.relation_score[relation] * 0.6 + float(
                        self.relation_weight[relation]) * 0.4
                pagerank_score = node['properties']['pagerank_score']
                pagerank_score = math.log10(pagerank_score * 10000000) / math.log10(10000000.000021683)
                # NOTE: 10000000.000169972(V3.1.46)  (3.0.-10000000.000014795)
                if 'frequency' not in node['properties']:
                    frequency = 1
                else:
                    frequency = node['properties']['sum_frequency']
                frequency_score = math.log10(frequency) / math.log10(10820730)
                # NOTE: 10877974 (V3.1.46) (3.0.-  10021884)
                node_score = pagerank_score * 0.7 + frequency_score * 0.3
                sum_score = relation_score * node_score
                nodes[sum_score] = orig_relations[index]
            nodes = sorted(nodes.items(), key=lambda x: x[0], reverse=True)
            return nodes
        except KeyError:
            return []

    def cul_in_related_node(self, concept):
        in_relations = self.get_in_relation_info_by_concept(concept)
        return self.__cul_related_node([(x, r) for x, r, n in in_relations], in_relations)

    def cul_out_related_node(self, concept):
        out_relations = self.get_out_relation_info_by_concept(concept)
        return self.__cul_related_node([(x, r) for n, r, x in out_relations], out_relations)

    def get_upper_concept(self, concept):
        node = self.get_node_by_concept(concept)
        if not node:
            return []
        relations = self.get_out_relation_info_by_concept(concept)
        upper_concepts = []
        for relation in relations:
            if relation[1] == 'is a' or relation[1] == 'instance of' or relation[1] == 'subclass of':
                upper_concepts.append(relation[2])
        return upper_concepts

    def get_common_upper_concept(self, concept1, concept2):
        node1_upper_concepts = self.get_upper_concept(concept1)
        node2_upper_concepts = self.get_upper_concept(concept2)
        upper_concepts = []
        for upper_concept in node1_upper_concepts:
            if upper_concept in node2_upper_concepts:
                upper_concepts.append(upper_concept)
        return upper_concepts
