Agent/examples_how/overall_derivation/utils/conditional_ratio_calc.py

"""
条件概率计算工具：
1）计算某个人设树节点在父节点下的条件概率；
2）计算某个 pattern 的条件概率。
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

# 已推导列表：每项为 (已推导的选题点, 推导来源人设树节点)，如 ("分享","分享")、("柴犬","动物角色")
# 推导来源人设树节点的 post_ids 在计算条件概率时从人设树中读取
DerivedItem = tuple[str, str]


def _tree_dir(account_name: str, base_dir: Path | None = None) -> Path:
    """人设树目录：../input/{account_name}/原始数据/tree/（相对本文件所在目录）。"""
    if base_dir is not None:
        return base_dir / account_name / "原始数据" / "tree"
    return Path(__file__).resolve().parent.parent / "input" / account_name / "原始数据" / "tree"


def _load_trees(account_name: str, base_dir: Path | None = None) -> list[tuple[str, dict]]:
    """加载该账号下所有维度的人设树。返回 [(维度名, 根节点 dict), ...]。"""
    td = _tree_dir(account_name, base_dir)
    if not td.is_dir():
        return []
    result = []
    for p in td.glob("*.json"):
        try:
            with open(p, "r", encoding="utf-8") as f:
                data = json.load(f)
            # 文件格式为 { "实质": { ... } } 或 { "形式": { ... } }
            for dim_name, root in data.items():
                if isinstance(root, dict):
                    result.append((dim_name, root))
                break
        except Exception:
            continue
    return result


def _post_ids_of(node: dict) -> list[str]:
    """从树节点中取出 _post_ids，无则返回空列表。"""
    return list(node.get("_post_ids") or [])


def _build_node_index(account_name: str, base_dir: Path | None = None) -> dict[str, tuple[list[str], list[str]]]:
    """
    遍历所有维度的人设树，建立 节点名 -> (该节点 post_ids, 父节点 post_ids)。
    同一节点名在多个分支出现时，保留第一次遇到的（保证父子一致）。
    """
    index: dict[str, tuple[list[str], list[str]]] = {}
    for _dim, root in _load_trees(account_name, base_dir):
        parent_pids = _post_ids_of(root)

        def walk(parent_ids: list[str], node_dict: dict) -> None:
            for name, child in (node_dict.get("children") or {}).items():
                if not isinstance(child, dict):
                    continue
                if name not in index:
                    index[name] = (_post_ids_of(child), list(parent_ids))
                walk(_post_ids_of(child), child)

        walk(parent_pids, root)
    return index


def _derived_post_ids_from_sources(
    derived_list: list[DerivedItem],
    index: dict[str, tuple[list[str], list[str]]],
) -> set[str]:
    """根据 derived_list 中的「推导来源人设树节点」在人设树中的 post_ids 取交集，得到已推导的帖子集合。"""
    common: set[str] | None = None
    for _topic_point, source_node in derived_list:
        if source_node not in index:
            continue
        pids = set(index[source_node][0])
        if common is None:
            common = pids
        else:
            common &= pids
    return common if common is not None else set()


def calc_node_conditional_ratio(
    account_name: str,
    derived_list: list[DerivedItem],
    tree_node_name: str,
    base_dir: Path | None = None,
) -> float:
    """
    计算人设树节点 N 在父节点 P 下的条件概率。

    参数:
        account_name: 账号名称
        derived_list: 已推导列表，每项 (已推导的选题点, 推导来源人设树节点)
        tree_node_name: 人设树节点 N 的名称（字符串匹配）
        base_dir: 可选，input 根目录；不传则使用相对本文件的 ../input

    计算规则:
        已推导的帖子集合 = 各「推导来源人设树节点」在人设树中的 post_ids 的交集（方法内从树读取）
        分子 = |已推导的帖子集合 ∩ N 的 post_ids|
        分母 = |已推导的帖子集合 ∩ P 的 post_ids|
        条件概率 = 分子/分母，且 ≤1；分母为 0 时返回 1。
    """
    index = _build_node_index(account_name, base_dir)
    derived_post_ids = _derived_post_ids_from_sources(derived_list, index)
    if tree_node_name not in index:
        return 1.0
    n_pids, p_pids = index[tree_node_name]
    set_n = set(n_pids)
    set_p = set(p_pids)
    den = len(derived_post_ids & set_p)
    if den == 0:
        return 0.0
    num = len(derived_post_ids & set_n)
    return min(1.0, num / den)


def _pattern_nodes_and_post_count(pattern: dict[str, Any]) -> tuple[list[str], int]:
    """从 pattern 中解析出节点列表和 post_count。支持 nodes + post_count 或 i + post_count。"""
    nodes = pattern.get("nodes")
    if nodes is not None and isinstance(nodes, list):
        nodes = [str(x).strip() for x in nodes if x]
    else:
        raw = pattern.get("i") or pattern.get("pattern_str") or ""
        nodes = [x.strip() for x in str(raw).replace("+", " ").split() if x.strip()]
    post_count = int(pattern.get("post_count", 0))
    return nodes, post_count


def calc_pattern_conditional_ratio(
    account_name: str,
    derived_list: list[DerivedItem],
    pattern: dict[str, Any],
    base_dir: Path | None = None,
) -> float:
    """
    计算某个 pattern 的条件概率。

    参数:
        account_name: 账号名称
        derived_list: 已推导列表，每项 (已推导的选题点, 推导来源人设树节点)
        pattern: 至少包含节点列表与 post_count。
                 - 节点列表: key 为 "nodes"（list）或 "i"（字符串，用 + 连接）
                 - post_count: 该 pattern 的帖子数量，作为分子
        base_dir: 可选，input 根目录

    计算规则:
        取 pattern 中「已被推导」的节点（其名称出现在 derived 的推导来源中），
        在人设树中取这些节点的 post_ids 的交集作为分母；
        分子 = pattern.post_count。
        条件概率 = 分子/分母，且 ≤1；分母为 0 时返回 1。
    """
    pattern_nodes, post_count = _pattern_nodes_and_post_count(pattern)
    if not pattern_nodes or post_count <= 0:
        return 1.0

    derived_sources = set(source for _post, source in derived_list)
    # pattern 中已被推导的节点
    derived_pattern_nodes = [n for n in pattern_nodes if n in derived_sources]
    if not derived_pattern_nodes:
        return 1.0

    index = _build_node_index(account_name, base_dir)
    # 仅使用在人设树中存在的「已被推导」节点，取它们在树中的 post_ids 的交集
    derived_in_tree = [n for n in derived_pattern_nodes if n in index]
    if not derived_in_tree:
        return 1.0
    common: set[str] | None = None
    for name in derived_in_tree:
        pids = set(index[name][0])
        if common is None:
            common = pids
        else:
            common &= pids
    if common is None or len(common) == 0:
        return 1.0
    den = len(common)
    return min(1.0, post_count / den)


def _test_with_user_example() -> None:
    """
    使用你提供的测试数据：已推导 (分享|分享)、(柴犬|动物角色)；
    人设树节点：恶作剧；pattern：分享+动物角色+创意表达 post_count=2。
    推导来源的 post_ids 在方法内部从人设树读取。
    """
    account_name = "家有大志"
    # 已推导列表：(已推导的选题点, 推导来源人设树节点)
    derived_list: list[DerivedItem] = [
        ("分享", "分享"),
        # ("柴犬", "动物角色"),
    ]

    # 1）人设树节点「恶作剧」的条件概率
    r_node = calc_node_conditional_ratio(account_name, derived_list, "恶作剧")
    print(f"1) 人设树节点「恶作剧」条件概率: {r_node}")

    # 2）pattern 分享+动物角色+创意表达 post_count=2 的条件概率
    pattern = {"i": "分享+动物角色+创意表达", "post_count": 2}
    r_pattern = calc_pattern_conditional_ratio(account_name, derived_list, pattern)
    print(f"2) pattern 分享+动物角色+创意表达 (post_count=2) 条件概率: {r_pattern}")


if __name__ == "__main__":
    _test_with_user_example()