# -*- coding: utf-8 -*-

# ======================= 导入需要使用的Python库 ====================================

# 一些含有常用函数的库
import codecs
from collections import defaultdict
from pandas import DataFrame
import pandas as pd
# 结巴分词的库
import jieba
# 生成图像的库
import matplotlib.pyplot as plt
# 生成词云图的库
import wordcloud
# 生成柱状图的库
import numpy as np
# 生成人际关系图的库
import networkx as nx

# ======================= 路径设置 ====================================

# 存放需要处理的文本的目录，使用txt格式
TEXT_PATH = './文章/文本.txt'

# 各类词典的路径
DICT_PATH = './词典/人名.txt'  # 用于保存所有的人物名称
DICT_PATH_NZ = './词典/专有名词.txt'  # 用于保存非人名的专有名词
DICT_PATH_RELE = './词典/人物关系.txt'  # 用于保存人物之间的关系
DICT_PATH_NO = './词典/停用词.txt'  # 用于保存停用词
SYNONYMOUS_DICT_PATH = './词典/同义词.txt'  # 同义词的词库

# 输出成果的路径
SAVE_NODE_PATH = './结果/节点.txt' # 保存登场人物和登场次数
SAVE_EDGE_PATH = './结果/关联.txt' # 保存人物之间的关系权重
SAVE_CLOUD_PATH = './结果/词云图.png' # 生成文本中所用的词汇的频率词云图
SAVE_TIMES_PATH = './结果/人物登场次数图.png' # 生成人物登场次数的柱状图
SAVE_RELET_PATH = './结果/人物关系图.png' # 生成文本中登场的人物

# ======================= 词典创建 ====================================

# 让结巴分词加载自定义词典
jieba.load_userdict(DICT_PATH)
jieba.load_userdict(DICT_PATH_NZ)

# 创建同义词表，每行是一系列同义词，用空格分割，第一个词为导出词
COMBINE_LIST = {}
for line in open(SYNONYMOUS_DICT_PATH, "r", encoding='utf-8'):
    seperate_word = line.strip().split(" ")
    num = len(seperate_word)
    for i in range(1, num):
        COMBINE_LIST[seperate_word[i]] = seperate_word[0]

# 创建停用词列表
STOP_LIST = [line.strip() for line in open(DICT_PATH_NO, encoding='UTF-8').readlines()]

# 创建人名和关系列表
NAME_LIST = {} # 储存人物登场次数
RELATIONSHIP_LIST = {} # 储存人物关系
NAME_LINE = [] # 按照段落储存每段内的人物关系
with open(DICT_PATH, "r", encoding="utf8") as f:
    # 将角色姓名存入列表nameList
    nameList = f.read().replace("\n","").split(' 10 nr')

# ======================= 文本处理 ====================================

# 同义词和停用词处理
def text_editor(word):
    # 同义词处理
    if word in COMBINE_LIST:
        return COMBINE_LIST[word]
    # 停用词处理 and 去除单字
    if word not in STOP_LIST and (len(word) > 1 or word in nameList):
        return word
    return ""

# 名字处理
def text_name(word):
    # 当分词不在姓名列表nameList时认为该词不是人名
    if word not in nameList:
        return
    # 处理名字
    NAME_LINE[-1].append(word)             # 为当前段的环境增加一个人物
    if NAME_LIST.get(word) is None:            # 如果该人名在姓名字典中对应的权值为空（还没有这个键值对）
        NAME_LIST[word] = 0                # 则创建该键值对，参考实例test1.py
        RELATIONSHIP_LIST[word] = {}
    NAME_LIST[word] += 1                       # 该人物出现次数加 1

# 创建词云列表
CLOUD_LIST = ""

# 打开文章文本
with codecs.open(TEXT_PATH, "r", "utf8") as f:
    # 按行处理文本
    for line in f.readlines():
        # 为新读入的一段添加该段的人物名称列表
        NAME_LINE.append([])
        # 对当行进行分词
        for word in jieba.lcut(line):
            # 返回同义词替换和去掉停用词后的句子
            word = text_editor(word)
            # 若词语处理后为空则略过
            if word == "" or word == " ":
                continue
            # 名字处理
            text_name(word)
            # 词云增加
            CLOUD_LIST += word + " "

# 生成人物关系列表
# 按行处理
for line in NAME_LINE:
    # 获取每段中的任意两个人
    for name1 in line:                    
        for name2 in line:
            if name1 == name2:
                continue
            if RELATIONSHIP_LIST[name1].get(name2) is None:
                # 若两人尚未同时出现则新建项
                RELATIONSHIP_LIST[name1][name2] = 1
            else:
                # 两人共同出现次数加 1
                RELATIONSHIP_LIST[name1][name2] = RELATIONSHIP_LIST[name1][name2] + 1

