Source code for Karana.KUtils.MultibodyTUI.graph

"""Classes for working with directed graphs."""

from abc import ABC, abstractmethod
from typing import TypeVar, Generic, cast

from Karana.Dynamics import PhysicalBody, Multibody
from Karana.Frame import Frame, FrameContainer

T = TypeVar("T")




class GraphAdapter(ABC, Generic[T]):
    """Common interface for various types of directed graphs.

    The purpose of this class is to define a common interface for
    wrapper classes that adapt different graphs, each with their own
    interface, to a single common interface.
    """



    @abstractmethod
    def parents(self, node: T, /) -> list[T]:
        """Get the list of parents for a given node."""




    @abstractmethod
    def children(self, node: T, /) -> list[T]:
        """Get the list of children for a given node."""




    @abstractmethod
    def nodes(self) -> list[T]:
        """Get the list of all nodes in the graph."""




    def descendants(self, node: T, /) -> list[T]:
        """Get the list of all descendants for a given node."""
        # Use a dictionary instead of set to preserve ordering
        visited = {}
        stack = []
        stack.append(node)
        results = []
        # do a dfs traversal
        while stack:
            curr = stack.pop()
            for child in self.children(curr):
                if child not in visited:
                    visited[child] = None
                    stack.append(child)
        return list(visited)




    def ancestors(self, node: T, /) -> list[T]:
        """Get the list of all ancestors for a given node."""
        # Use a dictionary instead of set to preserve ordering
        visited = {}
        stack = []
        stack.append(node)
        results = []
        # do a dfs traversal
        while stack:
            curr = stack.pop()
            for parent in self.parents(curr):
                if parent not in visited:
                    visited[parent] = None
                    stack.append(parent)
        return list(visited)






class BodyGraphAdapter(GraphAdapter[PhysicalBody]):
    """Multibody wrapper implementing the GraphAdapter interface."""

    def __init__(self, multibody: Multibody):
        """Create an instance of BodyGraphAdapter.

        Parameters
        ----------
        multibody : Multibody
            The Multibody to create the GraphAdapter interface for.
        """
        self.multibody = multibody



    def parents(self, body: PhysicalBody) -> list[PhysicalBody]:
        """Get the parents of a body.

        Parameters
        ----------
        body : PhysicalBody
            The body to get the parents for.

        Returns
        -------
        list[PhysicalBody]
            The parents of the given body.
        """
        if body == self.multibody.virtualRoot():
            return []
        parent = body.physicalParentBody()
        return [parent]




    def children(self, body: PhysicalBody) -> list[PhysicalBody]:
        """Get the children of a body.

        Parameters
        ----------
        body : PhysicalBody
            The body to get the children of.

        Returns
        -------
        list[PhysicalBody]
            The children of a the given body.
        """
        # return body.childBodies()
        return body.multibody().childrenBodies(body)




    def nodes(self) -> list[PhysicalBody]:
        """Get all the bodies in the multibody (these are all the nodes of the graph).

        Returns
        -------
        list[PhysicalBody]
            All the bodies in the Multibody.
        """
        return [
            cast(PhysicalBody, self.multibody.virtualRoot())
        ] + self.multibody.sortedPhysicalBodiesList()






class FrameGraphAdapter(GraphAdapter[Frame]):
    """FrameContainer wrapper implementing the GraphAdapter interface."""

    def __init__(self, frame_container: FrameContainer):
        """Create an instance of FrameGraphAdapter.

        Parameters
        ----------
        frame_container : FrameContainer
            The frame container to create the GraphAdapter interface for.
        """
        self.frame_container = frame_container



    def parents(self, frame: Frame) -> list[Frame]:
        """Get the parents of a Frame.

        Parameters
        ----------
        frame : Frame
            The Frame to get the parents of.

        Returns
        -------
        list[Frame]
            The parent frame of the given frame.
        """
        parent = frame.parent()
        if parent is None:
            return []
        return [parent]




    def children(self, frame: Frame) -> list[Frame]:
        """Get the children of a Frame.

        Parameters
        ----------
        frame : Frame
            The Frame to get the children of.

        Returns
        -------
        list[Frame]
            The children of the provided Frame.
        """
        # Not ideal, but there doesn't seem to be a way to get the
        # children directly
        children = []
        for candidate in self.frame_container.frames():
            if candidate.parent() == frame:
                children.append(candidate)
        return children




    def nodes(self) -> list[Frame]:
        """Get all Frames of the FrameContainer.

        Returns
        -------
        list[Frame]
            All Frames of the FrameContainer.
        """
        return self.frame_container.frames()