mantis/readtree.py

from .utilities import prRed, prGreen, prPurple, prWhite, prOrange, \
                        wrapRed, wrapGreen, wrapPurple, wrapWhite, wrapOrange
from .utilities import  gen_nc_input_for_data

    


def grp_node_reroute_common(nc, nc_to, all_nc):
    # we need to do this: go  to the to-node
    # then reroute the link in the to_node all the way to the beginning
    # so that the number of links in "real" nodes is unchanged
    # then the links in the dummy nodes need to be deleted
    for inp_name, inp in nc.inputs.items():
        # assume each input socket only has one input for now
        if inp.is_connected:
            while (inp.links):
                in_link = inp.links.pop()
                from_nc = in_link.from_node
                from_socket = in_link.from_socket
                links = []
                from_links = from_nc.outputs[from_socket].links.copy()
                while(from_links):
                    from_link = from_links.pop()
                    if from_link == in_link:
                        from_link.die()
                        continue # DELETE the dummy node link 
                    links.append(from_link)
                from_nc.outputs[from_socket].links = links
                down = nc_to.outputs[inp_name]
                for downlink in down.links:
                    downlink.from_node = from_nc
                    downlink.from_socket = from_socket
                    from_nc.outputs[from_socket].links.append(downlink)
                    if hasattr(downlink.to_node, "reroute_links"):
                        downlink.to_node.reroute_links(downlink.to_node, all_nc)
                in_link.die()

def reroute_links_grp(nc, all_nc):
    if nc.inputs:
        if (nc_to := all_nc.get( ( *nc.signature, "NodeGroupInput") )):
            grp_node_reroute_common(nc, nc_to, all_nc)
        else:
            raise RuntimeError("internal error: failed to enter a node group ")

def reroute_links_grpout(nc, all_nc):
    if (nc_to := all_nc.get( ( *nc.signature[:-1],) )):
        grp_node_reroute_common(nc, nc_to, all_nc)
    else:
        raise RuntimeError("error leaving a node group (maybe you are running the tree from inside a node group?)")

# FIXME I don't think these signatures are unique.
def insert_lazy_parents(nc):
    from .link_containers import LinkInherit
    from .base_definitions import NodeLink
    inherit_nc = None
    if nc.inputs["Relationship"].is_connected:
        link = nc.inputs["Relationship"].links[0]
        # print(nc)
        from_nc = link.from_node
        if from_nc.node_type in ["XFORM"] and link.from_socket in ["xForm Out"]:
            inherit_nc = LinkInherit(("MANTIS_AUTOGENERATED", *nc.signature[1:], "LAZY_INHERIT"), nc.base_tree)
            for from_link in from_nc.outputs["xForm Out"].links:
                if from_link.to_node == nc and from_link.to_socket == "Relationship":
                    break # this is it
            from_link.to_node = inherit_nc; from_link.to_socket="Parent"
            from_link.to_node.inputs[from_link.to_socket].is_linked=True
            
            links=[]
            while (nc.inputs["Relationship"].links):
                to_link = nc.inputs["Relationship"].links.pop()
                if to_link.from_node == from_nc and to_link.from_socket == "xForm Out":
                    continue # don't keep this one
                links.append(to_link)
                to_link.from_node.outputs[from_link.from_socket].is_linked=True
            
            nc.inputs["Relationship"].links=links
            link=NodeLink(from_node=inherit_nc, from_socket="Inheritance", to_node=nc, to_socket="Relationship")
            inherit_nc.inputs["Parent"].links.append(from_link)
            
            inherit_nc.parameters = {
                                     "Parent":None,
                                     "Inherit Rotation":True,
                                     "Inherit Scale":'FULL',
                                     "Connected":False,
                                    }
            # because the from node may have already been done.
            init_connections(from_nc)
            init_dependencies(from_nc)
            init_connections(inherit_nc)
            init_dependencies(inherit_nc)
    return inherit_nc


from_name_filter = ["Driver", ]

to_name_filter = [
                   "Custom Object xForm Override",
                   "Custom Object",
                   "Deform Bones"
                 ]


