mantis/utilities/node_common.py

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

from bpy.app import version as bpy_version


# Utility to print links nicely for end user.
def print_link_list(links):
    ret_str = ""
    for link in links:
        ret_str+= link.__repr__() + '\n'
    return ret_str

# I've moved this to a separate function because the trace functions
# are kind of hard to understand and the cycle check breaks the flow
# of the code. 
def detect_cycle_in_trace(link, link_circle):
    if link.is_hierarchy:
        if link in link_circle:
            raise GraphError( "There is a cycle in the following links:\n" +
                                print_link_list(link_circle))
        link_circle.append(link)

# NOTE this is a dupe of trace_single_line, and it must be used
#      at any time before the tree is parsed (to check for cycles)
#      I don't want or need to pay the performance cost to check this everywhere
def trace_single_line_safe(mantis_node, input_name, link_index=0):
    # DO: refactor this for new link class
    """Traces a line to its input."""
    nodes = [mantis_node]
    # a test for cyclic connections
    link_circle = []
    # Trace a single line
    if (socket := mantis_node.inputs.get(input_name) ):
        while (socket.is_linked):
            link = socket.links[link_index]; link_index = 0
            # should throw an error if it is in the set
            if detect_cycle_in_trace(link, link_circle):
                return
            if (socket := link.from_node.outputs.get(link.from_socket)):
                nodes.append(socket.node)
                if socket.can_traverse:
                    socket = socket.traverse_target
                else: # this is an output.
                    break
            else:
                break
    return nodes, socket


def trace_single_line(mantis_node, input_name, link_index=0):
    # DO: refactor this for new link class
    """Traces a line to its input."""
    nodes = [mantis_node]
    # a test for cyclic connections
    if (socket := mantis_node.inputs.get(input_name) ):
        while (socket.is_linked):
            link = socket.links[link_index]; link_index = 0
            if (socket := link.from_node.outputs.get(link.from_socket)):
                nodes.append(socket.node)
                if socket.can_traverse:
                    socket = socket.traverse_target
                else: # this is an output.
                    break
            else:
                break
    return nodes, socket


def trace_single_line_up(mantis_node, output_name,):
    """I use this to get the xForm from a link node."""
    nodes = [mantis_node]
    if hasattr(mantis_node, "outputs"):
        # Trace a single line
        if (socket := mantis_node.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] # TODO: find out if this is wise.
                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_line_up_branching(mantis_node : MantisNode, output_name : str,
        break_condition : Callable = lambda mantis_node : False):
    """ Returns all leaf nodes at the ends of branching lines from an output."""
    leaf_nodes = []
    if hasattr(mantis_node, "outputs"):
        # Trace a single line
        if (socket := mantis_node.outputs.get(output_name) ):
            check_sockets={socket}
            while (check_sockets):
                # This is bad, but it's efficient for nodes that only expect
                #  one path along the given line
                socket = check_sockets.pop()
                for link in socket.links:
                    other = link.to_node.inputs.get(link.to_socket)
                    if (other):
                        socket = other
                        if break_condition(socket.node):
                            leaf_nodes.append(socket.node)
                        elif socket.is_input and socket.can_traverse:
                            check_sockets.add(socket.traverse_target)
                        else: # this is an input.
                            leaf_nodes.append(socket.node)
    return leaf_nodes


