mantis/node_container_common.py

from mantis.utilities import (prRed, prGreen, prPurple, prWhite,
                              prOrange,
                              wrapRed, wrapGreen, wrapPurple, wrapWhite,
                              wrapOrange,)
from mantis.base_definitions import GraphError, CircularDependencyError
# BE VERY CAREFUL
# the x_containers files import * from this file
# so all the top-level imports are carried over


def get_socket_value(node_socket):
    if node_socket.bl_idname in  ['RelationshipSocket', 'xFormSocket']:
        return None
    elif hasattr(node_socket, "default_value"):
        from .utilities import to_mathutils_value
        default_value_type = type(node_socket.default_value)
        math_val = to_mathutils_value(node_socket)
        if math_val:
            return math_val
        # maybe we can use it directly.. ?
        elif ( (default_value_type == str) or (default_value_type == bool) or
             (default_value_type == float) or (default_value_type == int) ):
            return node_socket.default_value
        return None

# TODO: unify the fill_paramaters for auto-gen nodes
def fill_parameters(nc):
    from .utilities import get_node_prototype
    if (nc.signature[0] is "MANTIS_AUTOGENERATED"):
        return None
    else:
        np = get_node_prototype(nc.signature, nc.base_tree)
    if ( not np ):
        raise RuntimeError(wrapRed("No node prototype found for... %s" % ( [nc.base_tree] + list(nc.signature[1:]) ) ) )
    for key in nc.parameters.keys():
        node_socket = np.inputs.get(key)
        # if (nc.signature[0] is "MANTIS_AUTOGENERATED"):
            # node_socket = None
            # for node_socket in np.inputs:
                # if node_socket.identifier == nc.signature[-1]:
                    # break
        if nc.parameters[key] is not None:
            continue # will be filled by the node itself
        if not node_socket:
            #maybe the node socket has no name
            if ( ( len(np.inputs) == 0) and ( len(np.outputs) == 1) ):
                # this is a simple input node.
                node_socket = np.outputs[0]
            elif key == 'Name': # for Links we just use the Node Label
                nc.parameters[key] = np.label
                continue
            else:
                pass
        if node_socket:
            if node_socket.bl_idname in  ['RelationshipSocket', 'xFormSocket']:
                continue
            elif hasattr(node_socket, "default_value"):
                if (value := get_socket_value(node_socket)) is not None:
                    nc.parameters[key] = value
                else:
                    raise RuntimeError(wrapRed("No value found for " + nc.__repr__() + " when filling out node parameters for " + np.name + "::"+node_socket.name))
            else:
                pass



def evaluate_input(node_container, input_name):
    if not (node_container.inputs.get(input_name)):
        # just return the parameter, there is no socket associated
        return node_container.parameters.get(input_name)
    trace = trace_single_line(node_container, input_name)
    # this should give a key error if there is a problem
    #  it is NOT handled here because it should NOT happen
    prop = trace[0][-1].parameters[trace[1].name]
    return prop

def check_for_driver(node_container, input_name, index = None):
    prop = evaluate_input(node_container, input_name)
    if (index is not None):
        prop = prop[index]
    # This should work, even for None.
    return (prop.__class__.__name__ == 'MantisDriver')

def trace_node_lines(node_container):
    """ Tells the depth of a node within the node tree. """
    node_lines = []
    if hasattr(node_container, "inputs"):
        for key, socket in node_container.inputs.items():
            # Recrusive search through the tree.
            #  * checc each relevant input socket in the node
            #  * for EACH input, find the node it's connected to
            #    * repeat from here until you get all the lines
            if ( ( key in ["Relationship", "Parent", "Input Relationship", "Target"])
                          and (socket.is_connected) ):
                # it is necesary to check the key because of Link nodes,
                #   which don't really traverse like normal.
                # TODO: see if I can refactor this to make it traverse
                other = socket.from_node
                if (other):
                    other_lines = trace_node_lines(other)
                    if not other_lines:
                        node_lines.append([other])
                    for line in other_lines:
                        node_lines.append( [other] + line )
    return node_lines


