mantis/schema_solve.py

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



class SchemaSolver:
    def __init__(self, schema_dummy, nodes, prototype, signature=None,):
        self.all_nodes = nodes # this is the parsed tree from Mantis
        self.node = schema_dummy
        self.tree = prototype.node_tree
        self.uuid = self.node.uuid
        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.prototype.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.init_schema_links()
        self.set_index_strings()

        # 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=lambda a : -a.multi_input_sort_id)

        for ui_node in self.tree.nodes:
            # first we need to fill the parameters of the schema nodes.
            # 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):
                # 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
                mantis_node.fill_parameters(ui_node)
            # HACK to make Group Nodes work
            if ui_node.bl_idname == "NodeGroupInput":
                from .schema_containers import SchemaConstInput
                mantis_node = SchemaConstInput(signature=signature, base_tree=self.node.base_tree, parent_schema_node=self.node)
                self.schema_nodes[signature] = mantis_node
                mantis_node.fill_parameters(ui_node)
            if ui_node.bl_idname == "NodeGroupOutput":
                from .schema_containers import SchemaConstOutput
                mantis_node = SchemaConstOutput(signature=signature, base_tree=self.node.base_tree, parent_schema_node=self.node)
                self.schema_nodes[signature] = mantis_node
                mantis_node.fill_parameters(ui_node)
    
    def set_index_strings(self):
        self.index_str = lambda : '.'+str(self.uuid)+'.'+str(self.index).zfill(4)
        self.prev_index_str = lambda : '.'+str(self.uuid)+'.'+str(self.index-1).zfill(4)
        if self.is_node_group:
            self.index_str=lambda : ''
            self.prev_index_str=lambda : ''
    
    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
            parent_name='Constant'
            if item.parent.name != '': # in an "prphan" item this is left blank , it is not None or an AttributeError.
                parent_name = item.parent.name
            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()
    

    

    def gen_solve_iteration_mantis_nodes(self, frame_mantis_nodes, unprepared):
        for prototype_ui_node in self.tree.nodes:
            mantis_node_name = prototype_ui_node.name
            index_str = self.index_str()
            mContext=self.node.mContext
            if isinstance(prototype_ui_node, SchemaUINode):
                continue  # IGNORE the schema interface nodes, we already made them in __init__()
                # they are reused for each iteration.
            elif prototype_ui_node.bl_idname in ['NodeFrame', 'NodeReroute']:
                continue  # IGNORE stuff that is purely UI - frames, reroutes.
            elif prototype_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 prototype of the mantis node which we make for this iteration
            # for Schema sub-nodes ... they need a prototype to init.
            if prototype_mantis_node.node_type in ['DUMMY', 'DUMMY_SCHEMA']:
                # We stored the prototype ui_node when creating the Mantis node.
                ui_node = prototype_mantis_node.prototype
                # 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, prototype=ui_node,
                        ui_signature =  prototype_mantis_node.signature)
                    # now let's copy the links from the prototype 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
                        mantis_node.inputs.init_sockets(['Schema Length']) # add a Schema Length socket
                        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 prototype 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, prototype=ui_node)
            else:
                mantis_node = prototype_mantis_node.__class__(signature, prototype_mantis_node.base_tree)
            frame_mantis_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, prototype_ui_node)
            mantis_node.fill_parameters(prototype_ui_node)
            # be sure to pass on the Mantis Context to them
            mantis_node.mContext=mContext
            


    def handle_link_from_index_input(self, index, frame_mantis_nodes, ui_link):
        _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 to_node.node_type in ['DUMMY', 'DUMMY_SCHEMA']:
            from .utilities import gen_nc_input_for_data
            nc_cls = gen_nc_input_for_data(ui_link.from_socket)
            if (nc_cls): #HACK
                unique_name = "".join([
                    ui_link.to_socket.node.name+self.index_str(),
                    ui_link.from_socket.name,
                    ui_link.from_socket.identifier,
                    "==>",
                    ui_link.to_socket.name,
                    ui_link.to_socket.identifier,
                    ])
                sig = ("MANTIS_AUTOGENERATED", *self.tree_path_names[1:-1], unique_name)
                nc_from = frame_mantis_nodes.get(sig)
                if not nc_from:
                    nc_from = nc_cls(sig, self.node.base_tree)
                # ugly! maybe even a HACK!
                nc_from.inputs = {}
                from .base_definitions import NodeSocket
                nc_from.outputs = {ui_link.from_socket.name:NodeSocket(name = ui_link.from_socket.name, node=nc_from)}
                nc_from.parameters = {ui_link.from_socket.name:index}
                frame_mantis_nodes[sig]=nc_from
                from_node = nc_from
                self.solved_nodes[sig]=from_node
                _connection = from_node.outputs[ui_link.from_socket.name].connect(node=to_node, socket=ui_link.to_socket.identifier)
            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("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("I was expecting there to be an incoming connection here for Schema Length")

    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_mantis_socket = 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 = frame_mantis_nodes[ (*self.autogen_path_names, to_mantis_node.signature[-1]+self.index_str()) ]
            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_mantis_socket)
            # 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]
        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_to_array_input_get(self, frame_mantis_nodes, ui_link, index):
        from_name, to_name = get_link_in_out(ui_link)
        from_nc = frame_mantis_nodes[(*self.autogen_path_names, from_name+self.index_str())]
        to_nc = self.schema_nodes[(*self.tree_path_names, to_name)]
        # this only needs to be done once:
        if 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_nc = self.all_nodes[(*self.tree_path_names, from_name)]
            # I am not sure about this!
            existing_link = old_nc.outputs[ui_link.from_socket.name].links[0]
            existing_link.die()
        #
        connection = from_nc.outputs[ui_link.from_socket.name].connect(node=to_nc, socket=ui_link.to_socket.name)
    
    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)
            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==')
        # 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 # just set it directly since it is getting set to None somewhere (I should find out where tho)
        # end HACK
        if index == expose_when:
            for outgoing in self.constant_out[ui_link.to_socket.name]:
                to_node = outgoing.to_node
                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.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
        for outgoing in self.array_output_connections[ui_link.to_socket.identifier]:
            # print (type(outgoing))
            from .schema_containers import SchemaIndex
            from_name = get_link_in_out(ui_link)[0]
            from_node = frame_mantis_nodes[ (*self.autogen_path_names, from_name+self.index_str()) ]
            if not from_node:
                from_node = self.schema_nodes[(*self.tree_path_names, from_ui_node.bl_idname)]
            to_node = outgoing.to_node

            if isinstance(from_node, SchemaIndex): # I think I need to dedup this stuff
                # print("INDEX")
                from .utilities import gen_nc_input_for_data
                nc_cls = gen_nc_input_for_data(ui_link.from_socket)
                if (nc_cls): #HACK
                    sig = ("MANTIS_AUTOGENERATED", *self.tree_path_names[1:-1], self.index_str(), ui_link.from_socket.name, ui_link.from_socket.identifier)
                    nc_from = nc_cls(sig, self.node.base_tree)
                    # ugly! maybe even a HACK!
                    nc_from.inputs = {}
                    from .node_container_common import NodeSocket
                    nc_from.outputs = {ui_link.from_socket.name:NodeSocket(name = ui_link.from_socket.name, node=nc_from)}
                    from .node_container_common import get_socket_value
                    if ui_link.from_socket.name in ['Index']:
                        nc_from.parameters = {ui_link.from_socket.name:index}
                    else:
                        nc_from.parameters = {ui_link.from_socket.name:self.solve_length}
                    frame_mantis_nodes[sig]=nc_from
                    from_node = nc_from
                    self.solved_nodes[sig]=from_node

            # I have a feeling that something bad will happen if both of these conditions (above and below) are true
            if to_node.node_type == 'DUMMY_SCHEMA' and to_node.prepared:
                other_stem = ('SCHEMA_AUTOGENERATED', *to_node.signature[1:])
                from .utilities import get_node_prototype
                other_schema_np = get_node_prototype(to_node.signature, to_node.base_tree)
                other_schema_tree = other_schema_np.node_tree
                for n in other_schema_tree.nodes:
                    if n.bl_idname not in ["SchemaArrayInput", "SchemaArrayInputGet"]:
                        continue
                    out = n.outputs[outgoing.to_socket]
                    for l in out.links:
                        other_index_str = lambda : '.'+str(to_node.uuid)+'.'+str(index).zfill(4)
                        # get it by [], we want a KeyError if this fails
                        try:
                            out_node = self.all_nodes[(*other_stem, l.to_node.name+other_index_str())]
                        except KeyError as e:
                            for n in self.all_nodes:
                                if len(n) > len(other_stem)+1: break
                                for elem in other_stem:
                                    if elem not in n: break
                                else:
                                    print(n)
                            raise e
                        connection = from_node.outputs[ui_link.from_socket.name].connect(node=out_node, socket=l.to_socket.name)
            else:
                connection = from_node.outputs[ui_link.from_socket.name].connect(node=to_node, socket=outgoing.to_socket)

    def handle_link_from_array_input_get(self, frame_mantis_nodes, index, ui_link, from_ui_node ):
        get_index = index
        from_node = self.schema_nodes[(*self.tree_path_names, from_ui_node.bl_idname)]
        from .utilities import cap, wrap
        get_index = from_node.evaluate_input("Index", index)
        oob = from_node.evaluate_input("OoB Behaviour")
        # we must assume that the array has sent the correct number of links
        if oob == 'WRAP':
            get_index = wrap(get_index, len(self.array_input_connections[ui_link.from_socket.identifier])-1, 0)
        if oob == 'HOLD':
            get_index = cap(get_index, len(self.array_input_connections[ui_link.from_socket.identifier])-1)
        try:
            incoming = self.array_input_connections[ui_link.from_socket.identifier][get_index]
        except IndexError:
            raise RuntimeError(wrapRed("Dummy! You need to make it so Mantis checks if there are enough Array inputs! It should probably have a Get Index!"))
        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 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()
        # forbid some nodes - they aren't necessary to solve the schema & cause problems.
        from .readtree import execution_error_cleanup
        while unprepared:
            nc = unprepared.pop()
            if sum([dep.prepared for dep in nc.hierarchy_dependencies]) == len(nc.hierarchy_dependencies):
                try:
                    nc.bPrepare()
                except Exception as e:
                    execution_error_cleanup(nc, e)
                if nc.node_type == 'DUMMY_SCHEMA':
                    self.solve_nested_schema(nc)
            elif nc.node_type == 'DUMMY_SCHEMA' and not self.test_is_sub_schema(nc):
                forbidden.add(nc)
                continue # do NOT add this as a dependency.
            else: # Keeping this for-loop as a fallback, it should never add dependencies though
                can_add_me = True
                for dep in nc.hierarchy_dependencies:
                    if not dep.prepared and dep not in unprepared:
                        if dep in forbidden:
                            can_add_me=False
                            forbidden.add(nc) # forbid the parent, too
                            continue
                        prOrange(f"Adding dependency... {dep}")
                        unprepared.appendleft(dep)
                if can_add_me:
                    unprepared.appendleft(nc) # 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_definitions import (SchemaIndex,
                                        SchemaArrayInput,
                                        SchemaArrayInputGet,
                                        SchemaArrayInputAll,
                                        SchemaArrayOutput,
                                        SchemaConstInput,
                                        SchemaConstOutput,
                                        SchemaOutgoingConnection,
                                        SchemaIncomingConnection,)
        
        from .utilities import clear_reroutes, link_node_containers

        self.set_index_strings()
        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
        awaiting_prep_stage = []
        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, self.index)
                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):
                self.held_links.append(ui_link)
                continue 
            # HOLD these links until prep is done a little later
            if isinstance(to_ui_node, (SchemaConstOutput, NodeGroupOutput)) or isinstance(to_ui_node, SchemaArrayOutput) or \
                isinstance(from_ui_node, SchemaArrayInputGet):
                awaiting_prep_stage.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_node_containers(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 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
        
        self.prepare_nodes(unprepared)
        
        while(awaiting_prep_stage):
            ui_link = awaiting_prep_stage.pop()
            to_ui_node = ui_link.to_socket.node; from_ui_node = ui_link.from_socket.node
            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):
                self.handle_link_to_array_output(frame_mantis_nodes, self.index, ui_link, to_ui_node, from_ui_node)
            if isinstance(from_ui_node, SchemaArrayInputGet):
                self.handle_link_from_array_input_get(frame_mantis_nodes, self.index, ui_link, from_ui_node )
            # end seciton
        return frame_mantis_nodes
    
    def solve_nested_schema(self, schema_nc):
        """ 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_nc.prepared == False:
            all_nodes = self.all_nodes.copy()
            ui_node = schema_nc.prototype
            length = schema_nc.evaluate_input("Schema Length")
            tree = ui_node.node_tree
            if schema_nc.prototype.bl_idname == "MantisNodeGroup":
                prOrange(f"Expanding Node Group {tree.name} in node {schema_nc}.")
            else:
                prOrange(f"Expanding schema {tree.name} in node {schema_nc} with length {length}.")
            solver = SchemaSolver(schema_nc, all_nodes, ui_node, schema_nc.ui_signature)
            solved_nodes = solver.solve()
            schema_nc.prepared = True
            for k,v in solved_nodes.items():
                self.solved_nodes[k]=v
        return solver


    def finalize(self, frame_nc):
        from .schema_definitions 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 =frame_nc[(*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_nc in self.nested_schemas.items():
                self.solve_nested_schema(schema_nc)


        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, nc in frame_mantis_nodes.items():
                if nc.node_type == 'DUMMY_SCHEMA':
                    self.nested_schemas[sig] = nc
        self.finalize(frame_mantis_nodes)
        self.node.solver = self
        self.node.prepared = True
        return self.solved_nodes