# For this, we basically just evaluate the driver once and use the result
def find_driven_default_value(mantis_node, socket_name):
    from ..misc_nodes import UtilityDriver, UtilitySwitch
    from .utilities import get_ui_node
    xf_inputs_values_var_names = []

    node_line, _last_socket = trace_single_line(mantis_node, socket_name)
    driver_node = None
    driver_type = 'SUM'
    keys = [ (0, 0), (1, 1) ]
    for driver_node in node_line:
        if driver_node.node_type == 'DRIVER': break
    else: # it's not being driven
        return mantis_node.evaluate_input(socket_name)
    
    if isinstance(driver_node, UtilityDriver):
        driver_type = driver_node.evaluate_input("Driver Type")
        # then get the variables
        for i, input in enumerate(driver_node.inputs):
            # unicode for 'a'
            var_name = chr(i - 2 + 96)
            if i == 0 : continue # the Driver Type doesn't interest us
            elif i ==1 : continue # for now. we do need the fCurve
            else:
                if not input.is_linked: continue # nothing to do here
            var_node    = input.links[0].from_node
            # NOTE: this var_node will not be prepared on time unless we sort the graph by depth
            #       and it doesn't seem to be possible to do that without breaking drivers... so
            # HACK it's safe to do this, but VERY ugly and bad.
            if not var_node.prepared:
                var_node.bPrepare()
            # end HACK
            var_data    = var_node.parameters["Driver Variable"]
            var_type    = var_data['type']
            var_channel = var_data['channel']
            prop_name = None; value = None 
            xf1 = var_node.GetxForm(index=1)
            match var_type:
                case "SINGLE_PROP":
                    prop_name = var_node.evaluate_input("Property")
                    # we'll get the property value in the next step
                case "TRANSFORMS":
                    matrix1 = xf1.evaluate_input['Matrix']
                    # Mantis ensures that these are identity at creation time
                    if "LOC" in var_channel or "ROT" in var_channel:
                            value = 0
                    else: # Scale, defaults to 1
                            value = 1
                case "ROTATION_DIFF":
                    xf2 = var_node.GetxForm(index=2)
                    matrix1 = xf1.evaluate_input['Matrix']
                    matrix2 = xf2.evaluate_input['Matrix']
                    value = matrix1.to_quaternion().rotation_difference(matrix2.to_quaternion()).angle
                case "LOC_DIFF":
                    xf2 = var_node.GetxForm(index=2)
                    matrix1 = xf1.evaluate_input['Matrix']
                    matrix2 = xf2.evaluate_input['Matrix']
                    value = (matrix1.to_translation() - matrix2.to_translation()).length
                case "CONTEXT_PROP":
                    raise NotImplementedError("Mantis Does not yet support Context Property")
                    pass
            xf_inputs_values_var_names.append((xf1, prop_name, value, var_name))

        # Now we'll get the keyframe points and sort them, too!
        if driver_node.inputs['fCurve'].is_linked:
            fCurve_node = driver_node.inputs['fCurve'].links[0].from_node
        extrapolation_mode = 'CONSTANT'
        keys = []
        if fCurve_node:
            extrapolation_mode = fCurve_node.evaluate_input("Extrapolation Mode")
            for l in fCurve_node.inputs['Keyframes'].links:
                key_node = l.from_node
                if key_node:
                    keys.append( (key_node.evaluate_input("Frame"),  key_node.evaluate_input("Value")) )
        keys = sorted( keys, key = lambda a : a[0]) # sort by the frame

    elif isinstance(driver_node, UtilitySwitch):
        # Reimplementing the GetxForm from the switch here
        # because it is faster to inline this and store the trace result
        trace = trace_single_line(driver_node, "Parameter" )
        for node in trace[0]:
            if node.node_type == 'XFORM':
                xf = node                    # reuse trace HERE
                prop_name = trace[1].name    # saves us from tracing again.
                xf_inputs_values_var_names.append(( xf, prop_name, None, "a")); break
        else:
            raise RuntimeError(f"ERROR: Cannot Set the default of the driven value {socket_name} \n"
                                f"in node {mantis_node}")
    # We're gonna assume the xForm is sorting these right
    # but we'll check the name if it is cheap
    default_values = {}
    for xf, prop_name, value, var_name in xf_inputs_values_var_names:
        i=0
        for output in xf.outputs:
            if 'xForm Out' == output.name: continue
            if output.name == prop_name: break
            i+=1
        node_line, _last_socket = trace_single_line(xf, "Custom Properties", i)
        custom_prop_node = node_line[-1]
        if value:
            default_values[var_name] = value; continue
        if prop_name:
            if custom_prop_node.evaluate_input("Name") == prop_name:
                custom_prop_type = custom_prop_node.evaluate_input("Type")
                default_prop_name = ""
                match custom_prop_type:
                    case "BOOL":
                        default_prop_name = 'Default (Bool)'
                    case "FLOAT":
                        default_prop_name = 'Default (Float)'
                    case "INT":
                        default_prop_name = 'Default (Int)'
                    case _:
                        raise RuntimeError(f"ERROR: Cannot Set the default of the driven value {socket_name} \n"
                                           f"in node {mantis_node}"
                                           f"Reason: Custom Property {prop_name} is not a valid type." )
                default_values[var_name] = float(custom_prop_node.evaluate_input(default_prop_name))
            else:
                raise RuntimeError(f"ERROR: failed to get proeprty {prop_name} for node {xf}, \n"
                                   f"cannot Set the default of the driven value {socket_name} \n"
                                   f"in node {mantis_node}"
                                    "Reason: Cannot Find Custom Property Node")
    frame = 0.0
    if default_values:
        match driver_type:
            case "AVERAGE":
                frame = sum(default_values.values())/len(default_values)
            case "SUM":
                frame = sum(default_values.values())
            case "MIN":
                frame = min(default_values.values())
            case "MAX":
                frame = max(default_values.values())
            case "SCRIPTED":
                raise NotImplementedError (f"ERROR: Cannot Set the default of the driven value {socket_name} \n"
                                        f"in node {mantis_node} \n"
                                            "Reason: Driver Expressions are not yet implemented.")
    # Now we apply the fCurve
    # For now, We'll assume LINEAR keys TODO: come back to this during Pose update!
    if len(keys) >= 2:
        previous_frame = None
        next_frame = None
        previous_value = None
        next_value = None
        for key in keys:
            if key[0] <= frame: previous_frame, previous_value = key
            if key[0] >= frame: next_frame, next_value = key
            if previous_frame is not None and next_frame is not None: break
        else: # if our frame is out-of-bounds we need a slope to extrapolate on.
            if previous_frame and not next_frame:
                next_frame, next_value = previous_frame, previous_value
                if extrapolation_mode == 'LINEAR':
                    previous_frame, previous_value = keys[-2] # we need the second last key to extrapolate
            elif next_frame and not previous_frame:
                previous_frame, previous_value = next_frame, next_value
                if extrapolation_mode == 'LINEAR':
                    next_frame, next_value = keys[1] # we need the second key to extrapolate
            else:   
                raise RuntimeError (f"ERROR: Cannot Set the default of the driven value {socket_name} \n"
                                    f"in node {mantis_node} \n"
                                    "Reason: Cannot place value on driver fCurve.")
        if previous_frame == next_frame: # the frame is exactly on a key
            return previous_value
        # this is just a simple lerp because we are assuming the keys are linear for now
        result = previous_value + ((frame - previous_frame) / (next_frame - previous_frame)) * (next_value - previous_value)
    elif len(keys) ==1:
        result = key[1] # flat curve, take the value of the only key
    else:
        raise RuntimeError (f"ERROR: Cannot Set the default of the driven value {socket_name} \n"
                            f"in node {mantis_node} \n"
                             "Reason: Invalid fCurve has 0 keyframes")
    return result
        
