from neo4j import GraphDatabase
from enum import Enum

class RelationshipType(Enum):
    PREREQUISITE = "PREREQUISITE"  # 先修知识
    PART_OF = "PART_OF"            # 属于/包含
    IS_A = "IS_A"                  # 子类
    RELATED_TO = "RELATED_TO"      # 相关
    TESTS = "TESTS"                # 考察
    REQUIRES_SKILL = "REQUIRES_SKILL"  # 需要能力
    ASSESSES_SKILL = "ASSESSES_SKILL"  # 评估能力
    MASTERY = "MASTERY"            # 掌握程度

class KnowledgeGraphBuilder:
    def __init__(self, uri, user, password):
        self._driver = GraphDatabase.driver(uri, auth=(user, password))
    
    def close(self):
        self._driver.close()
    
    def clear_database(self):
        with self._driver.session() as session:
            session.run("MATCH (n) DETACH DELETE n")
    
    def create_concept_node(self, name, description, node_type="Concept"):
        with self._driver.session() as session:
            session.run(
                "MERGE (n:{} {{name: $name}}) "
                "SET n.description = $description".format(node_type),
                name=name, description=description
            )
    
    def create_relationship(self, from_node, to_node, rel_type, properties=None):
        with self._driver.session() as session:
            query = (
                "MATCH (a {{name: $from_name}}), (b {{name: $to_name}}) "
                "MERGE (a)-[r:{}]->(b)".format(rel_type.value)
            )
            if properties:
                query += " SET r += $properties"
            session.run(query, 
                       from_name=from_node, 
                       to_name=to_node,
                       properties=properties or {})
    
    def build_math_knowledge_graph(self):
        # 清空数据库
        self.clear_database()
        
        # 创建核心概念节点
        self.create_concept_node("数的认识", "数的基本概念和分类")
        self.create_concept_node("数的运算", "基本数学运算原理")
        
        # 创建具体知识点节点
        self.create_concept_node("自然数", "表示物体个数的数")
        self.create_concept_node("加法意义", "加法的基本定义和含义")
        self.create_concept_node("加法运算", "加法计算方法")
        
        # 创建技能节点
        self.create_concept_node("数值计算", "进行数学计算的能力", "Skill")
        
        # 创建关系
        self.create_relationship("自然数", "数的认识", RelationshipType.PART_OF)
        self.create_relationship("加法意义", "加法运算", RelationshipType.PREREQUISITE)
        self.create_relationship("加法运算", "数值计算", RelationshipType.REQUIRES_SKILL)
        
        # 示例：创建题目节点和关系
        self.create_concept_node("题目1", "计算3+5的值", "Question")
        self.create_relationship("题目1", "加法运算", RelationshipType.TESTS)
        self.create_relationship("题目1", "数值计算", RelationshipType.ASSESSES_SKILL)

if __name__ == "__main__":
    builder = KnowledgeGraphBuilder("bolt://localhost:7687", "neo4j", "DsideaL147258369")
    try:
        builder.build_math_knowledge_graph()
        print("知识图谱构建完成")
    finally:
        builder.close()