# TODO: modify this to work with multi-input nodes
def trace_single_line(node_container, input_name):
    # DO: refactor this for new link class
    """Traces a line to its input."""
    nodes = [node_container]
    if hasattr(node_container, "inputs"):
        # Trace a single line
        if (socket := node_container.inputs.get(input_name) ):
            while (socket.is_linked):
                link = socket.links[0] # inputs can only get one link.
                other = link.from_node.outputs.get(link.from_socket)
                if (other):
                    socket = other
                    if socket.can_traverse:
                        socket = socket.traverse_target
                        nodes.append(socket.node)
                    else: # this is an output.
                        nodes.append(socket.node)
                        break
                else:
                    break
    return nodes, socket


# this is same as the other, just flip from/to and in/out
def trace_single_line_up(node_container, output_name,):
    """ Tells the depth of a node within the node tree. """
    nodes = [node_container]
    if hasattr(node_container, "outputs"):
        # Trace a single line
        if (socket := node_container.outputs.get(output_name) ):
            while (socket.is_linked):
                # This is bad, but it's efficient for nodes that only expect
                #  one path along the given line
                link = socket.links[0]
                other = link.to_node.inputs.get(link.to_socket)
                if (other):
                    socket = other
                    if socket.can_traverse:
                        socket = socket.traverse_target
                        nodes.append(socket.node)
                    else: # this is an input.
                        nodes.append(socket.node)
                        break
                else:
                    break
    return nodes, socket

def trace_all_lines_up(nc, output_name):
    copy_items = {}
    for item in dir(nc):
        if "__" not in item:
            copy_items[item]=getattr(nc, item)
    # we want to copy it, BUT:
    copy_items["outputs"]:{output_name:nc.outputs[output_name]}
    # override outputs with just the one we care about.
    
    check_me = type('', (object,), copy_items)
    return get_depth_lines(check_me)[1]




def get_depth_lines(root):
    path, seek, nc_path = [0,], root, [root,]
    lines, nc_paths = {}, {}
    nc_len = len (root.hierarchy_connections)-1
    curheight=0
    while (path[0] <= nc_len):
        nc_path.append(nc_path[-1].hierarchy_connections[path[-1]])
        if (not (node_lines  := lines.get(nc_path[-1].signature, None))):
            node_lines = lines[nc_path[-1].signature] = set()
        if (not (node_paths  := nc_paths.get(nc_path[-1].signature, None))):
            node_paths = nc_paths[nc_path[-1].signature] = set()
        node_lines.add(tuple(path)); node_paths.add(tuple(nc_path))
        if nc_path[-1].hierarchy_connections:
            path.append(0); curheight+=1
        else:
            path[curheight] = path[curheight] + 1
            nc_path.pop()
            connected_nodes = nc_path[-1].hierarchy_connections
            if ( path[-1] <= len(connected_nodes)-1 ):
                seek = connected_nodes[path[-1]]
            elif curheight > 0:
                while(len(path) > 1):
                    path.pop(); curheight -= 1; path[curheight]+=1; nc_path.pop()
                    if ( (len(nc_path)>1) and path[-1] < len(nc_path[-1].hierarchy_connections) ):
                        break 
    return lines, nc_paths
    



def node_depth(lines):
    maxlen = 0
    for line in lines:
        if ( (l := len(line) ) > maxlen):
            maxlen = l
    return maxlen
    


#TODO rewrite this so it'll work with new nc_path thing
#  not a high priority bc this was debugging code for something that
#  works and has since ben refactored to work better
def printable_path(nc, path, no_wrap = False):
    string = ""; cur_nc = nc
    #DO: find out if the copy is necessary
    path = path.copy(); path.reverse()
    dummy = lambda a : a
    while path:
        wrap = dummy
        if not no_wrap:
            wrap=wrapWhite
            if (cur_nc.node_type == 'DRIVER'):
                wrap = wrapPurple
            elif (cur_nc.node_type == 'XFORM'):
                wrap = wrapOrange
            elif (cur_nc.node_type == 'LINK'):
                wrap = wrapGreen
        string += wrap(cur_nc.__repr__()) + " -> "
        try:
            cur_nc = get_from_path(cur_nc, [path.pop()] )
        except IndexError:
            string = string[:-4]
            return string
    string = string[:-4]
    return string
    # why is this not printing groups in brackets?