# this function exists to walk back to the first ancestor node of XFORM or LINK without
#  breaking the traverse behaviour which I still need for various reasons.
def get_parent_xForm_info(mantis_node, socket_name = 'Relationship'):
    from ..mantis_dataclasses import xForm_info
    # we re-implement the trace logic here for performance.
    # because this kind of connection has traversal (for Links to find their children), we must
    # simply seek back to the first available xForm or Link and get its xForm info.\
    parent_node = None
    if (socket := mantis_node.inputs.get(socket_name) ):
        while (socket.is_linked):
            link = socket.links[0]
            if (socket := link.from_node.outputs.get(link.from_socket)):
                if socket.node.node_type in ['LINK', 'XFORM']:
                    parent_node = socket.node; break
                if socket.can_traverse:
                    socket = socket.traverse_target
                else: # this is an output.
                    break
            else:
                break
    parent_xForm_info = None
    if parent_node:
        if parent_node.node_type == 'XFORM':
            parent_xForm_info = parent_node.parameters['xForm Out']
        else:
            if 'Output Relationship' in parent_node.parameters.keys():
                parent_xForm_info = parent_node.parameters['Output Relationship']
            else: # TODO refactor Inherit Node to have the same output as the other link nodes.
                parent_xForm_info = parent_node.parameters['Inheritance']
    if parent_xForm_info is None: parent_xForm_info = xForm_info() # default to empty
    return parent_xForm_info

