import jieba
import jieba.posseg as pseg
import re
import datetime


# 从输入的“公司名”中提取主体
def main_extract(input_str,stop_word,d_4_delete,d_city_province):
    # 开始分词并处理
    seg = pseg.cut(input_str)
    seg = remove_word(seg,stop_word,d_4_delete)
    seg_lst = city_prov_ahead(seg,d_city_province)
    return seg_lst

    
#TODO：实现公司名称中地名提前
def city_prov_ahead(seg,d_city_province):
    # print(seg)
    city_prov_lst = []
    seg_lst = []
    # TODO ...
    for s in seg:
        # print(s)
        if s.flag == 'ns' and s.word in d_city_province:
            city_prov_lst.append(s.word)
        else:
            seg_lst.append(s.word)
    return city_prov_lst+seg_lst




#TODO：替换特殊符号
def remove_word(seg,stop_word,d_4_delete):
    # TODO ...
    seg_lst = [j for j in [i for i in seg if i.word not in stop_word] if j.word not in d_4_delete]
    return seg_lst


# 初始化，加载词典
def my_initial():
    fr1 = open(r"../data/dict/co_City_Dim.txt", encoding='utf-8')
    fr2 = open(r"../data/dict/co_Province_Dim.txt", encoding='utf-8')
    fr3 = open(r"../data/dict/company_business_scope.txt", encoding='utf-8')
    fr4 = open(r"../data/dict/company_suffix.txt", encoding='utf-8')
    #城市名
    lines1 = fr1.readlines()
    d_4_delete = []
    d_city_province = [re.sub(r'(\r|\n)*','',line) for line in lines1]
    #省份名
    lines2 = fr2.readlines()
    l2_tmp = [re.sub(r'(\r|\n)*','',line) for line in lines2]
    d_city_province.extend(l2_tmp)
    #公司后缀
    lines3 = fr3.readlines()
    l3_tmp = [re.sub(r'(\r|\n)*','',line) for line in lines3]
    lines4 = fr4.readlines()
    l4_tmp = [re.sub(r'(\r|\n)*','',line) for line in lines4]
    d_4_delete.extend(l4_tmp)
    #get stop_word
    fr = open(r'../data/dict/stopwords.txt', encoding='utf-8')   
    stop_word = fr.readlines()
    stop_word_after = [re.sub(r'(\r|\n)*','',stop_word[i]) for i in range(len(stop_word))]
    stop_word_after[-1] = stop_word[-1]
    stop_word = stop_word_after
    return d_4_delete,stop_word,d_city_province

# TODO：测试实体统一用例
d_4_delete,stop_word,d_city_province = my_initial()
company_name = "河北银行股份有限公司"
lst = main_extract(company_name,stop_word,d_4_delete,d_city_province)
company_name = ''.join(lst)  # 对公司名提取主体部分，将包含相同主体部分的公司统一为一个实体
print(company_name)

input_str = '金固股份'
lst = main_extract(input_str,stop_word,d_4_delete,d_city_province)
company_name = ''.join(lst)  # 对公司名提取主体部分，将包含相同主体部分的公司统一为一个实体
print(company_name)

import pandas as pd
train_data = pd.read_csv('../data/info_extract/train_data.csv', encoding = 'gb2312', header=0)
train_data.head()

test_data = pd.read_csv('../data/info_extract/test_data.csv', encoding = 'gb2312', header=0)
test_data.head()

train_data.shape, test_data.shape

# 处理test数据，利用开源工具进行实体识别和并使用实体统一函数存储实体

import fool
import pandas as pd
from copy import copy


test_data = pd.read_csv('../data/info_extract/test_data.csv', encoding = 'gb2312', header=0)
test_data['ner'] = None
ner_id = 1001
ner_dict_new = {}  # 存储所有实体
ner_dict_reverse_new = {}  # 存储所有实体

