"""
变换计算：沿路径累积矩阵、从 root 计算到所有节点。不负责 I/O 和图构建。
"""
import numpy as np
from typing import Dict, List, Optional

from .transform import Transform


def get_transform_along_path(
    transforms: Dict[str, Transform],
    path: List[str],
) -> Optional[np.ndarray]:
    """
    沿 path 累积变换矩阵，path 为节点序列 [n0, n1, ..., nk]。
    若某条边不存在（正向与逆向都不存在）则返回 None。
    """
    if len(path) < 2:
        return np.eye(4)
    T = np.eye(4)
    for i in range(len(path) - 1):
        from_frame = path[i]
        to_frame = path[i + 1]
        key = f"{from_frame}->{to_frame}"
        if key in transforms:
            T = T @ transforms[key].to_matrix()
        else:
            key_inv = f"{to_frame}->{from_frame}"
            if key_inv in transforms:
                T = T @ transforms[key_inv].inverse().to_matrix()
            else:
                return None
    return T


def compute_all_from_root(
    transforms: Dict[str, Transform],
    graph: Dict[str, set],  # Set[str] per node
    root: str,
) -> Dict[str, np.ndarray]:
    """
    从 root 到图中所有可达节点的变换矩阵。
    调用方需保证 root 在 graph 中；否则由调用方抛异常。
    """
    from .graph_ops import find_path_bfs

    result = {root: np.eye(4)}
    for node in graph:
        if node == root:
            continue
        path = find_path_bfs(graph, root, node)
        if path is None:
            continue
        T = get_transform_along_path(transforms, path)
        if T is not None:
            result[node] = T
    return result