# *** # *** # *** # *** # *** # *** # *** # *** # *** # *** # *** # *** # *** # *** #
#                                  DATA FROM NODES                                  #
# *** # *** # *** # *** # *** # *** # *** # *** # *** # *** # *** # *** # *** # *** #

from .base_definitions import replace_types, NodeSocket

# TODO: investigate whether I can set the properties in the downstream nodes directly.
#       I am doing this in Schema Solver and it seems to work quite efficiently.
def make_connections_to_ng_dummy(base_tree, tree_path_names, local_nc, all_nc, nc_to):
    np = nc_to.prototype
    for inp in np.inputs:
        nc_from = None
        if inp.bl_idname in  ['WildcardSocket']:
            continue # it isn't a real input so I don't think it is good to check it.
        to_s = inp.identifier
        if not inp.is_linked: # make an autogenerated NC for the inputs of the group node
            if inp.bl_idname in ['xFormSocket']:
                continue
            from .node_container_common import get_socket_value
            nc_cls = gen_nc_input_for_data(inp)
            if (nc_cls):
                sig = ("MANTIS_AUTOGENERATED", *tree_path_names, np.name, inp.name, inp.identifier)
                nc_from = nc_cls(sig, base_tree)
                # ugly! maybe even a HACK!
                nc_from.inputs = {}
                nc_from.outputs = {inp.name:NodeSocket(name = inp.name, node=nc_from)}
                nc_from.parameters = {inp.name:get_socket_value(inp)}
                # 
                local_nc[sig] = nc_from; all_nc[sig] = nc_from
                from_s = inp.name
            else:
                prRed("No available auto-generated class for input", *tree_path_names, np.name, inp.name)
            nc_from.outputs[from_s].connect(node=nc_to, socket=to_s, sort_id=0)

def gen_node_containers(base_tree, current_tree, tree_path_names, all_nc, local_nc, dummy_nodes, group_nodes, schema_nodes ):
    from .internal_containers import DummyNode
    for ui_node in current_tree.nodes:
        if ui_node.bl_idname in ["NodeFrame", "NodeReroute"]:
            continue # not a Mantis Node
        if ui_node.bl_idname in ["NodeGroupInput", "NodeGroupOutput"]:
            # we only want ONE dummy in/out per tree_path, so use the bl_idname to make a Dummy node
            sig = (None, *tree_path_names, ui_node.bl_idname)
            if not local_nc.get(sig):
                nc = DummyNode( signature=sig , base_tree=base_tree, prototype=ui_node )
                local_nc[sig] = nc; all_nc[sig] = nc; dummy_nodes[sig] = nc
                if ui_node.bl_idname in ["NodeGroupOutput"]:
                    nc.reroute_links = reroute_links_grpout
        elif ui_node.bl_idname in  ["MantisNodeGroup", "MantisSchemaGroup"]:
            nc = DummyNode( signature= (sig := (None, *tree_path_names, ui_node.name) ), base_tree=base_tree, prototype=ui_node )
            local_nc[sig] = nc; all_nc[sig] = nc; dummy_nodes[sig] = nc
            make_connections_to_ng_dummy(base_tree, tree_path_names, local_nc, all_nc, nc)
            if ui_node.bl_idname == "MantisNodeGroup":
                group_nodes.append(nc)
                nc.reroute_links = reroute_links_grp
            else:
                group_nodes.append(nc)
                schema_nodes[sig] = nc
        # if it wasn't the types we ignore or the types we make a Dummy for, use this to catch all non-special cases.
        elif (nc_cls := ui_node.mantis_class):
            sig = (None, *tree_path_names, ui_node.name)
            if ui_node.bl_idname in replace_types:
                sig = (None, *tree_path_names, ui_node.bl_idname)
                if local_nc.get(sig):
                    continue # already made
            nc = nc_cls( sig , base_tree)
            local_nc[sig] = nc; all_nc[sig] = nc
        else:
            nc = None
            prRed(f"Can't make nc for.. {ui_node.bl_idname}")
        # this should be done at init
        if nc.signature[0] not in ['MANTIS_AUTOGENERATED'] and nc.node_type not in ['SCHEMA', 'DUMMY', 'DUMMY_SCHEMA']:
            nc.fill_parameters()

