# conda activate rag
# pip install py2neo pyvis
from py2neo import Graph
from pyvis.network import Network
from Config import Config

# 连接Neo4j
graph = Graph(Config.NEO4J_URI, auth=Config.NEO4J_AUTH)

# 查询所有知识点节点和关系
# query = """
# MATCH (n:KnowledgePoint)-[r]->(m)
# RETURN n, r, m
# """
query = """
MATCH (n:KnowledgePoint)
OPTIONAL MATCH (n)-[r1:PREREQUISITE]->(p)
OPTIONAL MATCH (n)-[r2:RELATED]->(r)
OPTIONAL MATCH (n)-[r3:HAS_SUB_POINT]->(s)
RETURN n, 
       COLLECT(DISTINCT {type: 'PREREQUISITE', node: p, rel: r1}) AS prerequisites,
       COLLECT(DISTINCT {type: 'RELATED', node: r, rel: r2}) AS relateds,
       COLLECT(DISTINCT {type: 'HAS_SUB_POINT', node: s, rel: r3}) AS sub_points
"""
data = graph.run(query).data()

# 创建网络图
net = Network(height="750px", width="100%", notebook=True, cdn_resources='in_line')

# 添加节点和边
# 修改后的添加节点和边逻辑
for item in data:
    node_label = item['n'].get('name', '未命名')
    net.add_node(item['n'].identity, 
                label=node_label,
                name=node_label,
                font={'size': 14, 'face': 'SimHei'})
    
    # 处理先修知识关系
    for pre in item['prerequisites']:
        if pre['node']:
            net.add_node(pre['node'].identity,
                        label=pre['node'].get('name', '未命名'),
                        name=pre['node'].get('name', '未命名'),
                        font={'size': 14, 'face': 'SimHei'})
            net.add_edge(item['n'].identity, pre['node'].identity, title='PREREQUISITE')
    
    # 处理相关知识关系
    for rel in item['relateds']:
        if rel['node']:
            net.add_node(rel['node'].identity,
                        label=rel['node'].get('name', '未命名'),
                        name=rel['node'].get('name', '未命名'),
                        font={'size': 14, 'face': 'SimHei'})
            net.add_edge(item['n'].identity, rel['node'].identity, title='RELATED')
    
    # 处理子知识点关系
    for sub in item['sub_points']:
        if sub['node']:
            net.add_node(sub['node'].identity,
                        label=sub['node'].get('name', '未命名'),
                        name=sub['node'].get('name', '未命名'),
                        font={'size': 14, 'face': 'SimHei'})
            net.add_edge(item['n'].identity, sub['node'].identity, title='HAS_SUB_POINT')

# 生成HTML文件
with open("knowledge_graph.html", "w", encoding='utf-8') as f:
    f.write(net.generate_html())
# 打开HTML文件
import webbrowser
webbrowser.open("knowledge_graph.html")