for i in range(len(test_data)):
    sentence = copy(test_data.iloc[i, 1])
    # TODO：调用fool进行实体识别，得到words和ners结果
    # TODO ...
    # print(sentence)
    words, ners = fool.analysis(sentence)
    
    ners[0].sort(key=lambda x:x[0], reverse=True)
    for start, end, ner_type, ner_name in ners[0]:
        if ner_type=='company' or ner_type=='person':
            
            company_main_name = ner_name    
            if ner_type=='company':
                # TODO：调用实体统一函数，存储统一后的实体
                # 并自增ner_id
                # TODO ...
                lst = main_extract(ner_name,stop_word,d_4_delete,d_city_province)
                company_main_name = ''.join(lst)  # 对公司名提取主体部分，将包含相同主体部分的公司统一为一个实体

            ner_dict_new[company_main_name] = ner_id
            ner_dict_reverse_new[ner_id] = company_main_name
            ner_id += 1            
            
            # 在句子中用编号替换实体名
            sentence = sentence[:start] + ' ner_' + str(ner_dict_new[company_main_name]) + '_ ' + sentence[end:]
    test_data.iloc[i, -1] = sentence

X_test = test_data[['ner']]

# 处理train数据，利用开源工具进行实体识别和并使用实体统一函数存储实体
train_data = pd.read_csv('../data/info_extract/train_data.csv', encoding = 'gb2312', header=0)
train_data['ner'] = None

for i in range(len(train_data)):
    # 判断正负样本
    if train_data.iloc[i,:]['member1']=='0' and train_data.iloc[i,:]['member2']=='0':
        sentence = copy(train_data.iloc[i, 1])
        # TODO：调用fool进行实体识别，得到words和ners结果
        # TODO ...
        words, ners = fool.analysis(sentence)
    
        ners[0].sort(key=lambda x:x[0], reverse=True)
        for start, end, ner_type, ner_name in ners[0]:
            if ner_type=='company' or ner_type=='person':
                # TODO：调用实体统一函数，存储统一后的实体
                # 并自增ner_id
                # TODO ...
                company_main_name = ner_name    
                if ner_type=='company':
                    # TODO：调用实体统一函数，存储统一后的实体
                    # 并自增ner_id
                    # TODO ...
                    lst = main_extract(ner_name,stop_word,d_4_delete,d_city_province)
                    company_main_name = ''.join(lst)  # 对公司名提取主体部分，将包含相同主体部分的公司统一为一个实体

                ner_dict_new[company_main_name] = ner_id
                ner_dict_reverse_new[ner_id] = company_main_name
                ner_id += 1   

                # 在句子中用编号替换实体名
                sentence = sentence[:start] + ' ner_' + str(ner_dict_new[company_main_name]) + '_ ' + sentence[end-1:]
        train_data.iloc[i, -1] = sentence
    else:
        # 将训练集中正样本已经标注的实体也使用编码替换
        sentence = copy(train_data.iloc[i,:]['sentence'])
        for company_main_name in [train_data.iloc[i,:]['member1'],train_data.iloc[i,:]['member2']]:
            # TODO：调用实体统一函数，存储统一后的实体
            # 并自增ner_id
            # TODO ...

            company_main_name = ner_name    
            if ner_type=='company':
                # TODO：调用实体统一函数，存储统一后的实体
                # 并自增ner_id
                # TODO ...
                lst = main_extract(ner_name,stop_word,d_4_delete,d_city_province)
                company_main_name = ''.join(lst)  # 对公司名提取主体部分，将包含相同主体部分的公司统一为一个实体

            ner_dict_new[company_main_name] = ner_id
            ner_dict_reverse_new[ner_id] = company_main_name
            ner_id += 1   

            # 在句子中用编号替换实体名
            sentence = re.sub(company_main_name, ' ner_%s_ '%(str(ner_dict_new[company_main_name])), sentence)
        train_data.iloc[i, -1] = sentence
        
y = train_data.loc[:,['tag']]
train_num = len(train_data)
X_train = train_data[['ner']]

# 将train和test放在一起提取特征
X = pd.concat([X_train, X_test])

# code
from sklearn.feature_extraction.text import TfidfTransformer  
from sklearn.feature_extraction.text import CountVectorizer  
from pyltp import Segmentor


# 实体符号加入分词词典
with open('../data/user_dict.txt', 'w') as fw:
    for v in ['一', '二', '三', '四', '五', '六', '七', '八', '九', '十']:
        fw.write( v + '号企业 ni\n')

