knowledge_graph.py

"""
V4 知识图谱 - Line + Edge 结构
- Entity: 实体节点，包含演化 Line
- Edge: 实体间关系
- Note: 原始对话，被 Entity 和 Edge 引用
- LineNode: 演化线上的节点（原子笔记）
- RelationVersion: 关系演化版本节点（与 LineNode 统一设计）

V4 改进：
- 新增 get_entity_notes_with_evolution() 方法
- 新增 get_notes_with_evolution_by_ids() 方法
- 新增 RelationVersion 节点，统一关系演化追踪
"""

import json
from typing import List, Dict, Optional, Set
from neo4j import GraphDatabase
from config import NEO4J_CONFIG, MTM_CONFIG, EMBEDDING_CONFIG
from embedding_manager import EmbeddingManager


class KnowledgeGraphV2:
    """
    V2 知识图谱

    图结构:
    - (:Entity)-[:HAS_LINE]->(:LineNode)-[:NEXT_STATE]->(:LineNode)  # 实体演化线
    - (:Entity)-[:HAS_RELATION]->(:RelationVersion)-[:TO]->(:Entity) # 关系演化链
    - (:RelationVersion)-[:NEXT_VERSION]->(:RelationVersion)         # 关系版本链
    - (:LineNode)-[:EVIDENCED_BY]->(:Note)                           # 实体证据链接
    - (:RelationVersion)-[:EVIDENCED_BY]->(:Note)                    # 关系证据链接

    设计统一性：
    - 实体状态演化：Entity -> LineNode 链 (CREATE/UPDATE/DELETE)
    - 关系状态演化：Entity -> RelationVersion 链 (CREATE/UPDATE/DELETE)
    - 两者都支持 valid 字段进行时序失效标记
    """

    def __init__(self, debug_log_path=None, embedding_manager=None):
        self.driver = GraphDatabase.driver(
            NEO4J_CONFIG["uri"],
            auth=(NEO4J_CONFIG["user"], NEO4J_CONFIG["password"])
        )
        # 数据库名称（支持多数据库隔离）
        self.database = NEO4J_CONFIG.get("database", "neo4j")
        # 别名映射缓存
        self._alias_map: Dict[str, str] = {}
        # 从 Neo4j 加载全局摘要
        self._global_summary: str = self._load_global_summary()
        # Debug 日志路径
        self.debug_log_path = debug_log_path

        # Embedding Manager - 向量化实体匹配
        self.use_embedding = EMBEDDING_CONFIG.get("use_embedding", True)
        if self.use_embedding:
            if embedding_manager is not None:
                # 使用共享的 embedding_manager（多线程场景）
                self.embedding_manager = embedding_manager
            else:
                # 创建新的 embedding_manager（单线程场景）
                self.embedding_manager = EmbeddingManager(
                    model_path=EMBEDDING_CONFIG["model_path"],
                    similarity_threshold=EMBEDDING_CONFIG["similarity_threshold"]
                )
                # 初始化时缓存所有已有实体的embedding
                self._refresh_embedding_cache()
        else:
            self.embedding_manager = None

    def close(self):
        self.driver.close()

    def _refresh_embedding_cache(self) -> None:
        """刷新embedding缓存 - 缓存所有实体的向量"""
        if not self.use_embedding or not self.embedding_manager:
            return
        entities = self.get_all_entities()
        self.embedding_manager.cache_all_entities(entities)
        if self.debug_log_path:
            with open(self.debug_log_path, "a", encoding="utf-8") as f:
                f.write(f"[DEBUG-KG] Cached {len(entities)} entity embeddings\n")

    # ==================== 基础操作 ====================

    def run_query(self, query: str, params: dict = None) -> List[dict]:
        with self.driver.session(database=self.database) as session:
            result = session.run(query, params or {})
            return [dict(record) for record in result]

    def run_write(self, query: str, params: dict = None) -> None:
        with self.driver.session(database=self.database) as session:
            session.run(query, params or {})

    # ==================== Note 操作 ====================

    def create_note(self, note: dict) -> None:
        """创建 Note 节点"""
        text = f"{note.get('user', '')} {note.get('assistant', '')}"
        self.run_write("""
            MERGE (n:Note {id: $id})
            SET n.session = $session,
                n.session_date = $session_date,
                n.user = $user,
                n.assistant = $assistant,
                n.text = $text,
                n.seq = $seq
        """, {
            "id": note["id"],
            "session": note.get("session", ""),
            "session_date": note.get("session_date", ""),
            "user": note.get("user", ""),
            "assistant": note.get("assistant", ""),
            "text": text,
            "seq": int(note["id"][1:]) if note["id"].startswith("N") else 0
        })

    def create_temporal_edges(self) -> None:
        """创建时序边 (NEXT)"""
        self.run_write("""
            MATCH (n1:Note), (n2:Note)
            WHERE n2.seq = n1.seq + 1
            MERGE (n1)-[:NEXT]->(n2)
        """)

    def get_note_by_id(self, note_id: str) -> Optional[dict]:
        result = self.run_query("MATCH (n:Note {id: $id}) RETURN n", {"id": note_id})
        return dict(result[0]["n"]) if result else None

    def get_notes_by_ids(self, note_ids: List[str]) -> List[dict]:
        result = self.run_query("""
            MATCH (n:Note) WHERE n.id IN $ids
            RETURN n ORDER BY n.seq ASC
        """, {"ids": note_ids})
        return [dict(r["n"]) for r in result]

    def get_all_notes(self) -> List[dict]:
        result = self.run_query("MATCH (n:Note) RETURN n ORDER BY n.seq ASC")
        return [dict(r["n"]) for r in result]

    def get_recent_notes(self, before_seq: int, limit: int = 10) -> List[dict]:
        """获取指定序号之前的最近 N 条 Notes"""
        result = self.run_query("""
            MATCH (n:Note)
            WHERE n.seq < $seq
            RETURN n
            ORDER BY n.seq DESC
            LIMIT $limit
        """, {"seq": before_seq, "limit": limit})
        return [dict(r["n"]) for r in reversed(result)]

    # ==================== Entity 操作 ====================

    def create_or_update_entity(
        self,
        name: str,
        display_name: str,
        entity_type: str,
        summary: str = "",
        heat: int = 1
    ) -> None:
        """创建或更新实体（不再使用aliases，改用embedding匹配）"""
        self.run_write("""
            MERGE (e:Entity {name: $name})
            SET e.display_name = $display_name,
                e.type = $entity_type,
                e.summary = CASE WHEN $summary <> '' THEN $summary ELSE COALESCE(e.summary, '') END,
                e.heat = COALESCE(e.heat, 0) + $heat
        """, {
            "name": name,
            "display_name": display_name,
            "entity_type": entity_type,
            "summary": summary,
            "heat": heat
        })

        # 更新缓存
        self._alias_map[name.lower()] = name

        # 缓存新实体的embedding
        if self.use_embedding and self.embedding_manager:
            self.embedding_manager.cache_entity_embedding(name, summary)

    def update_entity_summary(self, name: str, summary: str) -> None:
        """更新实体摘要"""
        self.run_write("""
            MATCH (e:Entity {name: $name})
            SET e.summary = $summary
        """, {"name": name, "summary": summary})

    def get_entity(self, name: str) -> Optional[dict]:
        """获取实体（支持别名）"""
        canonical = self.resolve_alias(name)
        if not canonical:
            return None
        result = self.run_query("MATCH (e:Entity {name: $name}) RETURN e", {"name": canonical})
        return dict(result[0]["e"]) if result else None

    def get_all_entities(self) -> List[dict]:
        """获取所有实体（按热度排序）"""
        result = self.run_query("""
            MATCH (e:Entity)
            RETURN e
            ORDER BY e.heat DESC
        """)
        return [dict(r["e"]) for r in result]

    def resolve_alias(self, name: str) -> Optional[str]:
        """解析别名到规范名称（精确匹配优先，embedding作为fallback）"""
        name_lower = name.lower()
        # 先检查缓存
        if name_lower in self._alias_map:
            return self._alias_map[name_lower]

        # 策略1: 精确匹配 e.name
        result = self.run_query("""
            MATCH (e:Entity)
            WHERE e.name = $name
            RETURN e.name as name
            LIMIT 1
        """, {"name": name_lower})

        if result:
            canonical = result[0]["name"]
            self._alias_map[name_lower] = canonical
            return canonical

        # 策略2: Embedding-based 相似度匹配（仅在未精确命中时启用）
        if self.use_embedding and self.embedding_manager:
            entities = self.get_all_entities()
            if entities:
                match, score = self.embedding_manager.find_best_match(
                    query=name,
                    candidates=[
                        {"name": e.get("name", ""), "summary": e.get("summary", "")}
                        for e in entities
                    ],
                    return_score=True
                )
                if match:
                    self._alias_map[name_lower] = match
                    if self.debug_log_path:
                        with open(self.debug_log_path, "a", encoding="utf-8") as f:
                            f.write(f"[DEBUG-KG] Embedding match: '{name}' -> '{match}' (score={score:.4f})\n")
                    return match

        # 未找到匹配
        if self.debug_log_path:
            with open(self.debug_log_path, "a", encoding="utf-8") as f:
                f.write(f"[DEBUG-KG] No match found for: '{name}'\n")
        return None

    # ==================== LineNode 操作 (演化线) ====================

    def add_line_node(
        self,
        entity_name: str,
        note_id: str,
        op: str,
        state_change: str,
        context: str,
        valid: bool = True
    ) -> None:
        """
        向实体的演化线添加节点

        Args:
            entity_name: 实体名称
            note_id: 关联的 Note ID
            op: CRUD 操作类型
            state_change: 状态变化描述
            context: 原子笔记摘要
            valid: 是否有效（被后续 UPDATE 覆盖的节点设为 False）
        """
        # ID 加上 op 类型避免同一 Note 里 CREATE/UPDATE/DELETE 冲突
        line_node_id = f"{entity_name}_{note_id}_{op}"
        # 从 note_id 解析 seq，用于排序
        note_seq = int(note_id[1:]) if note_id.startswith("N") else 0

        # 创建 LineNode（存储 seq 用于排序）
        self.run_write("""
            MERGE (ln:LineNode {id: $id})
            SET ln.entity = $entity,
                ln.note_id = $note_id,
                ln.seq = $seq,
                ln.op = $op,
                ln.state_change = $state_change,
                ln.context = $context,
                ln.valid = $valid,
                ln.created_at = timestamp()
        """, {
            "id": line_node_id,
            "entity": entity_name,
            "note_id": note_id,
            "seq": note_seq,
            "op": op,
            "state_change": state_change,
            "context": context,
            "valid": valid
        })

        # 链接到 Entity
        self.run_write("""
            MATCH (e:Entity {name: $entity})
            MATCH (ln:LineNode {id: $id})
            MERGE (e)-[:HAS_LINE]->(ln)
        """, {"entity": entity_name, "id": line_node_id})

        # 链接到 Note
        self.run_write("""
            MATCH (ln:LineNode {id: $id})
            MATCH (n:Note {id: $note_id})
            MERGE (ln)-[:EVIDENCED_BY]->(n)
        """, {"id": line_node_id, "note_id": note_id})

        # 链接到前一个 LineNode (用 seq 排序，保证时序正确)
        self.run_write("""
            MATCH (e:Entity {name: $entity})-[:HAS_LINE]->(prev:LineNode)
            WHERE prev.seq < $seq
            WITH prev ORDER BY prev.seq DESC LIMIT 1
            MATCH (ln:LineNode {id: $id})
            MERGE (prev)-[:NEXT_STATE]->(ln)
        """, {"entity": entity_name, "id": line_node_id, "seq": note_seq})

    def get_entity_line(self, entity_name: str, valid_only: bool = True) -> List[dict]:
        """获取实体的演化线"""
        canonical = self.resolve_alias(entity_name)
        if not canonical:
            return []

        valid_clause = "AND ln.valid = true" if valid_only else ""
        result = self.run_query(f"""
            MATCH (e:Entity {{name: $name}})-[:HAS_LINE]->(ln:LineNode)
            WHERE true {valid_clause}
            RETURN ln
            ORDER BY ln.seq ASC
        """, {"name": canonical})
        return [dict(r["ln"]) for r in result]

    def invalidate_line_nodes_before(self, entity_name: str, before_note_seq: int) -> int:
        """
        将实体在指定 Note 之前的 LineNode 标记为无效

        Returns:
            被标记为无效的节点数量
        """
        canonical = self.resolve_alias(entity_name)
        if not canonical:
            return 0

        result = self.run_query("""
            MATCH (e:Entity {name: $name})-[:HAS_LINE]->(ln:LineNode)-[:EVIDENCED_BY]->(n:Note)
            WHERE n.seq < $seq AND ln.valid = true
            SET ln.valid = false
            RETURN count(ln) as invalidated
        """, {"name": canonical, "seq": before_note_seq})
        return result[0]["invalidated"] if result else 0

    def get_entity_notes(self, entity_name: str, limit: int = 50) -> List[dict]:
        """获取实体相关的所有 Notes（通过 LineNode）"""
        canonical = self.resolve_alias(entity_name)
        if not canonical:
            return []
        result = self.run_query("""
            MATCH (e:Entity {name: $name})-[:HAS_LINE]->(ln:LineNode)-[:EVIDENCED_BY]->(n:Note)
            RETURN DISTINCT n
            ORDER BY n.seq ASC
            LIMIT $limit
        """, {"name": canonical, "limit": limit})
        return [dict(r["n"]) for r in result]

    def get_entity_notes_with_evolution(self, entity_name: str, limit: int = 50) -> List[dict]:
        """
        获取实体相关的 Notes，同时附带 LineNode 演化信息

        Returns:
            [{"note": {...}, "evolutions": [{"entity": "Jon", "op": "UPDATE", "state_change": "银行家→失业"}]}]
        """
        canonical = self.resolve_alias(entity_name)
        if not canonical:
            return []

        result = self.run_query("""
            MATCH (e:Entity {name: $name})-[:HAS_LINE]->(ln:LineNode)-[:EVIDENCED_BY]->(n:Note)
            WHERE ln.valid = true
            WITH n, collect({
                entity: ln.entity,
                op: ln.op,
                state_change: ln.state_change
            }) as evolutions
            RETURN n, evolutions
            ORDER BY n.seq ASC
            LIMIT $limit
        """, {"name": canonical, "limit": limit})

        return [{"note": dict(r["n"]), "evolutions": r["evolutions"]} for r in result]

    def get_notes_with_evolution_by_ids(self, note_ids: List[str]) -> List[dict]:
        """
        根据 Note IDs 获取 Notes，同时附带所有关联的 LineNode 演化信息

        Returns:
            [{"note": {...}, "evolutions": [{"entity": "Jon", "op": "UPDATE", "state_change": "..."}]}]
        """
        if not note_ids:
            return []

        result = self.run_query("""
            MATCH (n:Note)
            WHERE n.id IN $ids
            OPTIONAL MATCH (ln:LineNode)-[:EVIDENCED_BY]->(n)
            WHERE ln.valid = true
            WITH n, collect(
                CASE WHEN ln IS NOT NULL THEN {
                    entity: ln.entity,
                    op: ln.op,
                    state_change: ln.state_change
                } ELSE null END
            ) as raw_evolutions
            WITH n, [e IN raw_evolutions WHERE e IS NOT NULL] as evolutions
            RETURN n, evolutions
            ORDER BY n.seq ASC
        """, {"ids": note_ids})

        return [{"note": dict(r["n"]), "evolutions": r["evolutions"]} for r in result]

    # ==================== Relation 操作 (实体间关系 - RelationVersion 节点) ====================

    def create_relation(
        self,
        source: str,
        relation_type: str,
        target: str,
        note_id: str,
        context: str = "",
        op: str = "CREATE"
    ) -> None:
        """
        创建关系演化版本（RelationVersion 节点）

        设计与 LineNode 统一：
        - 每次关系变化创建新的 RelationVersion 节点
        - 通过 NEXT_VERSION 边连接形成版本链
        - 支持 valid 字段标记时序失效

        Args:
            source: 源实体名称
            relation_type: 关系类型
            target: 目标实体名称
            note_id: 关联的 Note ID
            context: 关系上下文描述
            op: 操作类型 (CREATE/UPDATE/DELETE)
        """
        note_seq = int(note_id[1:]) if note_id.startswith("N") else 0
        # 生成唯一 ID: source_TYPE_target_noteId_op
        rv_id = f"{source}_{relation_type}_{target}_{note_id}_{op}"

        # 1. 创建 RelationVersion 节点
        self.run_write("""
            MERGE (rv:RelationVersion {id: $id})
            SET rv.source = $source,
                rv.target = $target,
                rv.relation_type = $relation_type,
                rv.note_id = $note_id,
                rv.seq = $seq,
                rv.op = $op,
                rv.context = $context,
                rv.valid = true,
                rv.created_at = timestamp()
        """, {
            "id": rv_id,
            "source": source,
            "target": target,
            "relation_type": relation_type,
            "note_id": note_id,
            "seq": note_seq,
            "op": op,
            "context": context
        })

        # 2. 链接到源实体: Entity -[:HAS_RELATION]-> RelationVersion
        self.run_write("""
            MATCH (e:Entity {name: $source})
            MATCH (rv:RelationVersion {id: $id})
            MERGE (e)-[:HAS_RELATION]->(rv)
        """, {"source": source, "id": rv_id})

        # 3. 链接到目标实体: RelationVersion -[:TO]-> Entity
        self.run_write("""
            MATCH (rv:RelationVersion {id: $id})
            MATCH (t:Entity {name: $target})
            MERGE (rv)-[:TO]->(t)
        """, {"id": rv_id, "target": target})

        # 4. 链接到 Note: RelationVersion -[:EVIDENCED_BY]-> Note
        self.run_write("""
            MATCH (rv:RelationVersion {id: $id})
            MATCH (n:Note {id: $note_id})
            MERGE (rv)-[:EVIDENCED_BY]->(n)
        """, {"id": rv_id, "note_id": note_id})

        # 5. 链接到前一个 RelationVersion (同一 source-type-target 组合)
        self.run_write("""
            MATCH (e:Entity {name: $source})-[:HAS_RELATION]->(prev:RelationVersion)-[:TO]->(t:Entity {name: $target})
            WHERE prev.relation_type = $relation_type AND prev.seq < $seq
            WITH prev ORDER BY prev.seq DESC LIMIT 1
            MATCH (rv:RelationVersion {id: $id})
            MERGE (prev)-[:NEXT_VERSION]->(rv)
        """, {
            "source": source,
            "target": target,
            "relation_type": relation_type,
            "seq": note_seq,
            "id": rv_id
        })

        # 6. 如果是 UPDATE/DELETE，将之前的版本标记为无效
        if op in ("UPDATE", "DELETE"):
            self.invalidate_relation_versions_before(source, relation_type, target, note_seq)

    def invalidate_relation_versions_before(
        self,
        source: str,
        relation_type: str,
        target: str,
        before_seq: int
    ) -> int:
        """
        将指定关系在某时间点之前的版本标记为无效

        Returns:
            被标记为无效的版本数量
        """
        result = self.run_query("""
            MATCH (e:Entity {name: $source})-[:HAS_RELATION]->(rv:RelationVersion)-[:TO]->(t:Entity {name: $target})
            WHERE rv.relation_type = $relation_type AND rv.seq < $seq AND rv.valid = true
            SET rv.valid = false
            RETURN count(rv) as invalidated
        """, {
            "source": source,
            "target": target,
            "relation_type": relation_type,
            "seq": before_seq
        })
        return result[0]["invalidated"] if result else 0

    def get_relation_versions(
        self,
        source: str,
        relation_type: str,
        target: str,
        valid_only: bool = True
    ) -> List[dict]:
        """获取关系的版本链"""
        valid_clause = "AND rv.valid = true" if valid_only else ""
        result = self.run_query(f"""
            MATCH (e:Entity {{name: $source}})-[:HAS_RELATION]->(rv:RelationVersion)-[:TO]->(t:Entity {{name: $target}})
            WHERE rv.relation_type = $relation_type {valid_clause}
            RETURN rv
            ORDER BY rv.seq ASC
        """, {
            "source": source,
            "target": target,
            "relation_type": relation_type
        })
        return [dict(r["rv"]) for r in result]

    def get_entity_relations(self, entity_name: str, include_history: bool = False, valid_only: bool = True) -> List[dict]:
        """
        获取实体的所有关系（基于 RelationVersion 节点）

        Args:
            entity_name: 实体名称
            include_history: 是否包含完整版本历史
            valid_only: 是否只返回有效版本

        Returns:
            关系列表
        """
        canonical = self.resolve_alias(entity_name)
        if not canonical:
            return []

        valid_clause = "AND rv.valid = true" if valid_only else ""

        if include_history:
            # 返回完整版本历史
            result = self.run_query(f"""
                MATCH (e:Entity {{name: $name}})-[:HAS_RELATION]->(rv:RelationVersion)-[:TO]->(other:Entity)
                WHERE true {valid_clause}
                WITH other.name as related_entity,
                     rv.relation_type as relation,
                     'outgoing' as direction,
                     collect({{
                         note_id: rv.note_id,
                         seq: rv.seq,
                         op: rv.op,
                         context: rv.context,
                         valid: rv.valid
                     }}) as versions
                RETURN relation, related_entity, direction, versions
                UNION
                MATCH (other:Entity)-[:HAS_RELATION]->(rv:RelationVersion)-[:TO]->(e:Entity {{name: $name}})
                WHERE true {valid_clause}
                WITH other.name as related_entity,
                     rv.relation_type as relation,
                     'incoming' as direction,
                     collect({{
                         note_id: rv.note_id,
                         seq: rv.seq,
                         op: rv.op,
                         context: rv.context,
                         valid: rv.valid
                     }}) as versions
                RETURN relation, related_entity, direction, versions
            """, {"name": canonical})
            return [dict(r) for r in result]
        else:
            # 只返回最新有效版本
            result = self.run_query(f"""
                MATCH (e:Entity {{name: $name}})-[:HAS_RELATION]->(rv:RelationVersion)-[:TO]->(other:Entity)
                WHERE rv.valid = true
                WITH other.name as related_entity,
                     rv.relation_type as relation,
                     'outgoing' as direction,
                     rv.context as context,
                     rv.seq as seq
                ORDER BY seq DESC
                WITH relation, related_entity, direction, collect(context)[0] as context
                RETURN relation, related_entity, direction, context
                UNION
                MATCH (other:Entity)-[:HAS_RELATION]->(rv:RelationVersion)-[:TO]->(e:Entity {{name: $name}})
                WHERE rv.valid = true
                WITH other.name as related_entity,
                     rv.relation_type as relation,
                     'incoming' as direction,
                     rv.context as context,
                     rv.seq as seq
                ORDER BY seq DESC
                WITH relation, related_entity, direction, collect(context)[0] as context
                RETURN relation, related_entity, direction, context
            """, {"name": canonical})
            return [dict(r) for r in result]

    # ==================== 图遍历检索 ====================

    def find_common_connections(self, entity_names: List[str]) -> List[dict]:
        """找到多个实体的共同连接（交集查询）"""
        if len(entity_names) < 2:
            return []

        canonical_names = [self.resolve_alias(n) for n in entity_names if self.resolve_alias(n)]
        if len(canonical_names) < 2:
            return []

        # 两个实体的共同连接
        if len(canonical_names) == 2:
            result = self.run_query("""
                MATCH (e1:Entity {name: $name1})-[r1]->(common)<-[r2]-(e2:Entity {name: $name2})
                WHERE common:Entity OR common:Activity OR common:Location OR common:Event
                RETURN common, type(r1) as rel1, type(r2) as rel2
            """, {"name1": canonical_names[0], "name2": canonical_names[1]})
            return [{"common": dict(r["common"]), "rel1": r["rel1"], "rel2": r["rel2"]} for r in result]

        return []

    def graph_traversal(
        self,
        start_entities: List[str],
        max_hops: int = 2,
        limit: int = 50,
        current_seq: Optional[int] = None,
        apply_heat_decay: bool = True
    ) -> List[dict]:
        """
        图遍历检索：从起始实体出发，收集相关 Notes
        包括：直接关联 + 关系连接 + 共同连接 + 热度衰减

        Args:
            start_entities: 起始实体列表
            max_hops: 最大跳数
            limit: 返回 Note 数量上限
            current_seq: 当前时间点（用于计算热度衰减）
            apply_heat_decay: 是否应用热度衰减
        """
        canonical_names = [self.resolve_alias(n) for n in start_entities if self.resolve_alias(n)]
        if not canonical_names:
            return []

        # 获取当前最大 seq（如果未提供）
        if current_seq is None:
            max_seq_result = self.run_query("MATCH (n:Note) RETURN max(n.seq) as max_seq")
            current_seq = max_seq_result[0]["max_seq"] if max_seq_result else 100

        note_scores = {}
        heat_decay_factor = MTM_CONFIG["heat_decay_factor"]

        # 1. 直接关联的 Notes（通过 LineNode）+ 热度衰减
        if apply_heat_decay:
            direct = self.run_query("""
                MATCH (e:Entity)-[:HAS_LINE]->(ln:LineNode)-[:EVIDENCED_BY]->(n:Note)
                WHERE e.name IN $names AND ln.valid = true
                WITH e, ln, n, ($current_seq - n.seq) as age
                RETURN DISTINCT n,
                       1.0 * ($decay_factor ^ age) as score,
                       e.heat as entity_heat
            """, {
                "names": canonical_names,
                "current_seq": current_seq,
                "decay_factor": heat_decay_factor
            })
        else:
            direct = self.run_query("""
                MATCH (e:Entity)-[:HAS_LINE]->(ln:LineNode)-[:EVIDENCED_BY]->(n:Note)
                WHERE e.name IN $names AND ln.valid = true
                RETURN DISTINCT n, 1.0 as score, e.heat as entity_heat
            """, {"names": canonical_names})

        for r in direct:
            note_id = r["n"]["id"]
            # 综合考虑时间衰减分数和实体热度
            entity_heat = r.get("entity_heat", 1)
            final_score = r["score"] * (1 + entity_heat * 0.01)  # 热度加成

            if note_id not in note_scores or note_scores[note_id]["score"] < final_score:
                note_scores[note_id] = {"note": dict(r["n"]), "score": final_score}

        # 2. 通过关系连接的 Notes + 热度衰减
        if apply_heat_decay:
            related = self.run_query("""
                MATCH (e:Entity)-[]->(other:Entity)-[:HAS_LINE]->(ln:LineNode)-[:EVIDENCED_BY]->(n:Note)
                WHERE e.name IN $names AND NOT other.name IN $names AND ln.valid = true
                WITH other, n, ($current_seq - n.seq) as age
                RETURN DISTINCT n,
                       0.5 * ($decay_factor ^ age) as score,
                       other.heat as entity_heat
            """, {
                "names": canonical_names,
                "current_seq": current_seq,
                "decay_factor": heat_decay_factor
            })
        else:
            related = self.run_query("""
                MATCH (e:Entity)-[]->(other:Entity)-[:HAS_LINE]->(ln:LineNode)-[:EVIDENCED_BY]->(n:Note)
                WHERE e.name IN $names AND NOT other.name IN $names AND ln.valid = true
                RETURN DISTINCT n, 0.5 as score, other.heat as entity_heat
            """, {"names": canonical_names})

        for r in related:
            note_id = r["n"]["id"]
            entity_heat = r.get("entity_heat", 1)
            final_score = r["score"] * (1 + entity_heat * 0.01)

            if note_id not in note_scores or note_scores[note_id]["score"] < final_score:
                note_scores[note_id] = {"note": dict(r["n"]), "score": final_score}

        # 3. 共同连接的 Notes（如果有多个实体）
        if len(canonical_names) >= 2:
            common = self.find_common_connections(canonical_names)
            for c in common:
                if "common" in c and isinstance(c["common"], dict):
                    common_name = c["common"].get("name")
                    if common_name:
                        common_notes = self.get_entity_notes(common_name, limit=10)
                        for n in common_notes:
                            note_id = n["id"]
                            if note_id not in note_scores:
                                # 共同连接的 Notes 也应用衰减
                                age = current_seq - n.get("seq", 0)
                                decay_score = 0.8 * (heat_decay_factor ** age if apply_heat_decay else 1.0)
                                note_scores[note_id] = {"note": n, "score": decay_score}

        # 排序返回（分数高的优先，同分数则新的优先）
        sorted_notes = sorted(note_scores.values(), key=lambda x: (-x["score"], -x["note"]["seq"]))
        return [item["note"] for item in sorted_notes[:limit]]

    # ==================== 全局摘要 ====================

    def _load_global_summary(self) -> str:
        """从 Neo4j 加载全局摘要"""
        result = self.run_query("""
            MATCH (g:GlobalState {type: 'summary'})
            RETURN g.content as content
        """)
        return result[0]["content"] if result else ""

    def get_global_summary(self) -> str:
        return self._global_summary

    def set_global_summary(self, summary: str) -> None:
        self._global_summary = summary
        # 持久化到 Neo4j
        self.run_write("""
            MERGE (g:GlobalState {type: 'summary'})
            SET g.content = $content, g.updated_at = timestamp()
        """, {"content": summary})

    # ==================== 统计 ====================

    def get_stats(self) -> dict:
        stats = {}

        result = self.run_query("MATCH (n:Note) RETURN count(n) as c")
        stats["note_count"] = result[0]["c"] if result else 0

        result = self.run_query("MATCH (e:Entity) RETURN count(e) as c")
        stats["entity_count"] = result[0]["c"] if result else 0

        result = self.run_query("MATCH (ln:LineNode) RETURN count(ln) as c")
        stats["line_node_count"] = result[0]["c"] if result else 0

        # RelationVersion 统计
        result = self.run_query("MATCH (rv:RelationVersion) RETURN count(rv) as c")
        stats["relation_version_count"] = result[0]["c"] if result else 0

        result = self.run_query("MATCH (rv:RelationVersion) WHERE rv.valid = true RETURN count(rv) as c")
        stats["relation_version_valid"] = result[0]["c"] if result else 0

        result = self.run_query("MATCH ()-[r:NEXT]->() RETURN count(r) as c")
        stats["temporal_edges"] = result[0]["c"] if result else 0

        # 按类型统计实体
        result = self.run_query("""
            MATCH (e:Entity)
            RETURN e.type as type, count(e) as count
        """)
        for r in result:
            if r["type"]:
                stats[f"entity_{r['type']}"] = r["count"]

        # 按关系类型统计 RelationVersion
        result = self.run_query("""
            MATCH (rv:RelationVersion)
            RETURN rv.relation_type as type, count(rv) as count
        """)
        for r in result:
            if r["type"]:
                stats[f"relation_{r['type']}"] = r["count"]

        return stats

    # ==================== 清理 ====================

    def clear_all(self) -> None:
        self.run_write("MATCH (n) DETACH DELETE n")
        self._alias_map.clear()
        self._global_summary = ""

    def clear_entities_and_lines(self) -> None:
        """只清除实体、演化线和关系版本，保留 Notes"""
        self.run_write("MATCH (e:Entity) DETACH DELETE e")
        self.run_write("MATCH (ln:LineNode) DETACH DELETE ln")
        self.run_write("MATCH (rv:RelationVersion) DETACH DELETE rv")
        self.run_write("MATCH (g:GlobalState) DELETE g")  # 清除全局状态
        self._alias_map.clear()
        self._global_summary = ""  # 清空内存中的摘要


if __name__ == "__main__":
    kg = KnowledgeGraphV2()
    print("Stats:", kg.get_stats())
    kg.close()