import networkx as nx
import matplotlib.pyplot as plt

# 创建知识图谱
G = nx.Graph()

# 添加角色、地点、事件节点
characters = [
    "三藏", "孙悟空", "猪八戒", "沙僧", "镇元子", "妖精"
]
locations = [
    "白虎岭", "南山", "蛇回兽怕的白虎岭", "山岩", "南山山顶", "云端", "山坡"
]
events = [
    "三藏师徒上路", "镇元子与行者结为兄弟", "长老吃草还丹", "行者化缘",
    "遇到妖精变成的女子", "悟空识破妖精", "八戒误信妖精", "悟空打败妖精",
    "唐僧念紧箍咒"
]

# 添加节点
G.add_nodes_from(characters, type="角色")
G.add_nodes_from(locations, type="地点")
G.add_nodes_from(events, type="事件")

# 添加关系边
edges = [
    ("三藏", "孙悟空", "师徒关系"),
    ("三藏", "猪八戒", "师徒关系"),
    ("三藏", "沙僧", "师徒关系"),
    ("镇元子", "孙悟空", "兄弟"),
    ("妖精", "孙悟空", "敌对"),
    ("妖精", "三藏", "欲吃肉"),
    ("白虎岭", "妖精", "活动地点"),
    ("南山", "孙悟空", "觅桃"),
    ("三藏", "白虎岭", "遇到妖精"),
    ("悟空识破妖精", "孙悟空", "主导"),
    ("三藏", "唐僧念紧箍咒", "触发"),
    ("孙悟空", "唐僧念紧箍咒", "受影响"),
    ("猪八戒", "八戒误信妖精", "主导"),
    ("妖精", "八戒误信妖精", "欺骗")
]

# 将边添加到图中
for u, v, relation in edges:
    G.add_edge(u, v, label=relation)

# 绘制图形
pos = nx.spring_layout(G, seed=42)  # 节点布局
plt.figure(figsize=(15, 10))

# 绘制节点和边
nx.draw_networkx_nodes(G, pos, node_color="skyblue", node_size=2000)
nx.draw_networkx_edges(G, pos, edgelist=G.edges(), arrows=True, arrowstyle="->", width=2)

# 绘制节点标签
nx.draw_networkx_labels(G, pos, font_size=10, font_family="sans-serif")

# 绘制边标签
edge_labels = {(u, v): d["label"] for u, v, d in G.edges(data=True)}
nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_color="red")

plt.title("《西游记》知识图谱 - 三藏师徒与妖精之战")
plt.axis("off")
plt.show()