graph.py

#!/usr/bin/python

''' A basic graph representation with graphviz/DOT and railroad representations.
'''

from collections import defaultdict
from dataclasses import dataclass, field
from typing import Any, Iterable, Optional

from cs.ascii_art import (
    RRBase,
    RRChoice,
    RRMerge,
    RROptional,
    RRSequence,
    RRSplit,
    RRStack,
    RRTextBox,
    RR_START,
    RR_END,
)
from cs.gvutils import Graph as GVGraph, Node as GVNode
from cs.queues import ListQueue

from typeguard import typechecked

@dataclass
class Node:
  ''' A node in a `Graph`.
  '''
  name: Optional[str] = None
  in_edges: list["Edge"] = field(default_factory=list)
  out_edges: list["Edge"] = field(default_factory=list)
  attrs: dict = field(default_factory=dict)
  refobj: Any = None

  def __str__(self):
    return f'{self.__class__.__name__}:{id(self) if self.name is None else repr(self.name)}'

  __repr__ = __str__

  def __hash__(self):
    return id(self)

  def __eq__(self, other):
    return self is other

  def __lt__(self, other):
    return self.name < other.name

  @property
  def in_count(self):
    ''' The number of inbound edges.
    '''
    return len(self.in_edges)

  @property
  def out_count(self):
    ''' The number of outbound edges.
    '''
    return len(self.out_edges)

  @property
  def in_nodes(self):
    ''' The inbound `Node`s.
    '''
    return [edge.in_node for edge in self.in_edges]

  @property
  def in_node(self):
    ''' The inbound/parent node, if there is exactly one inbound edge.
    '''
    edge, = self.in_edges
    return edge.in_node

  @property
  def out_node(self):
    ''' The outbound/child node, if there is exactly one outbound edge.
    '''
    edge, = self.out_edges
    return edge.out_node

  @property
  def out_nodes(self):
    ''' The outbound `Node`s.
    '''
    return [edge.out_node for edge in self.out_edges]

  def as_GVNode(self):
    return GVNode(id=self.name or str(id(self)), attrs=dict(self.attrs))

  def as_dot(self, *, no_attrs=False):
    return self.as_GVNode().as_dot(no_attrs=no_attrs)

  def as_railroad(self):
    return RRTextBox(self.name or str(self))

  def reach_connectivity(
      self,
      reach=None,
      path=None,
      seen=None,
  ) -> dict['Node', dict['Node', list[list['Node']]]]:
    ''' Construct an adjacency graph for the entire network
        with a mapping per-Node to a mapping of each reachable
        destination Node to a list of the available routes as
        lists of the intermediate `Node`s.
    '''
    # TODO: this duplicates tail routes reachable in multiple
    # ways because they are traversed once for each prior route.
    # Eg X => B -> E lists 2 B->E routes because the B to E tail
    # route was traversed twice.
    if reach is None:
      reach = defaultdict(lambda: defaultdict(list))
    if path is None:
      path = []
    if seen is None:
      seen = set()
    for i, ancestor_node in enumerate(path):
      reach[ancestor_node][self].append(path[i + 1:])
    path.append(self)
    for out_node in self.out_nodes:
      if out_node not in path:
        out_node.reach_connectivity(reach, path, seen)
    path.pop()
    return reach

@dataclass
class Edge:
  ''' A directed edge between `Node`s in a `Graph`.
  '''
  in_node: Node
  out_node: Node
  name: Optional[str] = None
  attrs: dict = field(default_factory=dict)
  refobj: Any = None

  def __str__(self):
    return f'{self.__class__.__name__}:{id(self)}' if self.name is None else self.name

  def __hash__(self):
    return id(self)

  def __eq__(self, other):
    return self is other