# 初始化实例
segmentor = Segmentor()  
# 加载模型，加载自定义词典
segmentor.load_with_lexicon('/Users/zbh/data/ltp_data_v3.4.0/cws.model', '../data/user_dict.txt')  

# 加载停用词
fr = open(r'../data/dict/stopwords.txt', encoding='utf-8')   
stop_word = fr.readlines()
stop_word = [re.sub(r'(\r|\n)*','',stop_word[i]) for i in range(len(stop_word))]

# 分词
f = lambda x: ' '.join([word for word in segmentor.segment(x) if word not in stop_word and not re.findall(r'ner\_\d\d\d\d\_', word)])
corpus=X['ner'].map(f).tolist()

corpus[30:32]

from sklearn.feature_extraction.text import TfidfVectorizer
# TODO：提取tfidf特征
# TODO ...
tv = TfidfVectorizer(use_idf=True, smooth_idf=True, norm=None)
tv_fit = tv.fit_transform(corpus)

# -*- coding: utf-8 -*-
from pyltp import Parser
from pyltp import Segmentor
from pyltp import Postagger
import networkx as nx
import pylab
import re

postagger = Postagger() # 初始化实例
postagger.load_with_lexicon('/Users/zbh/data/ltp_data_v3.4.0/pos.model', '../data/user_dict.txt')  # 加载模型
segmentor = Segmentor()  # 初始化实例
segmentor.load_with_lexicon('/Users/zbh/data/ltp_data_v3.4.0/cws.model', '../data/user_dict.txt')  # 加载模型

def parse(s):
    """
    对语句进行句法分析，并返回句法结果
    """
    tmp_ner_dict = {}
    num_lst = ['一', '二', '三', '四', '五', '六', '七', '八', '九', '十']

    # 将公司代码替换为特殊称谓，保证分词词性正确
    for i, ner in enumerate(list(set(re.findall(r'(ner\_\d\d\d\d\_)', s)))):
        try:
            tmp_ner_dict[num_lst[i]+'号企业'] = ner
        except IndexError:
            # TODO：定义错误情况的输出
            # TODO ...
            
            
        s = s.replace(ner, num_lst[i]+'号企业')
    words = segmentor.segment(s)
    tags = postagger.postag(words)
    parser = Parser() # 初始化实例
    parser.load('/Users/zbh/data/ltp_data_v3.4.0/parser.model')  # 加载模型
    arcs = parser.parse(words, tags)  # 句法分析
    arcs_lst = list(map(list, zip(*[[arc.head, arc.relation] for arc in arcs])))
    
    # 句法分析结果输出
    parse_result = pd.DataFrame([[a,b,c,d] for a,b,c,d in zip(list(words),list(tags), arcs_lst[0], arcs_lst[1])], index = range(1,len(words)+1))
    parser.release()  # 释放模型
    # TODO：提取企业实体依存句法类型
    # TODO ...
    
    

    # 投资关系关键词
    key_words = ["收购","竞拍","转让","扩张","并购","注资","整合","并入","竞购","竞买","支付","收购价","收购价格","承购","购得","购进",
             "购入","买进","买入","赎买","购销","议购","函购","函售","抛售","售卖","销售","转售"]
    # TODO：*根据关键词和对应句法关系提取特征（如没有思路可以不完成）
    # TODO ...
    
    
    parser.release()  # 释放模型
    return your_result

tmp_ner_dict = {}
num_lst = ['一', '二', '三', '四', '五', '六', '七', '八', '九', '十']

s = X['ner'][:1][0]

s

# 将公司代码替换为特殊称谓，保证分词词性正确
for i, ner in enumerate(list(set(re.findall(r'(ner\_\d\d\d\d\_)', s)))):
    print(i, ner)
    try:
        tmp_ner_dict[num_lst[i]+'号企业'] = ner
    except IndexError:
        # TODO：定义错误情况的输出
        # TODO ...
        print('ner num is out off range')
        pass
        
    s = s.replace(ner, num_lst[i]+'号企业')