def setup_custom_property_inputs_outputs(mantis_node):
    from .utilities import get_ui_node
    if mantis_node.signature[0] == 'SCHEMA_AUTOGENERATED':
        from ..base_definitions import custom_props_types
        if mantis_node.__class__.__name__ not in custom_props_types:
            raise RuntimeError(wrapRed(f"Custom Properties not allowed for class {mantis_node.__class__.__name__}"))
        return
    else:
        ui_node = get_ui_node(mantis_node.ui_signature, mantis_node.base_tree)
    if ui_node:
        setup_custom_props_from_np(mantis_node, ui_node)
    else:
        prRed(f"Failed to setup custom properties for: {mantis_node.ui_signature}")

def setup_custom_props_from_np(mantis_node, ui_node):
    for inp in ui_node.inputs:
        if inp.identifier == "__extend__": continue
        if not (inp.name in mantis_node.inputs.keys()) :
            socket = NodeSocket(is_input = True, name = inp.name, node = mantis_node,)
            mantis_node.inputs[inp.name] = socket
            mantis_node.parameters[inp.name] = None
            for attr_name in ["min", "max", "soft_min", "soft_max", "description"]:
                try:
                    setattr(socket, attr_name, getattr(inp, attr_name))
                except AttributeError:
                    pass
    for out in ui_node.outputs:
        if out.identifier == "__extend__": continue
        if not (out.name in mantis_node.outputs.keys()) :
            mantis_node.outputs[out.name] = NodeSocket(is_input = False, name = out.name, node = mantis_node,)
            
def prepare_parameters(mantis_node):
    # some nodes add new parameters at runtime, e.g. Drivers
    # so we need to take that stuff from the mantis_nodes that have
    #  been executed prior to this node.
    for s_name, sock in mantis_node.inputs.items():
        if not (sock.is_linked):
            continue
        if (sock.name  in sock.links[0].from_node.parameters.keys()):
            mantis_node.parameters[s_name] = sock.links[0].from_node.parameters[sock.name]
    # should work, this is ugly.

def check_for_driver(mantis_node, input_name, value_index = None, array_index = 0, combo_property=False):
    prop = mantis_node.evaluate_input(input_name, array_index)
    if (value_index is not None):
        try:
            prop = prop[value_index]
        except KeyError: # We'll try to use the driver for all three channels.
            pass
        # TODO an Index or Key error here should have a reasonable error message
    return (prop.__class__.__name__ == 'MantisDriver')

def print_driver_item(mantis_node, b_object, original_prop):
    if mantis_node.node_type in ['LINK',]:
        printname  = wrapOrange(b_object.id_data.name)
    elif mantis_node.node_type in ['XFORM',]:
        printname  = wrapOrange(mantis_node.bGetObject().name)
    else:
        printname = wrapOrange(mantis_node)
    print("Adding driver %s to %s in %s" % (wrapPurple(original_prop), wrapWhite(mantis_node.signature[-1]), printname))

