mantis/execution/schema_solve.py

from ..utilities.utilities import (prRed, prGreen, prPurple, prWhite,
                              prOrange,
                              wrapRed, wrapGreen, wrapPurple, wrapWhite,
                              wrapOrange,)
from ..utilities.utilities import init_connections, init_dependencies, get_link_in_out
from ..base_definitions import (SchemaUINode, custom_props_types, \
    MantisNodeGroup, SchemaGroup, replace_types, GraphError, links_sort_key)
from ..utilities.node_common import setup_custom_props_from_np
# a class that solves Schema nodes
from bpy.types import NodeGroupInput, NodeGroupOutput, NodeSocket

from .readtree import execution_error_cleanup

class SchemaSolver:
    __slots__ = (
        "all_nodes", "node", "solved", "tree",
        "uuid", "error_popups", "signature",
        "schema_nodes", "solved_nodes",
        "incoming_connections", "outgoing_connections",
        "constant_in", "constant_out",
        "array_input_connections", "array_output_connections",
        "nested_schemas", "autogenerated_nodes",
        "held_links", "tree_path_names",
        "autogen_path_names", "is_node_group",
        "index_link", "solve_length", "index",
        "interfaces", "index_str", "prev_index_str",
        "current_interface", "next_interface"
    )
    def __init__(self, schema_dummy, nodes, ui_node, signature=None, error_popups=False):
        self.all_nodes = nodes # this is the parsed tree from Mantis
        self.node = schema_dummy
        self.node.solver = self
        self.solved = False
        self.tree = ui_node.node_tree
        self.uuid = self.node.uuid
        self.error_popups = error_popups
        if signature:
            self.signature = signature
        else:
            self.signature = self.node.signature
        self.schema_nodes={}; self.solved_nodes = {}
        self.incoming_connections = {}; self.outgoing_connections = {}
        self.constant_in = {}; self.constant_out = {}
        self.array_input_connections = {}; self.array_output_connections = {}
        self.nested_schemas = {}; self.autogenerated_nodes = {}
        self.held_links = []
        self.tree_path_names  = [*self.node.signature] # same tree as the schema node
        self.autogen_path_names = ['SCHEMA_AUTOGENERATED', *self.node.signature[1:]]
        self.is_node_group = False
        if self.node.ui_node.bl_idname == "MantisNodeGroup":
            self.is_node_group = True
        if self.node.inputs['Schema Length'].links:
            self.index_link = self.node.inputs['Schema Length'].links[0]
        else:
            self.index_link = None
        # this should never happen, but it's OK to check.
        if self.node.inputs["Schema Length"].is_linked:
            if (other := self.node.inputs['Schema Length'].links[0].from_node).prepared == False:
                raise RuntimeError(f"Schema Length cannot be determined for {self.node} because {other} is not prepared.")
        self.solve_length = self.node.evaluate_input("Schema Length")
        # I'm making this a property of the solver because the solver's data is modified as it solves each iteration
        self.index = 0
        self.interfaces = []

        prWhite(f"\nExpanding schema {self.tree.name} in node {self.node.signature}"
                 f" with length {self.solve_length}.")

        self.init_schema_links()
        self.generate_interface_nodes()
        self.set_schema_per_interation_constants()

        # Sort the multi-input nodes in reverse order of ID, this ensures that they are
        #   read in the order they were created
        for inp in self.node.inputs.values():
            inp.links.sort(key=links_sort_key)

        for ui_node in self.tree.nodes:
            # Make the mantis nodes for the special interface types
            # we use the bl_idname because all schema nodes should be single-instance
            signature = (*self.tree_path_names, ui_node.bl_idname)
            if isinstance(ui_node, (SchemaUINode, NodeGroupInput, NodeGroupOutput)):
                # We use the schema node's "natural signature" here because it represents
                # the "original" signature of the schema UI group node since this schema
                # solver may be in a nested schema, and its node's signature may have
                # uuid/index attached.
                get_sig = (*self.node.ui_signature, ui_node.bl_idname)
                if not (mantis_node := self.all_nodes.get(get_sig)):
                    raise RuntimeError(wrapRed(f"Not found: {get_sig}"))
                self.schema_nodes[signature] = mantis_node
            # HACK to make Group Nodes work
            if ui_node.bl_idname == "NodeGroupInput":
                from ..schema_nodes import SchemaConstInput
                mantis_node = SchemaConstInput(signature=signature, base_tree=self.node.base_tree, parent_schema_node=self.node)
                self.schema_nodes[signature] = mantis_node
            if ui_node.bl_idname == "NodeGroupOutput":
                from ..schema_nodes import SchemaConstOutput
                mantis_node = SchemaConstOutput(signature=signature, base_tree=self.node.base_tree, parent_schema_node=self.node)
                self.schema_nodes[signature] = mantis_node

    def init_schema_links(self,):
        """ Sort and store the links to/from the Schema group node."""
        for item in self.tree.interface.items_tree:
            if item.item_type == 'PANEL': continue
            from ..utilities.utilities import read_schema_type
            parent_name = read_schema_type(item)
            # just gonna try and make the 
            match parent_name:
                case 'Connection':
                    if item.in_out == 'INPUT':
                        if incoming_links := self.node.inputs[item.identifier].links:
                            self.incoming_connections[item.name] = incoming_links[0]
                        else:
                            self.incoming_connections[item.name] = None
                    else: # OUTPUT
                        if outgoing_links := self.node.outputs[item.identifier].links:
                            self.outgoing_connections[item.name] = outgoing_links.copy()
                        else:
                            self.outgoing_connections[item.name] = []
                case 'Constant':
                    if item.in_out == 'INPUT':
                        if constant_in_links := self.node.inputs[item.identifier].links:
                            self.constant_in[item.name] = constant_in_links[0]
                        else:
                            self.constant_in[item.name] = None
                    else: # OUTPUT
                        if constant_out_links := self.node.outputs[item.identifier].links:
                            self.constant_out[item.name] = constant_out_links.copy()
                        else:
                            self.constant_out[item.name] = []
                case 'Array':
                    if item.in_out == 'INPUT':
                        if item.identifier not in self.array_input_connections.keys():
                            self.array_input_connections[item.identifier]=[]
                        if in_links := self.node.inputs[item.identifier].links:
                            self.array_input_connections[item.identifier]=in_links.copy()
                            # I am tempted to put a check here, but it is sufficient to
                            # rely on hierarchy links ensuring the arrays are prepared.
                            # and testing here for un-prepared arrays will mess up schema
                            # relationships in non-hierarchy situations.
                            # Still, I'd like to have an easier time catching these problems.
                    else: # OUTPUT
                        if item.identifier not in self.array_output_connections.keys():
                            self.array_output_connections[item.identifier]=[]
                        if out_links := self.node.outputs[item.identifier].links:
                            self.array_output_connections[item.identifier] = out_links.copy()

    # Using lambdas here causes a lot of function calls, should be optimized out
    # since after all this only needs to change when the index updates
    def set_schema_per_interation_constants(self):
        # NOTE: I removed the UUID that has always been here
        #  I can't figure out why I needed it and as far as I can tell there
        #  really is no way for this to have name clashes. Tests work OK, too.
        self.index_str = '.'+str(self.index).zfill(4)
        self.prev_index_str = '.'+str(self.index-1).zfill(4)
        self.current_interface = self.interfaces[self.index]
        self.next_interface = self.interfaces[self.index + 1] # this is guaranteed
        if self.is_node_group:
            self.index_str = ''; self.prev_index_str = ''

    def generate_interface_nodes(self):
        from ..internal_nodes import GroupInterface
        # One input interface per iteration, and one output interface
        items_tree = self.tree.interface.items_tree
        mContext=self.node.mContext
        for i in range(self.solve_length): # one input per
            index_str = '.'+str(i).zfill(4)
            name = 'InputInterface'+index_str
            if i == 0: in_out = 'SCHEMA_INPUT'
            else:      in_out = 'SCHEMA_BETWEEN'
            signature = (*self.autogen_path_names, name)
            if not (interface := self.all_nodes.get( signature)):
                interface = GroupInterface(signature=signature, base_tree=self.tree, in_out=in_out, 
                            ui_node=None, items_tree=items_tree)
                interface.mContext = mContext
            # strip the AUTOGEN/SCHEMA tag if it is there
            mContext.interface_nodes[( None, *interface.signature[1:]) ] = interface
            interface.generated_at_schema_index = self.node.generated_at_schema_index + [i]
            # add the index and length socket constants
            interface.parameters['Index'] = i
            if i == 0: # TODO I am duplicating Length for the interface feature, this is not necessary
                interface.parameters['Length'] = self.solve_length
            # the GroupInterface will initialize with the right sockets, but we still have to link them

            if in_out == 'SCHEMA_INPUT': # Take the links into the schema and plug them here
                for item in items_tree:
                    if item.item_type == 'PANEL': continue
                    if (mantis_link := self.constant_in.get(item.name)):
                        mantis_link.from_node.outputs[mantis_link.from_socket].connect(interface, item.identifier)
                    # for mantis_link in self.array_input_connections.get(item.name, []):
                    #     mantis_link.from_node.outputs[mantis_link.from_socket].connect(interface, item.identifier)
                    if (mantis_link := self.incoming_connections.get(item.name)):
                        mantis_link.from_node.outputs[mantis_link.from_socket].connect(interface, item.identifier)
            self.interfaces.append(interface)
            # solved nodes will get added to the outer all_nodes later
            self.solved_nodes[signature]=interface

        # now do the output
        signature = (*self.autogen_path_names, "OutputInterface" + index_str)
        interface = GroupInterface(signature=signature, base_tree=self.tree,
                    in_out='OUTPUT', ui_node=None, items_tree=items_tree)
        self.interfaces.append(interface)
        interface.generated_at_schema_index = i
        self.solved_nodes[signature]=interface
        interface.mContext = mContext
        # strip the AUTOGEN/SCHEMA tag if it is there
        mContext.interface_nodes[( None, *interface.signature[1:]) ] = interface

    def is_node_deeper_nested(self, queried_node, compare_node):
        return len(compare_node.signature) < len(queried_node.signature)

    def gen_solve_iteration_mantis_nodes(self, frame_mantis_nodes, unprepared):
        for ui_node in self.tree.nodes:
            mantis_node_name = ui_node.name
            index_str = self.index_str
            mContext=self.node.mContext
            if isinstance(ui_node, SchemaUINode):
                continue  # IGNORE the schema interface nodes, we already made them in __init__()
                # they are reused for each iteration.
            elif ui_node.bl_idname in ['NodeFrame', 'NodeReroute']:
                continue  # IGNORE stuff that is purely UI - frames, reroutes.
            elif ui_node.bl_idname in ['NodeGroupInput', 'NodeGroupOutput']:
                continue # we converted these to Schema Nodes because they represent a Group input.
            signature = (*self.autogen_path_names, mantis_node_name+index_str)
            ui_signature=(*self.signature, mantis_node_name)
            prototype_mantis_node = self.all_nodes[ui_signature]
            # the prototype_mantis_node was generated inside the schema when we parsed the tree.
            # it is the ui_node of the mantis node which we make for this iteration
            # for Schema sub-nodes ... they need a ui_node to init.
            if prototype_mantis_node.node_type in ['DUMMY', 'DUMMY_SCHEMA']:
                # We stored the ui_node ui_node when creating the Mantis node.
                ui_node = prototype_mantis_node.ui_node
                # if prototype_mantis_node is a group or schema: TODO changes are needed elsewhere to make this easier to read. LEGIBILITY
                if ui_node.bl_idname in ["MantisNodeGroup", "MantisSchemaGroup"]:
                    mantis_node = prototype_mantis_node.__class__(
                        signature, prototype_mantis_node.base_tree, ui_node=ui_node,
                        ui_signature =  prototype_mantis_node.signature)
                    # now let's copy the links from the ui_node node
                    if ui_node.bl_idname in ["MantisNodeGroup"]:
                        mantis_node.prepared = False
                        mantis_node.node_type = 'DUMMY_SCHEMA' # we promote it to a schema for now
                        from ..base_definitions import NodeSocket
                        mantis_node.inputs['Schema Length'] = NodeSocket(is_input=True, node=mantis_node, name='Schema Length')
                        mantis_node.parameters['Schema Length'] = 1 # set the length to 1 since it is a single group instance
                        # we'll make the autogenerated nodes for constant inputs. It doesn't matter that there is technically
                        #  a ui_node available for each one  -- these are cheap and I want this to be easy.
                        from .readtree import make_connections_to_ng_dummy
                        make_connections_to_ng_dummy(self.node.base_tree, self.autogen_path_names, frame_mantis_nodes, self.all_nodes, mantis_node)
                else:
                    mantis_node = prototype_mantis_node.__class__(signature, prototype_mantis_node.base_tree, ui_node=ui_node)
            else:
                mantis_node = prototype_mantis_node.__class__(signature, prototype_mantis_node.base_tree)
            frame_mantis_nodes[mantis_node.signature] = mantis_node
            self.all_nodes[mantis_node.signature] = mantis_node
            mantis_node.ui_signature=ui_signature # set the natural signature to ensure we can access from the UI
            if mantis_node.prepared == False:
                unprepared.append(mantis_node)
            if mantis_node.__class__.__name__ in custom_props_types:
                setup_custom_props_from_np(mantis_node, ui_node)
            # be sure to pass on the Mantis Context to them.
            mantis_node.mContext=mContext
            mantis_node.generated_at_schema_index = self.node.generated_at_schema_index + [self.index]
            mantis_node.fill_parameters(ui_node)
            assert mantis_node.mContext is not None

    def handle_link_from_index_input(self, index, frame_mantis_nodes, ui_link):
        from ..base_definitions import can_remove_socket_for_autogen
        _from_name, to_name = get_link_in_out(ui_link)
        to_node = frame_mantis_nodes[ (*self.autogen_path_names, to_name+self.index_str) ]

        if (not can_remove_socket_for_autogen(to_node, ui_link.to_socket.name)) or \
                                    to_node.node_type in ['DUMMY', 'DUMMY_SCHEMA']:
            from_node = self.current_interface
            connection = from_node.outputs[ui_link.from_socket.name].connect(node=to_node, socket=ui_link.to_socket.identifier)
            connection.is_hierarchy = False # For now this is the easiest way to do this
            return
        # Since the index is already determined, it is safe to remove the socket and just keep the value.
        to_node.parameters[ui_link.to_socket.name] = index
        del to_node.inputs[ui_link.to_socket.name]

    def handle_link_from_schema_length_input(self, frame_mantis_nodes, ui_link):
        # see, here I can just use the schema node
        _from_name, to_name = get_link_in_out(ui_link)
        if to_name in replace_types:
            to_node = self.schema_nodes[(*self.autogen_path_names, to_name)]
        else:
            to_node = frame_mantis_nodes[(*self.autogen_path_names, to_name+self.index_str)]
        # this self.index_link is only used here?
        if self.index_link is None:
            # this should be impossible because the Schema gets an auto-generated Int input.
            raise NotImplementedError(" 241 This code should be unreachable. Please report this as a bug!")
        if (self.index_link.from_node):
            connection = self.index_link.from_node.outputs[self.index_link.from_socket].connect(node=to_node, socket=ui_link.to_socket.name)
        # otherwise we can autogen an int input I guess...?
        else:
            raise RuntimeError("247 I This code should be unreachable. Please report this as a bug!")

    def handle_link_from_incoming_connection_input(self, frame_mantis_nodes, ui_link):
        incoming = self.incoming_connections[ui_link.from_socket.name]
        if incoming is not None:
            from_node = incoming.from_node
            _from_name, to_name = get_link_in_out(ui_link)
            to_node = frame_mantis_nodes[ (*self.autogen_path_names, to_name+self.index_str) ]
            socket_name=ui_link.to_socket.name
            if to_node.node_type in [ 'DUMMY_SCHEMA' ]:
                socket_name=ui_link.to_socket.identifier
            connection = from_node.outputs[incoming.from_socket].connect(node=to_node, socket=socket_name)
            init_connections(from_node)

    def handle_link_to_outgoing_connection_output(self, frame_mantis_nodes, ui_link,):
        mantis_incoming_node = self.schema_nodes[*self.tree_path_names,  'SchemaIncomingConnection']
        for mantis_link in mantis_incoming_node.outputs[ui_link.to_socket.name].links:
            to_mantis_node, to_socket_name = mantis_link.to_node, mantis_link.to_socket
            from_name = get_link_in_out(ui_link)[0]
            from_mantis_node = self.solved_nodes[ (*self.autogen_path_names, from_name+self.prev_index_str) ]

            to_mantis_node_signature = ( *self.autogen_path_names,
                            to_mantis_node.signature[-1] + self.index_str )

            # we need to detect if the next node is in a group.
            # REMEMBER: at this point, nested groups haven't been solved yet.
            if to_mantis_node_signature not in frame_mantis_nodes.keys() and \
                    self.is_node_deeper_nested(to_mantis_node, from_mantis_node):
                to_mantis_node = frame_mantis_nodes[ (
                    *self.autogen_path_names, # the SCHEMA_AUTOGENERATED string
                    to_mantis_node.signature[-2]+ self.index_str ) ]
                # this connection was forbidden before, right? so this should be a safe assumption.
                assert to_mantis_node.node_type == 'DUMMY_SCHEMA', "I expected this to be a group/schema"
                # we need to get the identifier of the named socket now.
                to_socket_name = to_mantis_node.ui_node.inputs[to_socket_name].identifier
            else:
                to_mantis_node = frame_mantis_nodes[ to_mantis_node_signature ] # add the index string

            from_socket_name = ui_link.from_socket.name
            if from_mantis_node.node_type in ['DUMMY_SCHEMA']:
                from_socket_name = ui_link.from_socket.identifier
            connection = from_mantis_node.outputs[from_socket_name].connect(node=to_mantis_node, socket=to_socket_name)
            # We want to delete the links from the tree into the schema node.
            # TODO: this is not robust enough and I do not feel sure this is doing the right thing.
            if existing_link := self.incoming_connections[ui_link.to_socket.name]:
                if existing_link.to_node == self.node:
                    print ("INFO: Deleting...", existing_link)
                    if self.node.signature[-1] in existing_link.to_node.signature:
                        existing_link.die()
            # BUG may exist here.
            self.incoming_connections[ui_link.to_socket.name] = connection


    def handle_link_from_constant_input(self, frame_mantis_nodes, ui_link, to_ui_node):
        incoming = self.constant_in[ui_link.from_socket.name]
        if incoming is None:
            raise GraphError (f"{self.node} is missing a required input: {ui_link.from_socket.name}")
        from_node = incoming.from_node
        to_name = get_link_in_out(ui_link)[1]
        to_node = frame_mantis_nodes[(*self.autogen_path_names, to_name+self.index_str)]
        to_socket=ui_link.to_socket.name
        from ..base_definitions import MantisNodeGroup, SchemaGroup
        if isinstance(to_ui_node, (SchemaGroup, MantisNodeGroup)):
            to_socket=ui_link.to_socket.identifier
        connection = from_node.outputs[incoming.from_socket].connect(node=to_node, socket=to_socket)
        init_connections(from_node)

    def handle_link_from_array_input_get(self, frame_mantis_nodes, ui_link ):
        from_ui_node = ui_link.from_socket.node
        from_node = self.schema_nodes[(*self.node.ui_signature, from_ui_node.bl_idname)]
        from collections import deque
        unprepared = deque(from_node.hierarchy_dependencies)
        self.prepare_nodes(unprepared)
        from ..utilities.utilities import cap, wrap
        get_index = from_node.evaluate_input("Index", self.index) # get the most recent link
        # getting the link at self.index just saves the trouble of killing the old links
        # that are left because this node is reused in each iteration of the schema
        assert(get_index is not None), f"Cannot get index in {from_node}"
        oob = from_node.evaluate_input("OoB Behaviour")
        array_length = len(self.array_input_connections[ui_link.from_socket.identifier])-1
        if oob == 'WRAP':
            get_index = wrap(0, array_length+1, get_index)
        if oob == 'HOLD':
            get_index = cap(get_index, array_length)
        try:
            incoming = self.array_input_connections[ui_link.from_socket.identifier][get_index]
        except IndexError:
            raise RuntimeError(f"The array index {get_index} is out of bounds. Size: "
                               f"{len(self.array_input_connections[ui_link.from_socket.identifier])}")
        to_name = get_link_in_out(ui_link)[1]
        to_node = frame_mantis_nodes[(*self.autogen_path_names, to_name+self.index_str)]
        from_socket = incoming.from_node.outputs[incoming.from_socket]
        connection = from_socket.connect(to_node, ui_link.to_socket.name)
        # TODO: kill links  to the array get that are no longer needed
        # right now I inefficiently waste cpu-time solving extra nodes
        # because I took the lazy way out


    def handle_link_to_array_input_get(self, frame_mantis_nodes, ui_link):
        from_name, to_name = get_link_in_out(ui_link)
        from_mantis_node = frame_mantis_nodes[(*self.autogen_path_names, from_name+self.index_str)]
        to_mantis_node = self.schema_nodes[(*self.tree_path_names, to_name)]
        # this only needs to be done once:
        if self.index == 0: # BUG? HACK? TODO find out what is going on here.
            # Kill the link between the schema node group and the node connecting to it
            old_mantis_node = self.all_nodes[(*self.tree_path_names, from_name)]
            # I am not sure about this!
            existing_link = old_mantis_node.outputs[ui_link.from_socket.name].links[0]
            existing_link.die()
        #
        connection = from_mantis_node.outputs[ui_link.from_socket.name].connect(node=to_mantis_node, socket=ui_link.to_socket.name)
        connection.is_hierarchy = True # we have to mark this because links to Schema are not usually hierarchy (?)
        to_mantis_node.hierarchy_dependencies.append(from_mantis_node); from_mantis_node.hierarchy_connections.append(to_mantis_node)
        # TODO: review the wisdom of this default.

    def handle_link_from_array_input(self, frame_mantis_nodes, ui_link, index):
        get_index = index
        try:
            incoming = self.array_input_connections[ui_link.from_socket.identifier][get_index]
        except IndexError:
            if len(self.array_input_connections[ui_link.from_socket.identifier]) > 0:
                incoming = self.array_input_connections[ui_link.from_socket.identifier][0]
                # prOrange(incoming.from_node.node_type)
                if incoming.from_node.node_type not in ['DUMMY_SCHEMA']:
                    raise NotImplementedError(wrapRed("dev: make it so Mantis checks if there are enough Array inputs."))
                else: # do nothing
                    return
            else:
                raise RuntimeError(wrapRed("make it so Mantis checks if there are enough Array inputs!"))
        to_name = get_link_in_out(ui_link)[1]
        to_node = frame_mantis_nodes[(*self.autogen_path_names, to_name+self.index_str)]
        connection = incoming.from_node.outputs[incoming.from_socket].connect(node=to_node, socket=ui_link.to_socket.name)
        init_connections(incoming.from_node)

    def handle_link_from_array_input_all(self, frame_mantis_nodes, ui_link):
        all_links = self.array_input_connections[ui_link.from_socket.identifier]
        to_name = get_link_in_out(ui_link)[1]
        to_node = frame_mantis_nodes[(*self.autogen_path_names, to_name+self.index_str)]
        # connection = incoming.from_node.outputs[incoming.from_socket].connect(node=to_node, socket=ui_link.to_socket.name)
        for l in all_links:
            # we need to copy the link with the new from-node info
            from ..base_definitions import NodeLink
            to_socket_name=ui_link.to_socket.name
            if to_node.node_type in ['DUMMY_SCHEMA']:
                to_socket_name=ui_link.to_socket.identifier
            connection = NodeLink(l.from_node, l.from_socket, to_node, to_socket_name,
                                  l.multi_input_sort_id, l.sub_sort_id)
            to_node.flush_links()

    def handle_link_to_constant_output(self, frame_mantis_nodes, index, ui_link,  to_ui_node):
        to_node = self.schema_nodes[(*self.tree_path_names, to_ui_node.bl_idname)]
        expose_when = to_node.evaluate_input('Expose when N==')
        # this will be None for nested Groups that have ordinary Group Input and Group Output nodes.
        # and in that case we just do it no matter what
        if expose_when is not None and expose_when > self.solve_length-1:
            raise GraphError(f"The Schema has a constant output at index {expose_when}"
                             f" but it terminates at index {self.solve_length-1}. "
                             f"The Schema must have a length of at least {expose_when+1}.")
        from_name = get_link_in_out(ui_link)[0]
        #use it directly if it is a mantis node; this happens when the previous node was a Schema
        if isinstance( ui_link.from_socket, NodeSocket): # normal
            from_node = frame_mantis_nodes[(*self.autogen_path_names, from_name+self.index_str) ]
        else: # it's a mantis socket, set manually while looping though prepped links
            from_node = ui_link.from_socket.node
        from_socket_name = ui_link.from_socket.name
        if from_node.node_type == 'DUMMY_SCHEMA':
            if isinstance( ui_link.from_socket, NodeSocket): # normal
                from_socket_name = ui_link.from_socket.identifier
            else:
                from_socket_name = ui_link.from_socket.name
        # HACK here to force it to work with ordinary node groups, which don't seem to set this value correctly.
        if to_ui_node.bl_idname == "NodeGroupOutput":
            expose_when = index # the expose-when value is None because the GroupOutput doesn't have the socket.
        # end HACK
        if index == expose_when:
            for outgoing in self.constant_out[ui_link.to_socket.name]:
                to_node = outgoing.to_node
                connection = from_node.outputs[from_socket_name].connect(node=to_node, socket=outgoing.to_socket)

    def handle_link_from_subschema_to_output(self, frame_mantis_nodes, ui_link, to_ui_node):
        # wire the schema node itself to the output so it will push the connections when solving
        to_node = self.schema_nodes[(*self.tree_path_names, to_ui_node.bl_idname)]
        from_name = get_link_in_out(ui_link)[0]
        from_node = frame_mantis_nodes[ (*self.autogen_path_names, from_name+self.index_str) ]
        connection = from_node.outputs[ui_link.from_socket.identifier].connect(node=to_node, socket=ui_link.to_socket.name)

    def handle_link_from_array_type_to_array_out(self, original_ui_link, dummy_link):
        # this is so annoyingly specific lol
        from_node = dummy_link.from_mantis_node; from_socket_name=dummy_link.from_socket.name
        for outgoing in self.array_output_connections[original_ui_link.to_socket.identifier]:
            to_node = outgoing.to_node
            connection = from_node.outputs[from_socket_name].connect(node=to_node, socket=outgoing.to_socket)

    # WTF is even happening here?? TODO BUG HACK
    def handle_link_to_array_output(self, frame_mantis_nodes, index, ui_link, to_ui_node, from_ui_node):# if this duplicated code works, dedupe!
        to_node = self.schema_nodes[(*self.tree_path_names, to_ui_node.bl_idname)] # get it by [], we want a KeyError if this fails
        from_name = get_link_in_out(ui_link)[0]
        #use it directly if it is a mantis node; this happens when the previous node was a Schema
        if isinstance( ui_link.from_socket, NodeSocket): # normal
            from_node = frame_mantis_nodes[(*self.autogen_path_names, from_name+self.index_str) ]
        else: # it's a mantis socket, set manually while looping though prepped links
            from_node = ui_link.from_socket.node
        from_socket_name = ui_link.from_socket.name
        for outgoing in self.array_output_connections[ui_link.to_socket.identifier]:
            if not from_node:
                from_node = self.schema_nodes[(*self.tree_path_names, from_ui_node.bl_idname)]
            to_node = outgoing.to_node
            from ..schema_nodes import SchemaIndex
            if isinstance(from_node, SchemaIndex):
                signature = ("MANTIS_AUTOGENERATED", *self.tree_path_names[1:-1], self.index_str,
                             ui_link.from_socket.name, ui_link.from_socket.identifier)
                from_node = self.all_nodes.get(signature)
                if not from_node:
                    from .readtree import autogen_node
                    from_node = autogen_node(self.node.base_tree, ui_link.from_socket,
                                            signature, self.node.mContext)
                    if ui_link.from_socket.name in ['Index']:
                        from_node.parameters = {ui_link.from_socket.name:index}
                    else:
                        from_node.parameters = {ui_link.from_socket.name:self.solve_length}
                    frame_mantis_nodes[signature]=from_node; self.solved_nodes[signature]=from_node
                    self.all_nodes[signature]=from_node
            elif from_node.node_type == 'DUMMY_SCHEMA':
                if isinstance( ui_link.from_socket, NodeSocket): # normal
                    from_socket_name = ui_link.from_socket.identifier
                else:
                    from_socket_name = ui_link.from_socket.name

            if to_node.node_type == 'DUMMY_SCHEMA' and to_node.prepared:
                other_stem = ('SCHEMA_AUTOGENERATED', *to_node.signature[1:])
                other_schema_np = to_node.ui_node
                other_schema_tree = other_schema_np.node_tree
                for n in other_schema_tree.nodes:
                    if n.bl_idname not in ["SchemaArrayInput", "SchemaArrayInputGet", "SchemaArrayInputAll"]:
                        continue
                    out = n.outputs[outgoing.to_socket]
                    for l in out.links:
                        other_index_str = lambda : '.'+str(to_node.uuid)+'.'+str(index).zfill(4)
                        other_sig = (*other_stem, l.to_node.name+other_index_str())
                        out_node = self.all_nodes.get(other_sig)
                        if not out_node: # if it is a sub-schema, the other node is in the parent
                            out_node = to_node.solver.all_nodes.get(other_sig)
                        if not out_node:
                            raise KeyError(f"Cannot find node:  {other_sig}")
                        connection = from_node.outputs[from_socket_name].connect(node=out_node, socket=l.to_socket.name)
            else:
                connection = from_node.outputs[from_socket_name].connect(node=to_node, socket=outgoing.to_socket)


    def spoof_link_for_subschema_to_output_edge_case(self, ui_link, ):
        to_ui_node = ui_link.to_socket.node
        # ui_link is the link between the Schema & the output but we have mantis_links
        # connected up correctly. between the output and the schema's generated nodes.
        # so seek BACK from the output node and grab the from-node that is connected to
        #  the link. then modify the ui_link to point to that node.
        from ..base_definitions import DummyLink
        if not isinstance(ui_link, DummyLink): # make it a Dummy so i can modify it
            ui_link = DummyLink(ui_link.from_socket, ui_link.to_socket,
                                multi_input_sort_id=ui_link.multi_input_sort_id)
        to_node = self.schema_nodes[(*self.tree_path_names, to_ui_node.bl_idname)]
        for inp in to_node.inputs:
            if inp.name == ui_link.to_socket.name: # the most recent one is from
                l = inp.links[-1]; break # <----  this index of the outer schema
        ui_link.from_mantis_node = l.from_node; ui_link.from_socket = l.from_node.outputs[l.from_socket]
        return ui_link

    def test_is_sub_schema(self, other):
        for i in range(len(other.ui_signature)-1): # -1, we don't want to check this node, obviously
            if self.node.ui_signature[:i+1]:
                return False
        return True

    def prepare_nodes(self, unprepared):
        # At this point, we've already run a pretty exhaustive preperation phase to prep the schema's dependencies
        # So we should not need to add any new dependencies unless there is a bug elsewhere.
        # and in fact, I could skip this in some cases, and should investigate if profiling reveals a slowdown here.
        forbidden=set()
        e = None
        # forbid some nodes - they aren't necessary to solve the schema & cause problems.
        while unprepared:
            mantis_node = unprepared.pop()
            if mantis_node.node_type == 'DUMMY_SCHEMA' and not self.test_is_sub_schema(mantis_node):
                forbidden.add(mantis_node) # do NOT add this as a dependency.
            if mantis_node in forbidden: continue # trying to resolve dependencies for.
            if sum([dep.prepared for dep in mantis_node.hierarchy_dependencies]) == len(mantis_node.hierarchy_dependencies):
                try:
                    mantis_node.bPrepare()
                    if mantis_node.node_type == 'DUMMY_SCHEMA':
                        self.solve_nested_schema(mantis_node)
                except Exception as e:
                    raise execution_error_cleanup(mantis_node, e, show_error = False)
                    break
                if not mantis_node.prepared:
                    raise RuntimeError( f"{mantis_node} has failed to prepare."
                             " Please report this as a bug in mantis." )
            else: # Keeping this for-loop as a fallback, it should never add dependencies though
                can_add_me = True
                for dep in mantis_node.hierarchy_dependencies:
                    if not dep.prepared and dep not in unprepared:
                        if dep in forbidden:
                            can_add_me=False
                            forbidden.add(mantis_node) # forbid the parent, too
                            continue
                        unprepared.appendleft(dep)
                if can_add_me:
                    unprepared.appendleft(mantis_node) # just rotate them until they are ready.


    def solve_iteration(self):
        """ Solve an iteration of the schema.
             - 1 Create the Mantis Node instances that represent this iteration of the schema
             - 2 Connect the links from the entrypoint or previous iteration.
             - 3 Connect the constant and array links, and any link between nodes entirely within the tree
             - 4 Prepare the nodes that modify data (in case of e.g. array get index or nested schema length input)
             - 5 Connect the final prepared nodes
             and return the nodes that were created in this schema iteration (frame).
             This function also adds to held_links to pass data between iterations.
        """

        from ..schema_nodes_ui import (SchemaIndex,
                                        SchemaArrayInput,
                                        SchemaArrayInputGet,
                                        SchemaArrayInputAll,
                                        SchemaArrayOutput,
                                        SchemaConstInput,
                                        SchemaConstOutput,
                                        SchemaOutgoingConnection,
                                        SchemaIncomingConnection,)
        from ..utilities.utilities import clear_reroutes, link_mantis_nodes
        from ..base_definitions import array_output_types
        self.set_schema_per_interation_constants()
        frame_mantis_nodes = {}

        # Later we have to run bPrepare() on these guys, so we make the deque and fill it now.
        from collections import deque
        unprepared= deque()
        self.gen_solve_iteration_mantis_nodes(frame_mantis_nodes, unprepared)

        # This is where we handle node connections BETWEEN frames
        while(self.held_links):
            ui_link = self.held_links.pop()
            to_ui_node = ui_link.to_socket.node; from_ui_node = ui_link.from_socket.node
            if isinstance(to_ui_node, SchemaOutgoingConnection):
                self.handle_link_to_outgoing_connection_output(frame_mantis_nodes, ui_link)

        # Get the rerouted links from the graph. We don't really need to do this every iteration.
        # TODO: use profiling to determine if this is slow; if so: copy & reuse the data, refactor the pop()'s out.
        ui_links = clear_reroutes(list(self.tree.links))

        # Now we handle ui_links in the current frame, including those ui_links between Schema nodes and "real" nodes
        links_to_output = []
        array_input_get_link = []
        for ui_link in ui_links:
            to_ui_node = ui_link.to_socket.node; from_ui_node = ui_link.from_socket.node

            if isinstance(from_ui_node, SchemaIndex):
                if ui_link.from_socket.name == "Index":
                    self.handle_link_from_index_input(self.index, frame_mantis_nodes, ui_link)
                elif ui_link.from_socket.name == "Schema Length":
                    self.handle_link_from_schema_length_input(frame_mantis_nodes, ui_link)
                continue
            if isinstance(from_ui_node, SchemaIncomingConnection):
                if ui_link.from_socket.name in self.incoming_connections.keys():
                    self.handle_link_from_incoming_connection_input(frame_mantis_nodes, ui_link)
                continue
            if isinstance(from_ui_node, (SchemaConstInput, NodeGroupInput)):
                if ui_link.from_socket.name in self.constant_in.keys():
                    self.handle_link_from_constant_input( frame_mantis_nodes, ui_link, to_ui_node)
                continue
            if isinstance(to_ui_node, SchemaArrayInputGet):
                self.handle_link_to_array_input_get( frame_mantis_nodes, ui_link)
                continue
            if isinstance(from_ui_node, SchemaArrayInput):
                self.handle_link_from_array_input(frame_mantis_nodes, ui_link, self.index)
                continue
            if isinstance(from_ui_node, SchemaArrayInputAll):
                self.handle_link_from_array_input_all(frame_mantis_nodes, ui_link)
                continue
            # HOLD these links to the next iteration:
            if isinstance(to_ui_node, SchemaOutgoingConnection):
                if isinstance(from_ui_node, (MantisNodeGroup, SchemaGroup)):
                    self.handle_link_from_subschema_to_output(frame_mantis_nodes, ui_link, to_ui_node)
                self.held_links.append(ui_link) # is this wise? Why am I doing this?
                continue
            # HOLD these links until prep is done a little later
            if isinstance(to_ui_node, (SchemaConstOutput, NodeGroupOutput)) or isinstance(to_ui_node, SchemaArrayOutput):
                if isinstance(from_ui_node, (MantisNodeGroup, SchemaGroup)):
                    self.handle_link_from_subschema_to_output(frame_mantis_nodes, ui_link, to_ui_node)
                    # both links are desirable to create, so don't continue here
                if from_ui_node.bl_idname in array_output_types:
                    self.handle_link_from_subschema_to_output(frame_mantis_nodes, ui_link, to_ui_node)
                    # this one wires links around - we need to finish connecting it to the output
                    # before we can prepare it. Otherwise, it will send a link from *itself* instead of
                    # rewiring one of its *inputs* to the next node.
                    # We have to prep it, and then deal with the links between it and the array. pain.
                links_to_output.append(ui_link)
                continue
            if isinstance(from_ui_node, SchemaArrayInputGet):
                array_input_get_link.append(ui_link)
                continue

            # for any of the special cases, we hit a 'continue' block. So this connection is not special, and is made here.
            connection = link_mantis_nodes(self.autogen_path_names, ui_link,
                                frame_mantis_nodes, from_suffix=self.index_str,
                                to_suffix=self.index_str)

        for signature, node in frame_mantis_nodes.items():
            self.solved_nodes[signature]=node
            if node.node_type == "DUMMY_SCHEMA":
                # make sure to add the nodes to the group's sockets if the user set them directly
                from .readtree import make_connections_to_ng_dummy
                make_connections_to_ng_dummy(
                    self.node.base_tree,
                    self.autogen_path_names,
                    {}, # just pass an empty dict, this argument is not needed in this context
                    self.all_nodes,
                    node)
                from ..utilities.utilities import init_schema_dependencies
                init_schema_dependencies(node, self.all_nodes)
            else:
                init_dependencies(node) # it is hard to overstate how important this single line of code is

        # We have to prepare the nodes leading to Schema Length
        unprepared=deque()
        for node in frame_mantis_nodes.values():
            if node.node_type == 'DUMMY_SCHEMA' and (schema_len_in := node.inputs.get("Schema Length")):
                for l in schema_len_in.links:
                    unprepared.append(l.from_node)
            self.prepare_nodes(unprepared)

        # We have to prepare the nodes leading to Array Input Get
        for ui_link in array_input_get_link:
            from_name = get_link_in_out(ui_link)[0]
            # because this both provides and receives deps, it must be solved first.
            from_node = self.schema_nodes.get( (*self.node.ui_signature, ui_link.from_node.bl_idname) )
            self.handle_link_from_array_input_get(frame_mantis_nodes, ui_link )

        # Finally, we have to prepare nodes leading to outputs.
        for i in range(len(links_to_output)):
            unprepared=deque()
            ui_link = links_to_output[i]
            to_ui_node = ui_link.to_socket.node; from_ui_node = ui_link.from_socket.node
            # ugly workaround here in a very painful edge case...
            if isinstance(from_ui_node, (MantisNodeGroup, SchemaGroup)) or\
                                from_ui_node.bl_idname in array_output_types:
                ui_link=self.spoof_link_for_subschema_to_output_edge_case(ui_link)
                links_to_output[i] = ui_link
            from_name = get_link_in_out(ui_link)[0]
            signature = (*self.autogen_path_names, from_name+self.index_str)
            #use it directly if it is a mantis node; this happens when the previous node was a Schema
            if hasattr(ui_link, "from_node") and (from_node := self.schema_nodes.get( (*self.node.ui_signature, ui_link.from_node.bl_idname))):
                unprepared.append(from_node)
                unprepared.extend(from_node.hierarchy_dependencies)
            elif from_node := frame_mantis_nodes.get(signature):
                unprepared.append(from_node)
                unprepared.extend(from_node.hierarchy_dependencies)
            else:
                raise RuntimeError(" 671 there has been an error parsing the tree. Please report this as a bug.")
            self.prepare_nodes(unprepared) # prepare only the dependencies we need for this link
            # and handle the output by the specific type
            if isinstance(to_ui_node, (SchemaConstOutput, NodeGroupOutput)):
                self.handle_link_to_constant_output(frame_mantis_nodes, self.index, ui_link,  to_ui_node)
            if isinstance(to_ui_node, SchemaArrayOutput):
                if from_node.bl_idname in array_output_types:
                    from ..base_definitions import DummyLink
                    for l in from_node.rerouted:
                        new_link = DummyLink(
                            from_socket=l.from_node.outputs[l.from_socket],
                            to_socket=l.to_node.inputs[l.to_socket],
                            from_mantis_node=l.from_node, to_mantis_node=l.to_node,
                            multi_input_sort_id=l.multi_input_sort_id
                        )
                        self.handle_link_from_array_type_to_array_out(ui_link, new_link)
                else:
                    self.handle_link_to_array_output(frame_mantis_nodes, self.index, ui_link, to_ui_node, from_ui_node)
        return frame_mantis_nodes

    def solve_nested_schema(self, schema_mantis_node):
        """ Solves all schema node groups found in this Schema. This is a recursive function, which will
            solve all levels of nested schema - since this function is called by solver.solve().
        """
        solver=None
        if schema_mantis_node.prepared == False:
            all_nodes = self.all_nodes.copy()
            ui_node = schema_mantis_node.ui_node
            length = schema_mantis_node.evaluate_input("Schema Length")
            tree = ui_node.node_tree
            if schema_mantis_node.ui_node.bl_idname == "MantisNodeGroup":
                prOrange(f"Expanding Node Group {tree.name} in node {schema_mantis_node}.")
            else:
                prOrange(f"Expanding schema {tree.name} in node {schema_mantis_node} with length {length}.")
            solver = SchemaSolver(schema_mantis_node, all_nodes, ui_node, schema_mantis_node.ui_signature, error_popups=self.error_popups)
            solved_nodes = solver.solve()
            schema_mantis_node.prepared = True
            for k,v in solved_nodes.items():
                self.solved_nodes[k]=v
        return solver


    def finalize(self, mantis_nodes_in_frame):
        from ..schema_nodes_ui import (SchemaOutgoingConnection,)
        for i in range(len(self.held_links)):
            link = self.held_links.pop()
            to_np = link.to_socket.node; from_np = link.from_socket.node
            if isinstance(to_np, SchemaOutgoingConnection):
                if link.to_socket.name in self.outgoing_connections.keys():
                    if (outgoing_links := self.outgoing_connections[link.to_socket.name]) is None: continue
                    for outgoing in outgoing_links:
                        if outgoing:
                            to_node = outgoing.to_node
                            from_node =mantis_nodes_in_frame[(*self.autogen_path_names, from_np.name+self.index_str) ]
                            from_socket_name = link.from_socket.name
                            if from_node.node_type in ['DUMMY_SCHEMA']:
                                from_socket_name = link.from_socket.identifier
                            connection = from_node.outputs[from_socket_name].connect(node=to_node, socket=outgoing.to_socket)
                            # we need to kill the link between the Schema itself and the next node and update the deps. Otherwise:confusing bugs.
                            outgoing.die(); init_dependencies(to_node)
                    # else: # the node just isn't connected out this socket.


        # # solve all unsolved nested schemas...
        for schema_sig, schema_mantis_node in self.nested_schemas.items():
                self.solve_nested_schema(schema_mantis_node)


        for n in self.autogenerated_nodes.values():
            init_connections(n)
            for c in n.connections:
                init_dependencies(c)

        all_outgoing_links = []
        for conn in self.outgoing_connections.values():
            for outgoing in conn:
                all_outgoing_links.append(outgoing)
        for conn in self.constant_out.values():
            for outgoing in conn:
                all_outgoing_links.append(outgoing)
        for conn in self.array_output_connections.values():
            for outgoing in conn:
                all_outgoing_links.append(outgoing)

        for outgoing in all_outgoing_links:
            to_node = outgoing.to_node
            for l in to_node.inputs[outgoing.to_socket].links:
                if self.node == l.from_node:
                    l.die()

        for inp in self.node.inputs.values():
            for l in inp.links:
                init_connections(l.from_node) # to force it to have hierarchy connections with the new nodes.


    def solve(self):
        if self.solve_length < 1:
            from ..base_definitions import GraphError
            for o in self.node.outputs:
                if o.is_linked:
                    raise GraphError(f"ERROR: Schema {self} has a length"
                                     " of 0 but other nodes depend on it.")
            print (f"WARN: Schema {self} has a length of 0 or less and will not expand.")
            return {} # just don't do anything - it's OK to have a noop schema if it doesn't have dependencies.
        for index in range(self.solve_length):
            self.index = index
            frame_mantis_nodes = self.solve_iteration()
            for sig, mantis_node in frame_mantis_nodes.items():
                if mantis_node.node_type == 'DUMMY_SCHEMA':
                    self.nested_schemas[sig] = mantis_node
        self.finalize(frame_mantis_nodes)
        self.solved = True
        self.node.prepared = True
        prGreen(f"Schema declared {len(self.solved_nodes)} nodes.\n")
        return self.solved_nodes