words = segmentor.segment(s)
tags = postagger.postag(words)
parser = Parser() # 初始化实例
parser.load('/Users/zbh/data/ltp_data_v3.4.0/parser.model')  # 加载模型
arcs = parser.parse(words, tags)  # 句法分析
arcs_lst = list(map(list, zip(*[[arc.head, arc.relation] for arc in arcs])))

# arcs_lst

# 句法分析结果输出
parse_result = pd.DataFrame([[a,b,c,d] for a,b,c,d in zip(list(words),list(tags), arcs_lst[0], arcs_lst[1])], index = range(1,len(words)+1))
parser.release()  # 释放模型
# TODO：提取企业实体依存句法类型
# TODO ...

parse_result

# 投资关系关键词
key_words = ["收购","竞拍","转让","扩张","并购","注资","整合","并入","竞购","竞买","支付","收购价","收购价格","承购","购得","购进",
         "购入","买进","买入","赎买","购销","议购","函购","函售","抛售","售卖","销售","转售"]
# TODO：*根据关键词和对应句法关系提取特征（如没有思路可以不完成）
# TODO ...


parser.release()  # 释放模型
return your_result

def shortest_path(arcs_ret, source, target):
    """
    求出两个词最短依存句法路径，不存在路径返回-1
    arcs_ret：句法分析结果
    source：实体1
    target：实体2
    """
    G=nx.DiGraph()
    # 为这个网络添加节点...
    for i in list(arcs_ret.index):
        G.add_node(i)
    # TODO：在网络中添加带权中的边...（注意，我们需要的是无向边）
    # TODO ...
    

    try:
        # TODO：利用nx包中shortest_path_length方法实现最短距离提取
        # TODO ...
        
        
        return distance
    except:
        return -1

def get_feature(s):
    """
    汇总上述函数汇总句法分析特征与TFIDF特征
    """
    # TODO：汇总上述函数汇总句法分析特征与TFIDF特征
    # TODO ...
    
    
    return features

# 建立分类器进行分类
from sklearn.ensemble import RandomForestClassifier
from sklearn import preprocessing
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import GridSearchCV

# TODO：定义需要遍历的参数


# TODO：选择模型


# TODO：利用GridSearchCV搜索最佳参数


# TODO：对Test_data进行分类



# TODO：保存Test_data分类结果
# 答案提交在submit目录中，命名为info_extract_submit.csv和info_extract_entity.csv。
# info_extract_entity.csv格式为：第一列是实体编号，第二列是实体名（实体统一的多个实体名用“|”分隔）
# info_extract_submit.csv格式为：第一列是关系中实体1的编号，第二列为关系中实体2的编号。

 

from py2neo import Node, Relationship, Graph

graph = Graph(
    "http://localhost:7474", 
    username="neo4j", 
    password="person"
)

for v in relation_list:
    a = Node('Company', name=v[0])
    b = Node('Company', name=v[1])
    
    # 本次不区分投资方和被投资方，无向图
    r = Relationship(a, 'INVEST', b)
    s = a | b | r
    graph.create(s)
    r = Relationship(b, 'INVEST', a)
    s = a | b | r
    graph.create(s)

# TODO：查询某节点的3层投资关系

import jieba
import pandas as pd

# 找出test_data.csv中前25条样本所有的人物名称，以及人物所在文档的上下文内容
test_data = pd.read_csv('../data/info_extract/test_data.csv', encoding = 'gb2312', header=0)

# 存储人物以及上下文信息（key为人物ID，value为人物名称、人物上下文内容）
person_name = {}

# 观察上下文的窗口大小
window = 10  

# 遍历前25条样本
for i in range(25):
    sentence = copy(test_data.iloc[i, 1])
    words, ners = fool.analysis(sentence)
    ners[0].sort(key=lambda x:x[0], reverse=True)
    for start, end, ner_type, ner_name in ners[0]:
        if ner_type=='person':
            # TODO：提取实体的上下文
            



# 利用爬虫得到每个人物名称对应的URL
# TODO：找到每个人物实体的词条内容。

# TODO：将样本中人物上下文与爬取词条结果进行对比，选择最接近的词条。



# 输出结果
pd.DataFrame(result_data).to_csv('../submit/entity_disambiguation_submit.csv', index=False)