def data_from_tree(base_tree, tree_path, dummy_nodes, all_nc, all_schema):#
    # TODO: it should be relatively easy to make this use a while loop instead of recursion.
    local_nc, group_nodes = {}, []
    tree_path_names = [tree.name for tree in tree_path if hasattr(tree, "name")]
    if tree_path[-1]:
        current_tree = tree_path[-1].node_tree # this may be None.
    else:
        current_tree = base_tree
    #
    if current_tree: # the node-group may not have a tree set - if so, ignore it.
        from .utilities import clear_reroutes
        links = clear_reroutes(list(current_tree.links))
        gen_node_containers(base_tree, current_tree, tree_path_names, all_nc, local_nc, dummy_nodes, group_nodes, all_schema)
        
        from .utilities import link_node_containers
        for link in links:
            link_node_containers((None, *tree_path_names), link, local_nc)
        # Now, descend into the Node Groups and recurse
        for nc in group_nodes:
            data_from_tree(base_tree, tree_path+[nc.prototype], dummy_nodes, all_nc, all_schema)
    return dummy_nodes, all_nc, all_schema

from .utilities import check_and_add_root, init_connections, init_dependencies, init_schema_dependencies

def is_signature_in_other_signature(parent_signature, child_signature):
    # If the other signature is shorter, it isn't a child node
    if len(parent_signature) > len(child_signature):
        return False
    return parent_signature[0:] == child_signature[:len(parent_signature)]

def solve_schema_to_tree(nc, all_nc, roots=[]):
    from .utilities import get_node_prototype
    np = get_node_prototype(nc.signature, nc.base_tree)
    from .schema_solve import SchemaSolver
    tree = np.node_tree
    length = nc.evaluate_input("Schema Length")
    prOrange(f"Expanding schema {tree.name} in node {nc} with length {length}.")

    solver = SchemaSolver(nc, all_nc, np)
    solved_nodes = solver.solve()
    prWhite(f"Schema declared {len(solved_nodes)} nodes.")

    # maybe this should be done in schema solver. TODO invesitigate a more efficient way
    del_me = []
    for k, v in all_nc.items():
        # delete all the schema's internal nodes. The links have already been deleted by the solver.
        if v.signature[0] not in ['MANTIS_AUTOGENERATED'] and is_signature_in_other_signature(nc.signature, k):
            del_me.append(k)
    for k in del_me:
        del all_nc[k]
    for k,v in solved_nodes.items():
        all_nc[k]=v
        init_connections(v)
        check_and_add_root(v, roots, include_non_hierarchy=True)

    return solved_nodes

# *** # *** # *** # *** # *** # *** # *** # *** # *** # *** # *** # *** # *** # *** #
#                                  PARSE NODE TREE                                  #
# *** # *** # *** # *** # *** # *** # *** # *** # *** # *** # *** # *** # *** # *** #

schema_bl_idnames = [   "SchemaIndex",
                        "SchemaArrayInput",
                        "SchemaArrayInputGet",
                        "SchemaArrayInputAll",
                        "SchemaArrayOutput",
                        "SchemaConstInput",
                        "SchemaConstOutput",
                        "SchemaOutgoingConnection",
                        "SchemaIncomingConnection", 
                    ]