def get_parent(node_container, type = 'XFORM'):
    # type variable for selecting whether to get either 
    #   the parent xForm  or the inheritance node
    node_line, socket = trace_single_line(node_container, "Relationship")
    parent_nc = None
    for i in range(len(node_line)):
        # check each of the possible parent types.
        if ( (node_line[ i ].__class__.__name__ == 'LinkInherit') ):
            try: # it's the next one
                if (type == 'XFORM'):
                    return node_line[ i + 1 ]
                else: # type = 'LINK'
                    return node_line[ i ]
            except IndexError: # if there is no next one...
                return None # then there's no parent!
    return None
    # TODO!
    #
    # make this do shorthand parenting - if no parent, then use World
    #  if the parent node is skipped, use the previous node (an xForm)
    #  with default settings.
    # it is OK to generate a new, "fake" node container for this!

def get_target_and_subtarget(node_container, linkOb, input_name = "Target"):
    from bpy.types import PoseBone, Object, SplineIKConstraint, ArmatureModifier
    subtarget = ''; target = node_container.evaluate_input(input_name)
    if target:
        if (isinstance(target.bGetObject(), PoseBone)):
            subtarget = target.bGetObject().name
            target = target.bGetParentArmature()
        elif (isinstance(target.bGetObject(), Object) ):
            target = target.bGetObject()
        else:
            raise RuntimeError("Cannot interpret linkOb target!")
    
    if   (isinstance(linkOb, SplineIKConstraint)):
            if target and target.type not in ["CURVE"]:
                raise GraphError(wrapRed("Error: %s requires a Curve input, not %s" %
                                 (node_container, type(target))))
            # don't get a subtarget
            linkOb.target = target
    elif (isinstance(linkOb, ArmatureModifier)):
        linkOb.object = target
    elif (input_name == 'Target'): # this is sloppy, but it werks
        linkOb.target, linkOb.subtarget = target, subtarget
    elif (input_name == 'Pole Target'):
        linkOb.pole_target, linkOb.pole_subtarget = target, subtarget
    else: # this is really just for Armature Constraint targets...
        linkOb.target, linkOb.subtarget = target, subtarget


def setup_custom_props(nc):
    from .utilities import get_node_prototype    
    if np := get_node_prototype(nc.signature, nc.base_tree):
        for inp in np.inputs:
            if not (inp.name in nc.inputs.keys()) :
                nc.inputs[inp.name] = NodeSocket(is_input = True, name = inp.name, node = nc,)
                nc.parameters[inp.name] = None
        for out in np.outputs:
            if not (out.name in nc.outputs.keys()) :
                nc.outputs[out.name] = NodeSocket(is_input = False, name = out.name, node = nc,)
                #nc.parameters[out.name] = None
                # probably not something I want?
                # I don't think this supports in and out by the same name oof
    else:
        prRed(nc)
            
def prepare_parameters(nc):
    # some nodes add new parameters at runtime, e.g. Drivers
    # so we need to take that stuff from the node_containers that have
    #  been executed prior to this node.
    for s_name, sock in nc.inputs.items():
        if not (sock.is_linked):
            continue
        if (sock.name  in sock.links[0].from_node.parameters.keys()):
            nc.parameters[s_name] = sock.links[0].from_node.parameters[sock.name]
    # should work, this is ugly.


# TODO: this should handle sub-properties better
def evaluate_sockets(nc, c, props_sockets):
    # HACK I think I should do this in __init__
    if not hasattr(nc, "drivers"):
        nc.drivers = {}
    # end HACK
    for prop, (sock, default) in props_sockets.items():
        # c = nc.bObject
        # annoyingly, sometimes the socket is an array
        index = None
        if isinstance(sock, tuple):
            index = sock[1]; sock = sock[0]
        if (check_for_driver(nc, sock, index)):
            sock = (sock, index)
            original_prop = prop
            # TODO: deduplicate this terrible hack
            if ("." in prop): # this is a property of a property...
                sub_props = [c]
                while ("." in prop):
                    split_prop = prop.split(".")
                    prop = split_prop[1]
                    sub_prop = (split_prop[0])
                    if ("[" in sub_prop):
                        sub_prop, index = sub_prop.split("[")
                        index = int(index[0])
                        sub_props.append(getattr(sub_props[-1], sub_prop)[index])
                    else:
                        sub_props.append(getattr(sub_props[-1], sub_prop))
                setattr(sub_props[-1], prop, default)
            # this is really stupid
            else:
                setattr(c, prop, default)
            print("Adding driver %s to %s in %s" % (wrapPurple(original_prop), wrapWhite(nc.signature[-1]), wrapOrange(nc.GetxForm().bGetObject().name)))
            nc.drivers[sock] = original_prop
        else: # here we can do error checking for the socket if needed
            if (index is not None):
                setattr(c, prop, nc.evaluate_input(sock)[index])
            else:                    # 'mute' is better than 'enabled'
                # UGLY HACK          # because it is available in older
                if (prop == 'mute'): # Blenders.
                    setattr(c, prop, not nc.evaluate_input(sock))
                else:
                    setattr(c, prop, nc.evaluate_input(sock))

