Source code for semi_cr.core.lab.station.graphing.operations

import networkx as nx
from collections.abc import Iterable
from typing import Any

from semi_cr.core.lab.station.graphing.models import Node, Edge, GraphRelation
from semi_cr.core.lab.station.graphing.parsing import normalize_channel
from semi_cr.core.lab.station.graphing.ids import instrument_node_id


# LOW-LEVEL GRAPH OPERATIONS



def build_nx_graph(
    nodes: Iterable[Node],
    edges: Iterable[Edge],
) -> nx.MultiDiGraph:
    
    graph = nx.MultiDiGraph()

    for node in nodes:
        attrs = {
            "kind": node.kind,
            "name": node.name,
            **node.attrs,
        }

        if node.obj is not None:
            attrs["obj"] = node.obj

        graph.add_node(node.id, **attrs)

    for edge in edges:
        attrs = {"kind": edge.kind,
                 **edge.attrs}

        if edge.state is not None:
            attrs["state"] = edge.state

        graph.add_edge(
            edge.source,
            edge.target,
            **attrs,
        )

    return graph




def add_connection(
    graph: nx.MultiDiGraph,
    source: str,
    target: str,
    kind: str,
    state: str = "active",
    **attrs: Any,
) -> None:
    """
    Adds a directed physical edge.

    Direction convention:
        electronics -> matrix -> fischer -> converter -> filter -> sample -> device

    Therefore, reverse BFS from a device terminal walks outward toward electronics.
    """
    if source not in graph: 
        graph.add_node(source, kind="unknown", name=source)
    if target not in graph:
        graph.add_node(target, kind="unknown", name=target)

    graph.add_edge(
        source,
        target,
        kind=kind,
        state=state,
        **attrs,
    )




def ensure_instrument_terminal(
    graph: nx.MultiDiGraph,
    instrument_name: str,
    terminal_name: str,
    node_id: str,
) -> None:

    instrument_id = instrument_node_id(
        instrument_name
    )

    if instrument_id not in graph:
        graph.add_node(
            instrument_id,
            kind="instrument",
            config_name=instrument_name,
        )

    if node_id not in graph:
        graph.add_node(
            node_id,
            kind="instrument_terminal",
            instrument_name=instrument_name,
            terminal_name=normalize_channel(terminal_name),
        )

    if not graph.has_edge(
        instrument_id,
        node_id,
    ):
        graph.add_edge(
            instrument_id,
            node_id,
            kind=GraphRelation.CONTAINS,
        )