from .utilities import get_all_dependencies
def get_schema_length_dependencies(node, all_nodes={}):
    """ Get a list of all dependencies for the given node's length or array properties.
        This function will also recursively search for dependencies in its sub-trees.
    """
    deps = []
    prepare_links_to = ['Schema Length','Array', 'Index']
    def extend_dependencies_from_inputs(node):
        for inp in node.inputs.values():
            for l in inp.links:
                if "MANTIS_AUTOGENERATED" in l.from_node.signature: 
                    deps.extend([l.from_node]) # why we need this lol
                if inp.name in prepare_links_to:
                    deps.append(l.from_node)
                    deps.extend(get_all_dependencies(l.from_node))
    def deps_filter(dep): # remove any nodes inside the schema
        if len(dep.signature) > len(node.signature):
            for i in range(len(node.signature)):
                dep_sig_elem, node_sig_elem = dep.signature[i], node.signature[i]
                if dep_sig_elem != node_sig_elem: break # they don't match, it isn't an inner-node
            else: # remove this, it didn't break, meaning it shares signature with outer node
                return False # this is an inner-node
        return True
    # this way we can handle Schema and Array Get nodes with one function
    extend_dependencies_from_inputs(node)
    if node.node_type == 'DUMMY_SCHEMA':
        trees = [(node.prototype.node_tree, node.signature)] # this is UI data
        while trees:
            tree, tree_signature = trees.pop()
            for sub_ui_node in tree.nodes:
                if sub_ui_node.bl_idname in ['NodeReroute', 'NodeFrame']:
                    continue
                if sub_ui_node.bl_idname in schema_bl_idnames:
                    sub_node = all_nodes[(*tree_signature, sub_ui_node.bl_idname)]
                else:
                    sub_node = all_nodes[(*tree_signature, sub_ui_node.name)]
                if sub_node.node_type == 'DUMMY_SCHEMA':
                    extend_dependencies_from_inputs(sub_node)
                    trees.append((sub_node.prototype.node_tree, sub_node.signature))
    return list(filter(deps_filter, deps))


def parse_tree(base_tree):
    from uuid import uuid4
    base_tree.execution_id = uuid4().__str__() # set the unique id of this execution
   
    import time
    data_start_time = time.time()
    # annoyingly I have to pass in values for all of the dicts because if I initialize them in the function call
    #  then they stick around because the function definition inits them once and keeps a reference
    # so instead I have to supply them to avoid ugly code or bugs elsewhere
    # it's REALLY confusing when you run into this sort of problem. So it warrants four entire lines of comments!      
    dummy_nodes, all_mantis_nodes, all_schema =  data_from_tree(base_tree, tree_path = [None], dummy_nodes = {}, all_nc = {}, all_schema={})
    for dummy in dummy_nodes.values():    # reroute the links in the group nodes
        if (hasattr(dummy, "reroute_links")):
            dummy.reroute_links(dummy, all_mantis_nodes)
    prGreen(f"Pulling data from tree took {time.time() - data_start_time} seconds")
    
    start_time = time.time()
    roots, array_nodes = [], []
    from .misc_nodes import UtilityArrayGet
    for mantis_node in all_mantis_nodes.values():
        if mantis_node.node_type in ["DUMMY"]: # clean up the groups
            if mantis_node.prototype.bl_idname in ("MantisNodeGroup", "NodeGroupOutput"):
                continue
        # Initialize the dependencies and connections (from/to links) for each node.
        # we record & store it because using a getter is much slower (according to profiling)
        init_dependencies(mantis_node); init_connections(mantis_node)
        check_and_add_root(mantis_node, roots, include_non_hierarchy=True)
        # Array nodes need a little special treatment, they're quasi-schemas
        if isinstance(mantis_node, UtilityArrayGet):
            array_nodes.append(mantis_node)

    from collections import deque
    unsolved_schema = deque()
    solve_only_these = []; solve_only_these.extend(list(all_schema.values()))
    for schema in all_schema.values():
        # We can remove the schema that are inside another schema tree.
        for i in range(len(schema.signature)-1): # -1, we don't want to check this node, obviously
            if parent := all_schema.get(schema.signature[:i+1]):
                # This will be solved along with its parent schema.
                solve_only_these.remove(schema)
                break
        else:
            init_schema_dependencies(schema, all_mantis_nodes)
            solve_only_these.extend(get_schema_length_dependencies(schema, all_mantis_nodes))
            unsolved_schema.append(schema)
    for array in array_nodes:
        solve_only_these.extend(get_schema_length_dependencies(array))
    solve_only_these.extend(array_nodes)
    schema_solve_done = set()

    solve_only_these = set(solve_only_these)

    solve_layer = unsolved_schema.copy(); solve_layer.extend(roots)
    
    while(solve_layer):
        n = solve_layer.pop()
        if n not in solve_only_these: # removes the unneeded node from the solve-layer
            continue

        if n.signature in all_schema.keys():
            for dep in n.hierarchy_dependencies:
                if dep not in schema_solve_done and (dep in solve_only_these):
                    if dep.prepared: # HACK HACK HACK
                        continue 
                        # For some reason, the Schema Solver is able to detect and resolve dependencies outside
                        #  of solve_only_these. So I have to figure out why.
                    solve_layer.appendleft(n)
                    break
            else:
                solved_nodes = solve_schema_to_tree(n, all_mantis_nodes, roots)
                unsolved_schema.remove(n)
                schema_solve_done.add(n)
                for node in solved_nodes.values():
                    #
                    init_dependencies(node)
                    init_connections(node)
                    #
                    solve_layer.appendleft(node)
                for conn in n.hierarchy_connections:
                    if conn not in schema_solve_done and conn not in solve_layer:
                        solve_layer.appendleft(conn)
        else:
            for dep in n.hierarchy_dependencies:
                if dep not in schema_solve_done:
                    break
            else:
                n.bPrepare()
                schema_solve_done.add(n)
                for conn in n.hierarchy_connections:
                    if conn not in schema_solve_done and conn not in solve_layer:
                        solve_layer.appendleft(conn)
    if unsolved_schema:
        raise RuntimeError("Failed to resolve all schema declarations")
    # I had a problem with this looping forever. I think it is resolved... but I don't know lol

    all_mantis_nodes = list(all_mantis_nodes.values())
    kept_nc = {}
    while (all_mantis_nodes):
        nc = all_mantis_nodes.pop()
        if nc in array_nodes:
            continue
        if nc.node_type in ["DUMMY"]:
                continue
        # cleanup autogen nodes
        if nc.signature[0] == "MANTIS_AUTOGENERATED" and len(nc.inputs) == 0 and len(nc.outputs) == 1:
            output=list(nc.outputs.values())[0]
            value=list(nc.parameters.values())[0]   # IDEA modify the dependecy get function to exclude these nodes completely
            for l in output.links:
                to_node = l.to_node; to_socket = l.to_socket
                l.die()
                to_node.parameters[to_socket] = value
                del to_node.inputs[to_socket]
                init_dependencies(to_node)
            continue

        if (nc.node_type in ['XFORM']) and ("Relationship" in nc.inputs.keys()):
            if (new_nc := insert_lazy_parents(nc)):
                kept_nc[new_nc.signature]=new_nc
        kept_nc[nc.signature]=nc
    prWhite(f"Parsing tree took {time.time()-start_time} seconds.")
    prWhite("Number of Nodes: %s" % (len(kept_nc)))
    return kept_nc


