nix_builder/python_experiments/depth_first_search/inverted_tree.py

#!/usr/bin/env python
#
from __future__ import annotations

from collections import deque
from dataclasses import dataclass
from typing import Callable, Iterable


def main():
    activity_tree = ActivityTree()
    activity_tree.activities[0] = Activity(id=0, parent=0, text="*")
    activity_tree.activities[1] = Activity(id=1, parent=0, text="A")
    activity_tree.activities[2] = Activity(id=2, parent=1, text="B")
    activity_tree.activities[3] = Activity(id=3, parent=2, text="C")
    activity_tree.activities[4] = Activity(id=4, parent=2, text="D")
    activity_tree.activities[5] = Activity(id=5, parent=4, text="E")
    activity_tree.activities[6] = Activity(id=6, parent=4, text="F")
    activity_tree.activities[7] = Activity(id=7, parent=1, text="G")
    activity_tree.activities[8] = Activity(id=8, parent=1, text="H")
    activity_tree.activities[9] = Activity(id=9, parent=8, text="I")
    activity_tree.activities[10] = Activity(id=10, parent=9, text="J")

    match_iter = MatchIter(
        activity_tree,
        lambda act: act.id in (2, 5, 10),
        lambda act: act.id in (4,),
        lambda act: act.id in (1, 2, 3, 4, 5, 6, 7, 8, 9, 10),
    )

    while (depth_and_node := match_iter.get_next()) is not None:
        depth, matched, transparent, node = depth_and_node
        print("\t".join((str(depth), str(matched), str(transparent), node.text)))

    # while True:
    #     print("\n\n\nvvvvvvvvvv\n\n")
    #     node = match_iter.get_next()
    #     if node is None:
    #         break
    #     print(node.text)
    #     print("\n\n")
    #     match_iter.print_stack()
    #     print("\n\n\n^^^^^^^^^^\n\n")


class ActivityTree:
    activities: dict[int, Activity]

    def __init__(self) -> None:
        self.activities = {}


@dataclass
class Activity:
    id: int
    parent: int
    text: str


class MatchIter:
    activity_tree: ActivityTree
    stack: deque[StackEntry]
    is_match: Callable[[Activity], bool]
    is_transparent: Callable[[Activity], bool]
    is_alive: Callable[[Activity], bool]

    def __init__(
        self,
        activity_tree: ActivityTree,
        is_match: Callable[[Activity], bool],
        is_transparent: Callable[[Activity], bool],
        is_alive: Callable[[Activity], bool],
    ) -> None:
        self.activity_tree = activity_tree
        self.stack = deque([StackEntry(0)])
        self.is_match = is_match
        self.is_transparent = is_transparent
        self.is_alive = is_alive

    def get_next(self) -> tuple[int, bool, bool, Activity] | None:
        while len(self.stack) > 0:
            stack_entry = self.stack[-1]
            node = self.activity_tree.activities[stack_entry.node_id]

            if stack_entry.visited:
                if stack_entry.should_announce:
                    stack_entry_depth = len(
                        list(
                            entry
                            for entry in self.walk_up_stack(stack_entry)
                            if not entry.transparent
                        )
                    )
                    self.stack.pop()
                    return (
                        stack_entry_depth,
                        stack_entry.matched,
                        stack_entry.transparent,
                        node,
                    )
                self.stack.pop()
                continue

            if self.is_match(node):
                stack_entry.should_announce = True
                stack_entry.matched = True
                for parent in self.walk_up_stack(stack_entry):
                    if parent.node_id != 0 and not parent.should_announce:
                        parent.should_announce = True

            stack_entry.visited = True
            stack_entry.transparent = self.is_transparent(node)

            children = [
                c
                for c in self.activity_tree.activities.values()
                if c.parent == node.id and self.is_alive(c)
            ]
            for c in children:
                if c.id != 0:
                    self.stack.append(StackEntry(c.id))

        return None

    def print_stack(self) -> None:
        for stack_entry in self.stack:
            print(stack_entry.text(self.activity_tree))

    def walk_up_stack(self, base: StackEntry) -> Iterable[StackEntry]:
        found_start = False
        for i in range(len(self.stack)):
            idx = -1 - i
            stack_entry = self.stack[idx]
            if not found_start:
                if stack_entry is base:
                    found_start = True
                continue
            if stack_entry.visited:
                yield stack_entry


@dataclass
class StackEntry:
    node_id: int
    announced: bool = False
    should_announce: bool = False
    visited: bool = False
    matched: bool = False
    transparent: bool = False

    def text(self, activity_tree: ActivityTree) -> str:
        node = activity_tree.activities[self.node_id]
        return f"{node.text}\t{"annced" if self.announced else "-"}\t{"should" if self.should_announce else "-"}\t{"visited" if self.visited else "-"}\t{"matched" if self.matched else "-"}\t{"trans" if self.transparent else "-"}"


if __name__ == "__main__":
    main()