def finish_drivers(self):
    # gonna make this into a common function...
    drivers = []
    if not hasattr(self, "drivers"):
        return # HACK
    for driver, prop in self.drivers.items():
        #annoyingly, the driver may be a vector-driver...
        index = driver[1]; driver = driver[0]
        driver_trace = trace_single_line(self, driver)
        driver_provider, driver_socket = driver_trace[0][-1], driver_trace[1]
        if index is not None:
            driver = driver_provider.parameters[driver_socket.name][index].copy()
        else:
            driver = driver_provider.parameters[driver_socket.name].copy()
        if driver:
            # todo: deduplicate this terrible hack
            c = self.bObject # STUPID
            if ("." in prop): # this is a property of a property...
                sub_props = [c]
                while ("." in prop):
                    split_prop = prop.split(".")
                    prop = split_prop[1]
                    sub_prop = (split_prop[0])
                    if ("[" in sub_prop):
                        sub_prop, index = sub_prop.split("[")
                        index = int(index[0])
                        sub_props.append(getattr(sub_props[-1], sub_prop)[index])
                    else:
                        sub_props.append(getattr(sub_props[-1], sub_prop))
                driver["owner"] = sub_props[-1]
            else:
                driver["owner"] = self.bObject
            # prPurple("Successfully created driver for %s" % prop)
            driver["prop"] = prop
            drivers.append(driver)
        else:
            prRed("Failed to create driver for %s" % prop)
    from mantis.drivers import CreateDrivers
    CreateDrivers(drivers)


#Dummy classes for logic with node containers, they are not meant to do
#  each and every little thing the "real" Blender classes do.
class NodeLink:
    from_node = None
    from_socket = None
    to_node = None
    to_socket = None
    
    def __init__(self, from_node, from_socket, to_node, to_socket):
        self.from_node = from_node
        self.from_socket = from_socket
        self.to_node = to_node
        self.to_socket = to_socket
        self.from_node.outputs[self.from_socket].links.append(self)
        self.to_node.inputs[self.to_socket].links.append(self)
    
    def __repr__(self):
        return self.from_node.outputs[self.from_socket].__repr__() + " --> " + self.to_node.inputs[self.to_socket].__repr__()


class NodeSocket:

    @property # this is a read-only property.
    def is_linked(self):
        return len(self.links) > 0
        
    def __init__(self, is_input = False,
                 node = None, name = None,
                 traverse_target = None):
        self.can_traverse = False # to/from the other side of the parent node
        self.traverse_target = None # a reference to another Socket
        # will set in a sec
        self.node = node
        self.name = name
        self.is_input = is_input
        self.links = []
        if (traverse_target):
            set_traverse_target(traverse_target)
        
    def connect(self, node, socket):
        if  (self.is_input):
            to_node   = self.node; from_node = node
            to_socket = self.name; from_socket = socket
        else:
            from_node   = self.node; to_node   = node
            from_socket = self.name; to_socket = socket
        new_link = NodeLink(
                from_node,
                from_socket,
                to_node,
                to_socket)
        
        # if (from_node.signature[-2] in ["Chiral Identifier"] and
            # from_node.signature[-1] in ['Input_4']):
                # print(wrapRed("Connecting %s" % new_link),)
        # if (from_node.signature[-2] in ["Chiral Identifier"] and
            # from_node.signature[-1] in ['Input_3',]):
                # print(wrapRed("Connecting %s" % new_link),)
        
        return new_link
    
    def set_traverse_target(self, traverse_target):
        self.traverse_target = traverse_target
        self.can_traverse = True
        
    @property
    def is_connected(self):
        return len(self.links)>0
    
    
    def __repr__(self):
        return self.node.__repr__() + "::" + self.name