def switch_mode(mode='OBJECT', objects = []):
    active = None
    if objects:
        from bpy import context, ops
        active = objects[-1]
        context.view_layer.objects.active = active
        if (active):
            with context.temp_override(**{'active_object':active, 'selected_objects':objects}):
                ops.object.mode_set(mode=mode)
    return active

def execution_error_cleanup(node, exception, switch_objects = [] ):
    from bpy import  context
    if node:
        # TODO: see about zooming-to-node.
        base_tree = node.base_tree
        tree = base_tree
        try:
            pass
            space = context.space_data
            for name in node.signature[1:]:
                for n in tree.nodes: n.select = False
                n = tree.nodes[name]
                n.select = True
                tree.nodes.active = n
                if hasattr(n, "node_tree"):
                    tree = n.node_tree
        except AttributeError: # not being run in node graph
            pass
        finally:
            def error_popup_draw(self, context):
                self.layout.label(text=f"Error: {exception}")
                self.layout.label(text=f"see node: {node.signature[1:]}.")
            context.window_manager.popup_menu(error_popup_draw, title="Error", icon='ERROR')
    switch_mode(mode='OBJECT', objects=switch_objects)
    for ob in switch_objects:
        ob.data.pose_position = 'POSE'
    prRed(f"Error: {exception} in node {node}")
    return exception