# HACK This is the most hideous voodoo black magic fucntion in the entire program
#      but everything depends on it and it works 
def evaluate_sockets(mantis_node, b_object, props_sockets,):
    # this is neccesary because some things use dict properties for dynamic properties and setattr doesn't work
    def safe_setattr(ob, att_name, val):
        if ob.__class__.__name__ in ["NodesModifier"]:
            ob[att_name]=val
        elif b_object.__class__.__name__ in ["Key"]:
            if not val: val=0
            ob.key_blocks[att_name].value=val
        elif "]." in att_name:
            # it is of the form prop[int].prop2
            prop=att_name.split('[')[0]
            prop1=att_name.split('.')[1]
            get_index = int(att_name.split('[')[1][0])
            setattr(getattr(b_object, prop)[get_index], prop1, val)
        else:
            try:
                setattr(ob, att_name, val)
            except Exception as e:
                print ("ob, att_name, val")
                prRed(ob, att_name, val)
                raise e
    for prop, filler_item in props_sockets.items():
        if isinstance(filler_item, tuple):
            raise NotImplementedError(mantis_node)
        sock = filler_item.socket_name
        default = filler_item.default_value
        value_index = filler_item.value_index
        array_index = filler_item.array_index
        combo_property = filler_item.combo_property
        driver_key = (sock, value_index, array_index)
        if (check_for_driver(mantis_node, sock, value_index, array_index, combo_property)):
            original_prop = prop
            # TODO: deduplicate this terrible hack
            if ("." in prop) and not b_object.__class__.__name__ in ["Key"]: # this is a property of a property...
                sub_props = [b_object]
                while ("." in prop):
                    split_prop = prop.split(".")
                    prop = split_prop[1]
                    sub_prop = (split_prop[0])
                    if ("[" in sub_prop):
                        sub_prop, sub_index = sub_prop.split("[")
                        sub_prop_index = int(sub_index[0])
                        sub_props.append(getattr(sub_props[-1], sub_prop)[sub_prop_index])
                    else:
                        sub_props.append(getattr(sub_props[-1], sub_prop))
                safe_setattr(sub_props[-1], prop, default)
            # this is really stupid
            else:
                safe_setattr(b_object, prop, default)
            print_driver_item( mantis_node, b_object, original_prop)
            if b_object.__class__.__name__ in ["NodesModifier"]:
                mantis_node.drivers[driver_key] = "[\""+original_prop+"\"]" # lol. It is a dict element not a "true" property
            elif b_object.__class__.__name__ in ["Key"]:
                mantis_node.drivers[driver_key] = "key_blocks[\""+original_prop+"\"].value"
            else:
                mantis_node.drivers[driver_key] = original_prop
        else: # here we can do error checking for the socket if needed
            if (value_index is not None):
                safe_setattr(b_object, prop, mantis_node.evaluate_input(sock)[value_index])
            else:                    # 'mute' is better than 'enabled'
                # UGLY HACK          # because it is available in older
                if (prop == 'mute'): # Blenders.
                    safe_setattr(b_object, prop, not bool(mantis_node.evaluate_input(sock)))
                elif (prop == 'hide'): # this will not cast it for me, annoying.
                    safe_setattr(b_object, prop, bool(mantis_node.evaluate_input(sock)))
                else:
                    try:
                    # just wanna see if we are dealing with some collection
                    # check hasattr in case it is one of those ["such-and-such"] props, and ignore those
                        if hasattr(b_object, prop) and (not isinstance(getattr(b_object, prop), str)) and hasattr(getattr(b_object, prop), "__getitem__"):
                            # prGreen("Doing the thing")
                            for val_index, value in enumerate(mantis_node.evaluate_input(sock)):
                                # assume this will work, both because val should have the right number of elements, and because this should be the right data type.
                                from .drivers import MantisDriver
                                if isinstance(value, MantisDriver):
                                    getattr(b_object,prop)[val_index] = default[val_index]
                                    print("Adding driver %s to %s in %s" % (wrapPurple(prop), wrapWhite(mantis_node.signature[-1]), mantis_node))
                                    try:
                                        mantis_node.drivers[sock].append((prop, val_index))
                                    except:
                                        mantis_node.drivers[sock] = [(prop, val_index)]
                                else:
                                    getattr(b_object,prop)[val_index] =  value
                        else:
                            safe_setattr(b_object, prop, mantis_node.evaluate_input(sock))
                    except Exception as e:
                        prRed(b_object, mantis_node, prop, sock, mantis_node.evaluate_input(sock))
                        raise e