# 存储人物节点文件（node.txt）
with codecs.open(SAVE_NODE_PATH, "w", "utf8") as f:
    f.write("Id,Label,Weight\r\n")
    for name, times in NAME_LIST.items():
        f.write(name + "," + name + "," + str(times) + "\r\n")

# 存储人物关系文件（edge.txt）
with codecs.open(SAVE_EDGE_PATH, "w", "utf-8") as f:
    f.write("Source,Target,Weight\r\n")
    for name, edges in RELATIONSHIP_LIST.items():
        for v, w in edges.items():
            f.write(name + "," + v + "," + str(w) + "\r\n")

# ======================= 可视化生成 ====================================

# 避免中文乱码
plt.rcParams['font.sans-serif'] = ['SimHei']

# 词云图生成

# 参数都可以注释掉，但必须设置font_path
wc = wordcloud.WordCloud(
    width=1000,
    height=800,
    background_color="#ffffff",  # 设置背景颜色
    max_words=500,  # 词的最大数（默认为200）
    max_font_size=300,  # 最大字体尺寸
    min_font_size=10,  # 最小字体尺寸（默认为4）
    colormap='bone',  # string or matplotlib colormap, default="viridis"
    random_state=20,  # 设置有多少种随机生成状态，即有多少种配色方案
#    mask=plt.imread("C:/1.bmp"),  # 读取遮罩图片！！
    font_path='C:/Windows/Fonts/simhei.ttf'
)

my_wordcloud = wc.generate(CLOUD_LIST)

plt.imshow(my_wordcloud)
plt.axis("off")

# 保存图片文件
wc.to_file(SAVE_CLOUD_PATH)

print("词云图生成完毕")


# 人物登场频次柱状图生成

# 人物按照出现频次从高到低排序
NAME_LIST = sorted(NAME_LIST.items(), key=lambda d: d[1], reverse=True)
# 获取X轴和Y轴的数据
x = []
y = []
for name in NAME_LIST:
    x.append(name[0])
    y.append(name[1])

fig, ax = plt.subplots(figsize=(15, 5))
ax.bar(
    x=x,  # Matplotlib自动将非数值变量转化为x轴坐标
    height=y,  # 柱子高度，y轴坐标
    width=0.6,  # 柱子宽度，默认0.8，两根柱子中心的距离默认为1.0
    align="center",  # 柱子的对齐方式，'center' or 'edge'
    color="grey",  # 柱子颜色
    edgecolor="red",  # 柱子边框的颜色
    linewidth=2.0  # 柱子边框线的大小
)
ax.set_title("人物登场次数柱状图", fontsize=15)

# 一个常见的场景是：每根柱子上方添加数值标签
# 步骤：
# 1. 准备要添加的标签和坐标
# 2. 调用ax.annotate()将文本添加到图表
# 3. 调整样式，例如标签大小，颜色和对齐方式
xticks = ax.get_xticks()
for i in range(len(y)):
    xy = (xticks[i], y[i] * 1.03)
    s = str(y[i])
    ax.annotate(
        s=s,  # 要添加的文本
        xy=xy,  # 将文本添加到哪个位置
        fontsize=10,  # 标签大小
        color="blue",  # 标签颜色
        ha="center",  # 水平对齐
        va="baseline"  # 垂直对齐
    )

# 坐标轴名称
plt.xlabel("人物")
plt.ylabel("登场次数")

# 保存图片文件
plt.savefig(SAVE_TIMES_PATH)

print("物种柱状图生成完毕")


# 关系图生成

# 生成画布
plt.figure(figsize=(10, 7))
# G：图表，一个networkx图
# networkx提供四种图形结构：
# G = nx.Graph() 无多重边无向图
# G = nx.DiGraph() 无多重边有向图
# G = nx.MultiGraph() 有多重边无向图
# G = nx.MultiDiGraph() 有多重边有向图
# 一般只用第一种无向图就好，如果使用有向图，还需要考虑每一条边的方向
G = nx.Graph()

# 读取人物节点的原始数据
GET_Names = []
with open(SAVE_NODE_PATH, "r", encoding="utf8") as f:
    # 按行处理文本
    for line in f.readlines():
        GET_Names.append(line.replace("\n",""))
# 读取人物关系的原始数据
GET_Relationships = []
with open(DICT_PATH_RELE, "r", encoding="utf8") as f:
    # 按行处理文本
    for line in f.readlines():
        GET_Relationships.append(line.replace("\n",""))