def execute_tree(nodes, base_tree, context, error_popups = False):
    import bpy
    from time import time
    from .node_container_common import GraphError
    original_active = context.view_layer.objects.active
    start_execution_time = time()

    from collections import deque
    xForm_pass = deque()
    for nc in nodes.values():
        nc.prepared = False
        nc.executed = False
        check_and_add_root(nc, xForm_pass)

    executed = []

    # check for cycles here by keeping track of the number of times a node has been visited.
    visited={}
    check_max_len=len(nodes)**2 # seems too high but safe. In a well-ordered graph, I guess this number should be less than the number of nodes.
    max_iterations = len(nodes)**2
    i = 0
    switch_me = [] # switch the mode on these objects
    active = None # only need it for switching modes
    select_me = []
    try:
        while(xForm_pass):
            if i >= max_iterations:
                raise GraphError("There is probably a cycle somewhere in the graph.")
            i+=1    
            n = xForm_pass.pop()
            if visited.get(n.signature):
                visited[n.signature]+=1
            else:
                visited[n.signature]=0
            if visited[n.signature] > check_max_len:
                raise GraphError("There is a probably a cycle in the graph somewhere. Fix it!")
                # we're trying to solve the halting problem at this point.. don't do that.
                # TODO find a better way! there are algo's for this but they will require using a different solving algo, too
            if n.prepared:
                continue
            if n.node_type not in ['XFORM', 'UTILITY']:
                for dep in n.hierarchy_dependencies:
                    if not dep.prepared:
                        xForm_pass.appendleft(n) # hold it
                        break
                else:
                    n.prepared=True
                    executed.append(n)
                    for conn in n.hierarchy_connections:
                        if  not conn.prepared:
                            xForm_pass.appendleft(conn)
            else:
                for dep in n.hierarchy_dependencies:
                    if not dep.prepared:
                        break
                else:
                    try:
                        n.bPrepare(context)
                        if not n.executed:
                            n.bExecute(context)
                        if (n.__class__.__name__ == "xFormArmature" ):
                            ob = n.bGetObject()
                            switch_me.append(ob)
                            active = ob
                        if not (n.__class__.__name__ == "xFormBone" ) and hasattr(n, "bGetObject"):
                            ob = n.bGetObject()
                            if isinstance(ob, bpy.types.Object):
                                select_me.append(ob)
                    except Exception as e:
                        if error_popups:
                            raise execution_error_cleanup(n, e,)
                        else:
                            raise e
                    n.prepared=True
                    executed.append(n)
                    for conn in n.hierarchy_connections:
                        if  not conn.prepared:
                            xForm_pass.appendleft(conn)
        

        switch_mode(mode='POSE', objects=switch_me)
        if (active):
            with context.temp_override(**{'active_object':active, 'selected_objects':switch_me}):
                bpy.ops.object.mode_set(mode='POSE')

        for n in executed:
            try:
                n.bPrepare(context)
                if not n.executed:
                    n.bExecute(context)
            except Exception as e:
                if error_popups:
                    raise execution_error_cleanup(n, e,)
                else:
                    raise e


        switch_mode(mode='OBJECT', objects=switch_me)
        for ob in switch_me:
            ob.data.pose_position = 'POSE'
        # switch to pose mode here so that the nodes can use the final pose data
        # this will require them to update the depsgraph.
        
        for n in executed:
            try:
                n.bFinalize(context)
            except Exception as e:
                if error_popups:
                    raise execution_error_cleanup(n, e,)
                else:
                    raise e
        
        tot_time = (time() - start_execution_time)
        prGreen(f"Executed tree of {len(executed)} nodes in {tot_time} seconds")
        if (original_active):
            context.view_layer.objects.active = original_active
            original_active.select_set(True)
    except Exception as e:
        execution_error_cleanup(None, e, switch_me)
        if error_popups == False:
            raise e
    finally:
        context.view_layer.objects.active = active
        # clear the selection first.
        for ob in context.selected_objects:
            try:
                ob.select_set(False)
            except RuntimeError: # it isn't in the view layer
                pass
        for ob in select_me:
            try:
                ob.select_set(True)
            except RuntimeError: # it isn't in the view layer
                pass