from collections import deque

class DAG:
    def __init__(self):
        self.predecessors = {}
        self.successors = {}
        self.graph = {
            "predecessors":self.predecessors,
            "successors":self.successors,
        }

    def add_node(self, node: str):
        if node not in self.predecessors:
            self.predecessors[node] = set()
            self.successors[node] = set()

    def add_edge(self, u: str, v: str):
        # u -> v
        self.add_node(u)
        self.add_node(v)
        self.graph["predecessors"][v].add(u)
        self.graph["successors"][u].add(v)

    def find_predecessors(self, node: str):
        return self.predecessors.get(node, set())

    def find_successors(self, node: str):
        return self.successors.get(node, set())

    @staticmethod
    def get_graph_nodes(graph):
        return list(graph["predecessors"].keys())

    @staticmethod
    def topological_sort(graph):
        predecessors = graph["predecessors"]
        successors = graph["successors"]
        
        # 1. 复制一份入度统计，避免修改原图
        in_degree = {node: len(preds) for node, preds in predecessors.items()}
        
        # 2. 将所有入度为 0 的节点放入队列
        queue = deque([node for node, degree in in_degree.items() if degree == 0])
        
        topo_order = []
        
        # 3. 开始处理
        while queue:
            u = queue.popleft()
            topo_order.append(u)
            # 遍历该节点的所有后继节点
            for v in successors[u]:
                # 模拟移除边 u -> v，因此 v 的入度减 1
                in_degree[v] -= 1
                # 如果入度变为 0，则加入队列
                if in_degree[v] == 0:
                    queue.append(v)
        
        # 4. 检查是否存在环
        if len(topo_order) != len(predecessors):
            raise ValueError("Graph contains a cycle, topological sort is not possible.")
        return topo_order

    @staticmethod
    def get_predecessor_subgraph(graph, target_nodes):
        # 1. 查找所有祖先节点（包含目标节点本身）
        all_ancestors = set()
        stack = list(target_nodes)
        
        while stack:
            node = stack.pop()
            if node not in all_ancestors:
                all_ancestors.add(node)
                # 获取当前节点的所有前继，加入栈中继续向上寻找
                preds = graph["predecessors"].get(node, set())
                stack.extend(preds)
                
        # 2. 初始化子图结构
        sub_predecessors = {}
        sub_successors = {}
        
        # 3. 构建子图内容
        for node in all_ancestors:
            # 提取前继：只需过滤掉不在 all_ancestors 里的节点
            #（由于是前继子图，祖先的前继一定也在 all_ancestors 中，这里 set 过滤是为了严谨）
            sub_predecessors[node] = {p for p in graph["predecessors"].get(node, set()) if p in all_ancestors}
            # 提取后继：必须过滤，因为原图中的某些后继节点可能不是目标节点的祖先
            sub_successors[node] = {s for s in graph["successors"].get(node, set()) if s in all_ancestors}
            
        subgraph = {
            "predecessors": sub_predecessors,
            "successors": sub_successors
        }
        return subgraph

    @staticmethod
    def get_successor_subgraph(graph, target_nodes):
        # 1. 查找所有后代节点（包含目标节点本身）
        all_descendants = set()
        stack = list(target_nodes)
        
        while stack:
            node = stack.pop()
            if node not in all_descendants:
                all_descendants.add(node)
                # 获取当前节点的所有后继，加入栈中继续向下寻找
                succs = graph["successors"].get(node, set())
                stack.extend(succs)
                
        # 2. 初始化子图结构
        sub_predecessors = {}
        sub_successors = {}
        
        # 3. 构建子图内容
        for node in all_descendants:
            # 提取后继：过滤掉不在 all_descendants 里的节点
            # (理论上后代的后继一定在 all_descendants 中，此处过滤为保证结构纯净)
            sub_successors[node] = {s for s in graph["successors"].get(node, set()) if s in all_descendants}
            # 提取前继：必须过滤，因为原图中的某些前继节点可能不是目标节点的后代
            sub_predecessors[node] = {p for p in graph["predecessors"].get(node, set()) if p in all_descendants}
            
        subgraph = {
            "predecessors": sub_predecessors,
            "successors": sub_successors
        }
        return subgraph


if __name__ == "__main__":
    dag = DAG()

    
    dag.add_edge("a", "b")
    dag.add_edge("b", "c")
    dag.add_edge("1", "c")
    dag.add_edge("d", "e")
    dag.add_edge("c", "d")


    # print(DAG.topological_sort(dag.graph))

    # subgraph = DAG.get_predecessor_subgraph(dag.graph, ["b"])
    # print(DAG.topological_sort(subgraph))

    subgraph = DAG.get_successor_subgraph(dag.graph, ["a","1"])
    print(DAG.topological_sort(subgraph))