def finish_driver(mantis_node, b_object, driver_item, prop):
    try:
        driver_sock = driver_item[0]; index = driver_item[1]; array_index =driver_item[2]
    except Exception as e:
        print (wrapGreen(mantis_node), wrapOrange(b_object), wrapRed(driver_item), wrapWhite(prop),)
        raise e
    driver_trace = trace_single_line(mantis_node, driver_sock, array_index)
    driver_provider, driver_socket = driver_trace[0][-1], driver_trace[1]
    
    combo_property = False
    # try to detect combo properties
    if mantis_node.socket_templates:
        for st in mantis_node.socket_templates:
            if st.name == driver_sock:
                if isinstance(st.blender_property, tuple):
                    combo_property = True
                break

    if index is not None:
        try:
            driver = driver_provider.parameters[driver_socket.name][index].copy()
        except KeyError as e: # try single driver, it should still work fine.
            if combo_property:
                driver = driver_provider.parameters[driver_socket.name].copy()
            else:
                raise e
        # this is harmless and necessary for the weird ones where the property is a vector too
        if not combo_property:
            driver["ind"] = index
    else:
        driver = driver_provider.parameters[driver_socket.name].copy()
    if driver:
        socket_template = None
        if mantis_node.socket_templates:
            for socket_template in mantis_node.socket_templates:
                if socket_template.name == driver_sock: break
        if socket_template:
            if socket_template.is_animatable == False:
                raise RuntimeError(f"Driver is not allowed for {driver_sock} in node {mantis_node}")
        if ("." in prop) and mantis_node.__class__.__name__ != "DeformerMorphTargetDeform": # this is a property of a property...
            sub_props = [b_object]
            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]
        elif mantis_node.node_type in ['XFORM',] and mantis_node.__class__.__name__ in ['xFormBone']:
            # TODO: I really shouldn't have to hardcode this. Look into better solutions.
            # NOTE: VERSIONING in Blender 5.0 we have to use pose bone for this, too.
            # default to pose bone, it's usually what we want
            bone_col = mantis_node.bGetParentArmature().pose.bones
            from ..xForm_socket_templates import attributes_shared_by_bone_and_pose_bone
            if socket_template.prop_owner_group == 'bone':
                bone_col = mantis_node.bGetParentArmature().data.bones
            # this is a little bit clumsy. TODO clean up here
            if prop in attributes_shared_by_bone_and_pose_bone:
                # we want to use the pose bone for this...only if it has a driver
                bone_col = mantis_node.bGetParentArmature().pose.bones
            driver["owner"] = bone_col[b_object] # we use "unsafe" brackets instead of get() because we want to see any errors that occur
        # HACK having special cases here is indicitave of a deeper problem that should be refactored
        elif mantis_node.__class__.__name__ in ['xFormCurvePin'] and \
                      prop in ['offset_factor', 'forward_axis', 'up_axis']:
                driver["owner"] = b_object.constraints['Curve Pin']
        else:
            driver["owner"] = b_object
        driver["prop"] = prop
        return driver
    else:
        prOrange("Provider", driver_provider)
        prGreen("socket", driver_socket)
        print (index)
        prPurple(driver_provider.parameters[driver_socket.name])
        prRed("Failed to create driver for %s" % prop)
        return None

def finish_drivers(mantis_node):
    drivers = []
    if not hasattr(mantis_node, "drivers"):
        return # HACK
    for driver_item, prop in mantis_node.drivers.items():
        b_objects = [mantis_node.bObject]
        if mantis_node.node_type == 'LINK':
            b_objects = mantis_node.bObject # this is already a list
        for b_object in b_objects:
            # Mutating the driver_item here is the WORST THING I HAVE EVER SEEN
            # HACK UGLY BUG TERRIBLE BAD CODE
            if isinstance(prop, list):
                for sub_item in prop:
                        drivers.append(finish_driver(mantis_node, b_object, (driver_item, sub_item[1], 0), sub_item[0]))
                else:
                    drivers.append(finish_driver(mantis_node, b_object, (driver_item, sub_item[1], 0), sub_item[0]))
            else:
                drivers.append(finish_driver(mantis_node, b_object, driver_item, prop))
    from .drivers import CreateDrivers
    CreateDrivers(drivers)