@dataclass
class Graph(Node):
  ''' A graph containing `Node`s and `Edge`s.
      This is also a subclass of `Node`, to support subgraphs.
  '''
  nodes: set[Node] = field(default_factory=set)
  _nodes_by_name: dict[str, set[Node]] = field(
      default_factory=lambda: defaultdict(set)
  )
  edges: set[Edge] = field(default_factory=set)

  def __str__(self):
    return self.as_dot()

  def __hash__(self):
    return id(self)

  def __eq__(self, other):
    return self is other

  def add_node(self, node: str | Node):
    ''' Add a `Node` to the `Graph`.
        If `node` is a string, promote it to a new `Node` with that name.
        Return the `Node` (because it may be a new `Node`).
    '''
    if isinstance(node, str):
      name = node
      named_nodes = self._nodes_by_name[name]
      if named_nodes:
        # check that there's just one Node with that name
        try:
          node, = named_nodes
        except ValueError:
          raise ValueError(
              f'ambiguous {name=}: multiple Nodes with that name already exist'
          )
      else:
        # no Nodes with that name - make one
        node = Node(node)
    assert node is not None
    self.nodes.add(node)
    self._nodes_by_name[node.name].add(node)
    return node

  @typechecked
  def __getitem__(self, node_name: str) -> Node:
    ''' Return the `Node` with `.name==node_name`.
    '''
    if node_name not in self._nodes_by_name:
      raise KeyError(node_name)
    node,=self._nodes_by_name[node_name]
    return node

  def add_edge(
      self, node1: str | Node, node2: str | Node, **edge_attrs
  ) -> Edge:
    ''' Add an `Edge` to the `Graph`, return the new `Edge`.
    '''
    node1 = self.add_node(node1)  # may promote str to Node
    assert node1 is not None
    node2 = self.add_node(node2)  # may promote str to Node
    assert node2 is not None
    edge = Edge(node1, node2, **edge_attrs)
    self.edges.add(edge)
    node1.out_edges.append(edge)
    node2.in_edges.append(edge)
    return edge

  def add_chain(
      self, node1: str | Node, *more_nodes: str | Node, **edge_attrs
  ) -> list[Edge]:
    ''' Add a chain of `Node`s `Graph`, return a list of the resulting `Edge`s.
    '''
    more_nodes = list(more_nodes)
    edges = []
    while more_nodes:
      node2 = more_nodes.pop(0)
      edges.append(self.add_edge(node1, node2, **edge_attrs))
      node1 = node2
    return edges

  def as_GVGraph(
      self,
      *,
      digraph=True,
      strict=False,
      graph_attrs=None,
      node_attrs=None,
      edge_attrs=None,
  ):
    ''' Construct a `cs.gvutils.Graph` from this `Graph`.
    '''
    gvgraph = GVGraph(
        digraph=digraph,
        strict=strict,
        attrs=graph_attrs,
        node_attrs=node_attrs,
        edge_attrs=edge_attrs
    )
    gvnodes = {}
    for node in self.nodes:
      gvnode = node.as_GVNode()
      gvnodes[node] = gvnode
      gvgraph.add(gvnode)
    for edge in self.edges:
      gvgraph.join(gvnodes[edge.in_node], gvnodes[edge.out_node])
    return gvgraph

  def as_dot(
      self, *, fold=False, indent="", subindent="  ", graphtype=None, **gvkw
  ):
    ''' Construct a `cs.gvutils.Graph` and return it as a DOT string.
    '''
    gvgraph = self.as_GVGraph(**gvkw)
    return gvgraph.as_dot(
        fold=fold, indent=indent, subindent=subindent, graphtype=graphtype
    )

  def gvprint(self, **kw):
    gvpkw = {}
    attrs = {}
    node_attrs = {}
    edge_attrs = {}
    for k, v in kw.items():
      if k in ('file,fmt,layout,dataurl_encoding'):
        gvpkw[k] = v
      elif k.startswith('node_'):
        node_attrs[k.removeprefix('node_')] = v
      elif k.startswith('edge_'):
        node_attrs[k.removeprefix('edge_')] = v
      else:
        attrs[k] = v
    gvgraph = self.as_GVGraph(
        graph_attrs=attrs, node_attrs=node_attrs, edge_attrs=edge_attrs
    )
    gvgraph.print(**gvpkw)

  def partition_nodes(self) -> tuple[list, list, list]:
    ''' Partition `self.nodes` into a 3-tuple of `(roots,interior,tails)`.
        Each is a list. An isolated `Node` will appear in both `roots` and `tails`.
    '''
    roots = []
    tails = []
    interior = []
    for node in self.nodes:
      if not node.in_edges:
        roots.append(node)
      if not node.out_edges:
        tails.append(node)
      if node.in_edges and node.out_edges:
        interior.append(node)
    return roots, interior, tails

  def as_railroad(self) -> RRBase:
    ''' Return a railroad node for this `Graph`.
    '''
    ##rr_by_node = { node:node.as_railroad() for node in self.nodes}
    # A mapping of `Node` to railroad nodes containing them.
    root_nodes = []
    end_nodes = []
    counted_nodes = []
    # Collate root nodes (no in nodes), tail nodes (no out nodes)
    # and interior nodes (nodes with in and out nodes.

    for node in self.nodes:
      if not node.in_edges:
        root_nodes.append(node)
      if not node.out_edges:
        end_nodes.append(node)
      if node.in_edges and node.out_edges:
        counted_nodes.append((len(node.in_edges) * len(node.out_edges), node))
    # We start with the root nodes.
    # We construct sequences of singly connected nodes until we reach
    # a tail node or a node already mapped to a railroad node.
    rr_by_node = {}

    @typechecked
    def rr_from(root: Node) -> RRBase:
      ''' Produce an `RRSequence` encompassing the `graph` from `root`.
      '''
      rr = None
      rr_for_rrnode = {}

      def root_rr(rrnode):
        ''' Find the ancestral RR node from one of its interior RR nodes.
        '''
        while True:
          try:
            rrnode = rr_for_rrnode[rrnode]
          except KeyError:
            break
        return rrnode

      container_by_node = {}
      merge_by_node = defaultdict(trace(RRMerge))
      seq_by_node = defaultdict(RRSequence)
      seqs = []
      q = ListQueue([root], unique=True)
      for node in q:
        node0 = node
        seq = seq_by_node[node0]
        if rr is None:
          # the RR diagram starts from the first RRSequence
          rr = seq
        seqs.append(seq)
        # should we record the new sequence in an enclosing RRSplit?
        try:
          container = container_by_node.pop(node)
        except KeyError:
          if node is not root and node.in_count == 0:
            print(f'  unexpected {node.in_count=} for {node}')
        else:
          print(f'  container {container.desc} + seq {seq.desc}')
          container.append(seq)
          rr_for_rrnode[seq] = container

          # for merges, record that the sequence is enclosed in the merge
          rr_for_rrnode[seq] = container
        if node.in_count == 0:
          seq.append(RR_START)
        elif node.in_count > 1:
          # Start with a merge and queue the in_nodes.
          # Fetch the RRMerge which leads to this node.
          merge = merge_by_node[node]
          rr_for_rrnode[merge] = seq
          for lnode in node.in_nodes:
            # locate the start of the left node's linear chain (sequence)
            while lnode.in_count == 1 and lnode.in_node.out_count == 1:
              lnode = lnode.in_node
            try:
              lseq = seq_by_node[lnode]
            except KeyError:
              # no sequence for this node yet, but that's ok
              # we will include the sequence later when it's made
              assert lnode not in container_by_node
              container_by_node[lnode] = merge
              q.append(lnode)
              if lnode.name == 'b': breakpoint()
            else:
              # The origin sequence will always preexist.
              # If one of these is the origin sequence, make our sequence the origin.
              lroot = root_rr(lseq)
              # append the existing sequence right now
              merge.append(lroot)  ## was lseq
              if lroot is rr:
                rr = seq  ## was merge
              breakpoint()
          seq.append(merge)
        # append this Node's RR and all the following linear Nodes
        seq.append(node.as_railroad())
        while node.out_count == 1:
          next_node = node.out_node
          if next_node.in_count == 1:
            # continue the linear chain
            node = next_node
            seq.append(node.as_railroad())
          else:
            # the is a merge
            assert next_node.in_count > 1
            q.append(next_node)
            break

        # see why we stopped
        if node.out_count > 1:
          split = RRSplit()
          seq.append(split)
          rr_for_rrnode[split] = seq
          for out_node in node.out_nodes:
            print(f'  queue {out_node} for new Split from {node}')
            q.append(out_node)
            assert out_node not in container_by_node
            container_by_node[out_node] = split
            # connect the drawing hierarchy
            rseq = seq_by_node[out_node]
            rr_for_rrnode[out_node] = rseq
            rr_for_rrnode[rseq] = split
            rr_for_rrnode[split] = seq
        elif node.out_count == 1:
          # we must have hit a merge, queue it for consideration
          next_node = node.out_node
          assert next_node.in_count > 1
          merge = merge_by_node[next_node]
          print(f'  queue {node}.out_node:{next_node}, will be a merge')
        else:
          seq.append(RR_END)
      if container_by_node:
        print(
            f'  collating {len(container_by_node)} container_by_node entries'
        )
        for node, cont in container_by_node.items():
          seq = seq_by_node[node]
          print(
              f'    {node}->{cont.desc}: append({seq.desc=}:{seq.content[0].desc}...)'
          )
          cont.append(seq)
          if rr is seq:
            print(f'    seq is rr, make rr = container')
            rr = cont
        ##print(f'  HACK rr = last cont {cont.desc}')
        ##rr = cont
      return rr

    def rr_graph(self, root) -> RRBase:
      ''' New plan:

          Traverse the graph from a root, making a new graph where
          each node is a branch point and each edge has a refobj as
          an RR sequence or a nonbranching RR node.
          - discard cycles with a warning, for some future fixup mode
          - root nodes in the graph become an edge with a stub source node with refobj=None,
            likewise for tails

          Traverse the new simplified graph.
          A node with multiple inputs gets a left merge:
          - analyse the left graph for common ancestors?
      '''

      @typechecked
      def rr_path(path: list[Node]) -> RRBase:
        ''' Return an `RRBase` representing the `path`, a list of `Node`s.
        '''
        assert len(path) > 0
        if len(path) == 1:
          node = path[0]
          rr_nodes.add(node)
          return node.as_railroad()
        rr_nodes.update(path)
        return RRSequence([node.as_railroad() for node in path])

      @typechecked
      def rr_paths(paths: list[list[Node]]) -> RRBase:
        ''' Return an `RRBase` representing `paths`, a list of paths.
        '''
        assert len(paths) > 0
        if len(paths) == 1:
          return rr_path(paths[0])
        return RRChoice([rr_path(path) for path in paths])

      # construct a list of `RRSequence`s one for each path from each root
      rr_nodes = set()  # drawn nodes
      rrs = []  # resulting RR instances
      roots, interior, tails = self.partition_nodes()
      print("ROOTS:", roots)
      for root in roots:
        print("root", root)
        conns = root.reach_connectivity()
        # print out the adjacency mapping
        print("ORIGINAL ADJACENCY MAPPING for root", root.name)
        for src, by_dst in sorted(
            conns.items(),
            key=lambda src__by_dst: -max(max(len(path)
                                             for path in paths)
                                         for paths in src__by_dst[1].values())
        ):
          print("  src", src)
          if src in rr_nodes:
            print("    skip src, already drawn")
            continue
          for dst, paths in sorted(
              by_dst.items(),
              key=lambda dst__paths: -max(len(path) for path in dst__paths[1]),
          ):
            for path in paths:
              print(
                  "    dst", dst.name, "<- src", src.name,
                  [pnode.name for pnode in path]
              )
            if dst in rr_nodes:
              print("      skip dst, already drawn")
              continue
            optional = False
            if any(len(path) == 0 for path in paths):
              optional = True
              paths = list(filter(len, paths))
              if paths:
                rr = RRSequence(
                    [src.name, RROptional(rr_paths(paths)), dst.name]
                )
              else:
                rr = RRSequence([src.name, dst.name])
            else:
              if paths:
                rr = RRSequence([src.name, rr_paths(paths), dst.name])
              else:
                rr = RRSequence([src.name, dst.name])
            rrs.append(rr)
        breakpoint()
      print(len(rrs), "RR diagrams")
      if len(rrs) == 1:
        rr, = rrs
      else:
        rr = RRStack(rrs)
      print(rr)
      breakpoint()
      return rr

      # Prepare a connection graph where all the linear runs become
      # RRSequences and attached to a single Edge between the start
      # and end nodes.
      cnode_by_node = {}

      def cnode_for(node: Node):
        try:
          cnode = cnode_by_node[node]
        except KeyError:
          cnode = cnode_by_node[node] = Node(name=node.name, refobj=node)
        return cnode

      cgraph = Graph()
      q = ListQueue([root], unique=True)
      for node in q:
        cnode = cnode_for(node)
        for next_node in node.out_nodes:
          # gather the linear run between node and the next branching node
          seq = []
          while next_node.in_count == 1 and next_node.out_count == 1:
            seq.append(next_node)
            next_node = next_node.out_node
          next_cnode = cnode_for(next_node)
          cedge = cgraph.add_edge(cnode, next_cnode)
          cedge.objref = seq
          q.append(next_node)
      print("connection graph:")
      cgraph.gvprint(rankdir="LR")

      # there should be only 1 root node
      croot, = cgraph.partition_nodes()[0]
      cadj = connectivity(croot)
      # print out the adjacency mapping
      print("CGRAPH ADJACENCY MAPPING")
      for src, by_dst in cadj.items():
        for dst, paths in by_dst.items():
          for path in paths:
            print(src.name, dst.name, [pnode.name for pnode in path])

      # process each mapping in reverse order of the mapping
      # containing the longest path
      for src, by_dst in sorted(
          cadj.items(),
          key=lambda src__by_dst: -max(map(len, src__by_dst[1].values()))):
        print(src)
        for dst, paths in by_dst.items():
          print(" ", dst, list(map(len, paths)))

      breakpoint()

      # compose a stack of the possible paths
      for src, by_dst in cadj.items():
        assert len(by_dst) > 0
        src_node = cgraph[src.name]
        stack = []
        for dst, paths in by_dst.items():
          assert len(paths) > 0
          if len(paths) == 1:
            stack.append(paths[0])
          else:
            stack.append(paths)

    for root in root_nodes:
      rr = rr_graph(self, root)
      ##rr = trace(rr_from, retval=True)(root)
      ##pprint(rr)
      ##breakpoint()
      print(rr)
    return rr

if __name__ == '__main__':
  G = Graph("graph1")
  G.add_chain("d", "mid", "e")
  G.add_edge("d", "e")
  G.add_chain("x", "a", "y")
  G.add_chain("x", "00", "c", "y")
  G.add_edge("00", "00b")
  print(G.as_dot(fold=True))
  G.gvprint(rankdir='LR')
  rr = G.as_railroad()
  print(repr(rr))
  rr.print()