# 添加人物名称（节点）
i = 0
for name in GET_Names:
    # 跳过第一个表头
    if i == 0:
        i += 1
        continue
    G.add_node(name.split(",")[0], weight=int(name.split(",")[2]))
# 添加人物关系（边）
i = 0
with open(SAVE_EDGE_PATH, "r", encoding="utf8") as f:
    # 按行处理文本
    for line in f.readlines():
        # 跳过第一个表头
        if i == 0:
            i += 1
            continue
        # 初始化关系属性
        relation_name = ""
        # 获取相应的人物关系属性
        for relation in GET_Relationships:
            if (relation.split(",")[0] == line.split(",")[0] and relation.split(",")[1] == line.split(",")[1]) or (relation.split(",")[1] == line.split(",")[0] and relation.split(",")[0] == line.split(",")[1]):
                relation_name = relation.split(",")[2]
        # 添加边
        G.add_edge(line.split(",")[0], line.split(",")[1], weight=int(line.split(",")[2]), relationship=relation_name)

# 图的布局
# 建立布局，对图进行布局美化，networkx 提供的布局方式有：
#- circular_layout：节点在一个圆环上均匀分布
#- random_layout：节点随机分布
#- shell_layout：节点在同心圆上分布
#- spring_layout： 用Fruchterman-Reingold算法排列节点（样子类似多中心放射状）
#- spectral_layout：根据图的拉普拉斯特征向量排列节
# 这里使用Kamada Kawai布局的最优距离参数，并将非连接组件之间的距离设置为图中的最大距离
df = pd.DataFrame(index=G.nodes(), columns=G.nodes())
for row, data in nx.shortest_path_length(G):
    for col, dist in data.items():
        df.loc[row,col] = dist
df = df.fillna(df.max().max())
pos = nx.kamada_kawai_layout(G, dist=df.to_dict())

# 点
# 获取点的尺寸
nodeSize_0 = []
nodeSize = []
nodeColor = []
for weights in nx.get_node_attributes(G, "weight").values():
    # 开方以让数据更均匀一些
    nodeSize_0.append(int(weights) ** 0.5)
# 设置点位最细20px，最粗500px
min_weight = 100
max_weight = 300
# 斜率
k = (max_weight-min_weight)/(max(nodeSize_0)-min(nodeSize_0))
# 计算线条的粗细和颜色
for weights in nodeSize_0:
    nodeSize.append(1+k*(weights-min(nodeSize_0)))
    # 权重超过一半则为蓝色
    if weights > 0.5 * max(nodeSize_0):
        nodeColor.append('dodgerblue')
    else:
        nodeColor.append('lightgreen')
# 绘制节点
nx.draw_networkx_nodes(G, pos, alpha=1, node_size=nodeSize,node_shape='o',node_color=nodeColor)

# 边
# 获取边的权重列表
edgeWidth_0 = []
edgeWidth = []
edgeColor = []
for weights in nx.get_edge_attributes(G,'weight').values():
    # 开方以让数据更均匀一些
    edgeWidth_0.append(int(weights) ** 0.5)
# 设置线条最细1px，最粗10px
min_weight = 1
max_weight = 10
# 斜率
k = (max_weight-min_weight)/(max(edgeWidth_0)-min(edgeWidth_0))
# 计算线条的粗细和颜色
for weights in edgeWidth_0:
    edgeWidth.append(1+k*(weights-min(edgeWidth_0)))
    # 权重超过一半则为红色
    if weights > 0.5 * max(edgeWidth_0):
        edgeColor.append('lightcoral')
    else:
        edgeColor.append('lightgrey')
#pos:字典类型，节点作为键、位置作为值。位置是长度为2的序列
#edgelist：边缘元组的集合，只绘制指定的边，默认值为G.edges()
#width边的宽度，默认值为1.0
#alpha透明度，默认值为1.0（不透明），0为完全透明
#edge_color边的颜色，默认值为黑色
#style边的样式，默认值为实线。
nx.draw_networkx_edges(G, pos, width=edgeWidth, edge_color=edgeColor)

# 标签
# 线条上标注人物关系属性
labels = nx.get_edge_attributes(G,'relationship')
nx.draw_networkx_edge_labels(G,pos,edge_labels=labels, font_color='grey', font_size=8)
#font_size节点标签字体大小，默认值为12
nx.draw_networkx_labels(G, pos)

# 生成结果
plt.axis('off')
plt.title('人物关系图')
plt.rcParams['font.size'] = 10

# 保存图片文件
plt.savefig(SAVE_RELET_PATH)

print("人物关系图生成完毕")


# 显示图片
plt.show()