mantis/deformer_containers.py

from .node_container_common import *
from .xForm_containers import xFormGeometryObject
from .misc_nodes import InputExistingGeometryObject
from .base_definitions import MantisNode, MantisSocketTemplate
from .utilities import (prRed, prGreen, prPurple, prWhite, prOrange,
                        wrapRed, wrapGreen, wrapPurple, wrapWhite,
                        wrapOrange,)

from .deformer_socket_templates import *

from bpy.types import NodeTree

def TellClasses():
             
    return [ 
             DeformerArmature,
             DeformerHook,
             DeformerMorphTarget,
             DeformerMorphTargetDeform,
           ]


def trace_xForm_back(nc, socket):
    from .xForm_containers import xFormGeometryObject
    from .misc_nodes import InputExistingGeometryObject
    from bpy.types import Object
    if (trace := trace_single_line(nc, socket)[0] ) :
        for i in range(len(trace)): # have to look in reverse, actually
            if ( isinstance(trace[ i ], xFormGeometryObject ) ) or ( isinstance(trace[ i ], InputExistingGeometryObject ) ):
                return trace[ i ].bGetObject()
        raise GraphError(wrapRed(f"No other object found for {nc}."))


# semi-duplicated from link_containers
def GetxForm(nc):
    trace = trace_single_line_up(nc, "Deformer")
    for node in trace[0]:
        if (node.__class__ in [xFormGeometryObject, InputExistingGeometryObject]):
            return node
    raise GraphError("%s is not connected to a downstream xForm" % nc)

class MantisDeformerNode(MantisNode):
    def __init__(self, signature : tuple,
                 base_tree : NodeTree,
                 socket_templates : list[MantisSocketTemplate]=[]):
        super().__init__(signature, base_tree, socket_templates)
        self.node_type = 'LINK'
        self.prepared = True
    # we need evaluate_input to have the same behaviour as links.
    def evaluate_input(self, input_name, index=0):
        if ('Target' in input_name):
            socket = self.inputs.get(input_name)
            if socket.is_linked:
                return socket.links[0].from_node
            return None
            
        else:
            return super().evaluate_input(input_name, index)

class DeformerArmature(MantisDeformerNode):
    '''A node representing an armature deformer'''

    def __init__(self, signature, base_tree):
        super().__init__(signature, base_tree)
        inputs = [
            "Input Relationship",
            "Armature Object",
            "Blend Vertex Group",
            "Invert Vertex Group",
            "Preserve Volume",
            "Use Multi Modifier",
            "Use Envelopes",
            "Use Vertex Groups",
            "Skinning Method",
            "Deformer",
            "Copy Skin Weights From"
        ]
        outputs = [
          "Deformer"
        ]
        self.outputs.init_sockets(outputs)
        self.inputs.init_sockets(inputs)
        self.init_parameters(additional_parameters={"Name":None})
        self.set_traverse([("Deformer", "Deformer")])
        self.node_type = "LINK"
        self.prepared = True


    def GetxForm(self, socket="Deformer"):
        if socket == "Deformer":
            return GetxForm(self)
        else:
            trace_xForm_back(self, socket)
    
    # DUPLICATED FROM xForm_containers::xFormBone 
    # DEDUP HACK HACK HACK HACK HACK
    def bGetParentArmature(self):
        from .xForm_containers import xFormArmature
        from .misc_nodes import InputExistingGeometryObject
        from bpy.types import Object
        if (trace := trace_single_line(self, "Armature Object")[0] ) :
            for i in range(len(trace)):
                # have to look in reverse, actually
                if ( isinstance(trace[ i ], xFormArmature ) ):
                    return trace[ i ].bGetObject()
                elif ( isinstance(trace[i], InputExistingGeometryObject)):
                    if (ob := trace[i].bGetObject()).type == "ARMATURE":
                        return ob
        raise RuntimeError(f"Cannot find armature for node {self}")
        return None
        #should do the trick...

    def bExecute(self, bContext = None,):
        self.executed = True
    
    def initialize_vgroups(self,):
        ob = self.GetxForm().bGetObject()
        armOb = self.bGetParentArmature()
        for b in armOb.data.bones:
            if b.use_deform == False:
                continue
            vg = ob.vertex_groups.get(b.name)
            if not vg:
                vg = ob.vertex_groups.new(name=b.name)
                num_verts = len(ob.data.vertices)
                vg.add(range(num_verts), 0, 'REPLACE')
    
    def copy_weights(self):
        # we'll use modifiers for this, maybe use GN for it in the future tho
        import bpy
        ob = self.GetxForm().bGetObject()
        try:
            copy_from = self.GetxForm(socket="Copy Skin Weights From")
        except GraphError:
            copy_from = None
            prRed(f"No object found for copying weights in {self}, continuing anyway.")
        m = ob.modifiers.new(type="DATA_TRANSFER", name="Mantis_temp_data_transfer")
        m.object = None; m.use_vert_data = True
        m.data_types_verts = {'VGROUP_WEIGHTS'}
        m.vert_mapping = 'POLYINTERP_NEAREST'
        m.layers_vgroup_select_src = 'ALL'
        m.layers_vgroup_select_dst = 'NAME'
        m.object = copy_from
        # m.use_object_transform = False # testing reveals that this is undesirable - since the objects may not have their transforms applied.
        ob.modifiers.move(len(ob.modifiers)-1, 0)

        # ob.data = ob.data.copy()
        if False: #MAYBE the mouse needs to be in the 3D viewport, no idea how to set this in an override
            # TODO: figure out how to apply this, context is incorrect because armature is still in pose mode
            original_active = bpy.context.active_object
            original_mode = original_active.mode
            bpy.ops.object.mode_set(mode='OBJECT')
            with bpy.context.temp_override(**{'active_object':ob, 'selected_objects':[ob, copy_from]}):
                # bpy.ops.object.datalayout_transfer(modifier=m.name) # note: this operator is used by the modifier or stand-alone in the UI
                # the poll for this operator is defined in blender/source/blender/editors/object/object_data_transfer.cc
                # and blender/source/blender/editors/object/object_modifier.cc
                # bpy.ops.object.modifier_apply(modifier=m.name, single_user=True)
                bpy.ops.object.datalayout_transfer(data_type='VGROUP_WEIGHTS')
                bpy.ops.object.data_transfer(data_type='VGROUP_WEIGHTS')
            bpy.ops.object.mode_set(mode=original_mode)
         
    def bFinalize(self, bContext=None):
        prGreen("Executing Armature Deform Node")
        mod_name = self.evaluate_input("Name")
        d = self.GetxForm().bGetObject().modifiers.new(mod_name, type='ARMATURE')
        if d is None:
            raise RuntimeError(f"Modifier was not created in node {self} -- the object is invalid.")
        self.bObject = d
        d.object = self.bGetParentArmature()
        props_sockets = {
        'vertex_group'               : ("Blend Vertex Group", ""),
        'invert_vertex_group'        : ("Invert Vertex Group", ""),
        'use_deform_preserve_volume' : ("Preserve Volume", False),
        'use_multi_modifier'         : ("Use Multi Modifier", False),
        'use_bone_envelopes'         : ("Use Envelopes", False),
        'use_vertex_groups'          : ("Use Vertex Groups", False),
        }
        evaluate_sockets(self, d, props_sockets)
        #
        if (skin_method := self.evaluate_input("Skinning Method")) == "AUTOMATIC_HEAT":
            # This is bad and leads to somewhat unpredictable
            #  behaviour, e.g. what object will be selected? What mode?
            # also bpy.ops is ugly and prone to error when used in
            #  scripts. I don't intend to use bpy.ops when I can avoid it.
            import bpy
            self.initialize_vgroups()
            bContext.view_layer.depsgraph.update()
            ob = self.GetxForm().bGetObject()
            armOb = self.bGetParentArmature()
            deform_bones = []
            for pb in armOb.pose.bones:
                if pb.bone.use_deform == True:
                    deform_bones.append(pb)
            
            context_override = {
                                  'active_object':ob,
                                  'selected_objects':[ob, armOb],
                                  'active_pose_bone':deform_bones[0],
                                  'selected_pose_bones':deform_bones,}
            #
            # with bContext.temp_override(**{'active_object':armOb}):
            #     bpy.ops.object.mode_set(mode='POSE')
            #     bpy.ops.pose.select_all(action='SELECT')
            for b in armOb.data.bones:
                b.select = True
            with bContext.temp_override(**context_override):
                bpy.ops.paint.weight_paint_toggle()
                bpy.ops.paint.weight_from_bones(type='AUTOMATIC')
                bpy.ops.paint.weight_paint_toggle()
            for b in armOb.data.bones:
                b.select = False
                #
            # with bContext.temp_override(**{'active_object':armOb}):
            #     bpy.ops.object.mode_set(mode='POSE')
            #     bpy.ops.pose.select_all(action='DESELECT') 
            #     bpy.ops.object.mode_set(mode='OBJECT')
            # TODO: modify Blender to make this available as a Python API function.
        elif skin_method == "EXISTING_GROUPS":
            pass
        elif skin_method == "COPY_FROM_OBJECT":
            self.initialize_vgroups()
            self.copy_weights()


class DeformerHook(MantisDeformerNode):
    '''A node representing a hook deformer'''

    def __init__(self, signature, base_tree):
        super().__init__(signature, base_tree, HookSockets)
        # now set up the traverse target...
        self.init_parameters(additional_parameters={"Name":None})
        self.set_traverse([("Deformer", "Deformer")])
        self.prepared = True

    def driver_for_radius(self, object, hook, index, influence, bezier=True):
        """ Creates a driver to control the radius of a curve point with the hook."""
        from bpy.types import Bone, PoseBone
        var_template = {"owner":hook,
                        "name":"a",
                        "type":"TRANSFORMS",
                        "space":'WORLD_SPACE',
                        "channel":'SCALE_X',}
        var1_template = {"owner":hook.id_data,
                        "name":"b",
                        "type":"TRANSFORMS",
                        "space":'WORLD_SPACE',
                        "channel":'SCALE_X',}
        keys_template = [{"co":(0,0),
                          "interpolation": "LINEAR",
                          "type":"KEYFRAME",},
                         {"co":(1,influence),
                          "interpolation": "LINEAR",
                          "type":"KEYFRAME",},]
        if bezier:
            owner=object.data.splines[0].bezier_points
        else:
            owner=object.data.splines[0].points
        driver = {
            "owner":owner[index],
            "prop":"radius",
            "ind":-1,
            "extrapolation":"LINEAR",
            "type":"AVERAGE",
            "vars":[],
            "keys":keys_template,
        }
        if isinstance(hook, (Bone, PoseBone)):
            driver['type']='SCRIPTED'
            driver['expression']="(((1/b)*a)+((1/b_001)*a_001)+((1/b_002)*a_002))/3"
        from .drivers import CreateDrivers
        axes='XYZ'
        for i in range(3):
            var = var_template.copy()
            var["channel"]="SCALE_"+axes[i]
            driver["vars"].append(var)
            if isinstance(hook, (Bone, PoseBone)):
                var1=var1_template.copy()
                var1['channel']="SCALE_"+axes[i]
                driver['vars'].append(var1)
        CreateDrivers([driver])

    def GetxForm(self, socket="Deformer"):
        if socket == "Deformer":
            return GetxForm(self)
        else:
            trace_xForm_back(self, socket)
            
    def bExecute(self, bContext = None,):
        self.executed = True

    def bFinalize(self, bContext=None):
        from bpy.types import Bone, PoseBone, Object
        prGreen(f"Executing Hook Deform Node: {self}")
        mod_name = self.evaluate_input("Name")
        affect_radius = self.evaluate_input("Affect Curve Radius")
        auto_bezier = self.evaluate_input("Auto-Bezier")
        target_node = self.evaluate_input('Hook Target')
        target = target_node.bGetObject(); subtarget = ""
        props_sockets = self.gen_property_socket_map()
        if isinstance(target, Bone) or isinstance(target, PoseBone):
            subtarget = target.name; target = target.id_data
        ob=self.GetxForm().bGetObject()
        reuse = False
        for m in ob.modifiers:
            if  m.type == 'HOOK' and m.object == target and m.subtarget == subtarget:
                if self.evaluate_input("Influence") != m.strength:
                    continue # make a new modifier so they can have different strengths
                if ob.animation_data: # this can be None
                    drivers = ob.animation_data.drivers
                    for k in props_sockets.keys():
                        if driver := drivers.find(k):
                            # TODO: I should check to see if the drivers are the same...
                            break # continue searching for an equivalent modifier
                    else: # There was no driver - use this one.
                        d = m; reuse = True; break
                else: # use this one, there can't be drivers without animation_data.
                    d = m; reuse = True; break
        else:
            d = ob.modifiers.new(mod_name, type='HOOK')
            if d is None:
                raise RuntimeError(f"Modifier was not created in node {self} -- the object is invalid.")
        self.bObject = d
        self.get_target_and_subtarget(d, input_name="Hook Target")
        vertices_used=[]
        if reuse: # Get the verts in the list... filter out all the unneeded 0's
            vertices_used = list(d.vertex_indices)
            include_0 = 0 in vertices_used
            vertices_used = list(filter(lambda a : a != 0, vertices_used))
            if include_0: vertices_used.append(0)
        # now we add the selected vertex to the list, too
        vertex = self.evaluate_input("Curve Point Index")
        if ob.type == 'CURVE' and ob.data.splines[0].type == 'BEZIER' and auto_bezier:
            if affect_radius:
                self.driver_for_radius(ob, target_node.bGetObject(), vertex, d.strength)
            vertex*=3
            vertices_used.extend([vertex, vertex+1, vertex+2])
        else:
            vertices_used.append(vertex)
        # if we have a curve and it is NOT using auto-bezier for the verts..
        if ob.type == 'CURVE' and ob.data.splines[0].type == 'BEZIER' and affect_radius and not auto_bezier:
            print (f"WARN: {self}: \"Affect Radius\" may not behave as expected"
                    " when used on Bezier curves without Auto-Bezier")
            #bezier point starts at 1, and then every third vert, so 4, 7, 10...
            if vertex%3==1:
                self.driver_for_radius(ob, target_node.bGetObject(), vertex, d.strength)
        if ob.type == 'CURVE' and ob.data.splines[0].type != 'BEZIER' and \
                    affect_radius:
            self.driver_for_radius(ob, target_node.bGetObject(), vertex, d.strength, bezier=False)
            
        d.vertex_indices_set(vertices_used)
        evaluate_sockets(self, d, props_sockets)
        finish_drivers(self)
        # todo: this node should be able to take many indices in the future.
        # Also: I have a Geometry Nodes implementation of this I can use... maybe...


class DeformerMorphTarget(MantisDeformerNode):
    '''A node representing an armature deformer'''
    def __init__(self, signature, base_tree):
        super().__init__(signature, base_tree)
        inputs = [
            "Relative to",
            "Object",
            "Deformer",
            "Vertex Group",
        ]
        outputs = [
          "Deformer",
          "Morph Target",
        ]
        # now set up the traverse target...
        self.outputs.init_sockets(outputs)
        self.inputs.init_sockets(inputs)
        self.init_parameters(additional_parameters={"Name":None})
        self.set_traverse([("Deformer", "Deformer")])
        self.node_type = "LINK"
        self.prepared = True
    
    def GetxForm(self, trace_input="Object"):
        trace = trace_single_line(self, trace_input)
        for node in trace[0]:
            if (node.__class__ in [xFormGeometryObject, InputExistingGeometryObject]):
                return node
        raise GraphError("%s is not connected to an upstream xForm" % self)


    def bExecute(self, bContext = None,):
        prGreen("Executing Morph Target Node")

        ob = None; relative = None
        # do NOT check if the object exists here. Just let the next node deal with that.
        try:
            ob = self.GetxForm().bGetObject().name 
        except Exception as e: # this will and should throw an error if it fails
            ob = self.GetxForm().evaluate_input("Name")
        if self.inputs["Relative to"].is_linked:
            try:
                relative = self.GetxForm("Relative to").bGetObject().name
            except Exception as e: # same here
                prRed(f"Execution failed at {self}: no relative object found for morph target, despite link existing.")
                raise e

        vg = self.evaluate_input("Vertex Group") if self.evaluate_input("Vertex Group") else "" # just make sure it is a string
        
        mt={"object":ob, "vertex_group":vg, "relative_shape":relative}

        self.parameters["Morph Target"] = mt
        self.parameters["Name"] = ob # this is redundant but it's OK since accessing the mt is tedious
        self.executed = True



class DeformerMorphTargetDeform(MantisDeformerNode):
    '''A node representing an armature deformer'''

    def __init__(self, signature, base_tree):
        super().__init__(signature, base_tree)
        inputs = [
            "Deformer",
            "Use Shape Key",
            "Use Offset",
        ]
        outputs = [
          "Deformer",
        ]
        self.outputs.init_sockets(outputs)
        self.inputs.init_sockets(inputs)
        self.init_parameters(additional_parameters={"Name":None})
        self.set_traverse([("Deformer", "Deformer")])
        self.node_type = "LINK"
        self.prepared = True
        self.executed = True
        self.bObject = None
        setup_custom_props(self)

    def GetxForm(self):
        return GetxForm(self)
    
    # bpy.data.node_groups["Morph Deform.045"].nodes["Named Attribute.020"].data_type = 'FLOAT_VECTOR'
    # bpy.context.object.add_rest_position_attribute = True


    def gen_morph_target_modifier(self, context):
        # first let's see if this is a no-op
        targets = []
        for k,v in self.inputs.items():
            if "Target" in k:
                targets.append(v)
        if not targets:
            return # nothing to do here.
        
        # at this point we make the node tree
        from .geometry_node_graphgen import gen_morph_target_nodes
        m, props_sockets = gen_morph_target_nodes(
                            self.evaluate_input("Name"),
                            self.GetxForm().bGetObject(),
                            targets,
                            context,
                            use_offset=self.evaluate_input("Use Offset"))
        self.bObject = m
        evaluate_sockets(self, m, props_sockets)
        finish_drivers(self)

    def gen_shape_key(self, context): # TODO: make this a feature of the node definition that appears only when there are no prior deformers - and shows a warning!
        # TODO: the below works well, but it is quite slow. It does not seem to have better performence. Its only advantage is export to FBX.
        # there are a number of things I need to fix here
        #   - reuse shape keys if possible
        #   - figure out how to make this a lot faster
        #   - edit the xForm stuff to delete drivers from shape key ID's, since they belong to the Key, not the Object.
        # first check if we need to do anythign
        targets = []
        for k,v in self.inputs.items():
            if "Target" in k:
                targets.append(v)
        if not targets:
            return # nothing to do here
        from time import time
        start_time = time()
        from bpy import data
        xf = self.GetxForm()
        ob = xf.bGetObject()
        dg = context.view_layer.depsgraph
        dg.update()
        if xf.has_shape_keys == False:
            m = data.meshes.new_from_object(ob, preserve_all_data_layers=True, depsgraph=dg)
            ob.data = m
            ob.add_rest_position_attribute = True
            ob.shape_key_clear()
            ob.shape_key_add(name='Basis', from_mix=False)
        else:
            m = ob.data
        xf.has_shape_keys = True
        
        # using the built-in shapekey feature is actually a lot harder in terms of programming because I need...
            # min/max, as it is just not a feature of the GN version
            # to carry info from the morph target node regarding relative shapes and vertex groups and all that
            # the drivers may be more difficult to apply, too.
            # hafta make new geometry for the object and add shape keys and all that
            # the benefit to all this being exporting to game engines via .fbx

        # first make a basis shape key
        keys={}
        props_sockets={}
        for i, t in enumerate(targets):
            mt_node = t.links[0].from_node; sk_ob = mt_node.GetxForm().bGetObject()
            if sk_ob is None:
                sk_ob = data.objects.new(mt_node.evaluate_input("Name"), data.meshes.new_from_object(ob))
                context.collection.objects.link(sk_ob)
                prOrange(f"WARN: no object found for f{mt_node}; creating duplicate of current object ")
            sk_ob = dg.id_eval_get(sk_ob)
            mt_name = sk_ob.name
            vg = mt_node.parameters["Morph Target"]["vertex_group"]
            if vg: mt_name = mt_name+"."+vg
            
            sk = ob.shape_key_add(name=mt_name, from_mix=False)
            # the shapekey data is absolute point data for each vertex, in order, very simple

            # SERIOUSLY IMPORTANT:
               # use the current position of the vertex AFTER SHAPE KEYS AND DEFORMERS
               # easiest way to do it is to eval the depsgraph
               # TODO: try and get it without depsgraph update, since that may be (very) slow
            sk_m = sk_ob.data#data.meshes.new_from_object(sk_ob, preserve_all_data_layers=True, depsgraph=dg)
            for j in range(len(m.vertices)):
                sk.data[j].co = sk_m.vertices[j].co # assume they match
            # data.meshes.remove(sk_m)
            sk.vertex_group = vg
            sk.slider_min = -10
            sk.slider_max = 10
            keys[mt_name]=sk
            props_sockets[mt_name]= ("Value."+str(i).zfill(3), 1.0)
        for i, t in enumerate(targets):
            mt_node = t.links[0].from_node; sk_ob = mt_node.GetxForm().bGetObject()
            if sk_ob is None: continue
            if rel := mt_node.parameters["Morph Target"]["relative_shape"]:
                sk = keys.get(mt_name)
                sk.relative_key = keys.get(rel)
        
        self.bObject = sk.id_data
        evaluate_sockets(self, sk.id_data, props_sockets)
        finish_drivers(self)
        prWhite(f"Initializing morph target took {time() -start_time} seconds")
        

    def bFinalize(self, bContext=None):
        prGreen(f"Executing Morph Deform node {self}")
        # if there is a not a prior deformer then there should be an option to use plain 'ol shape keys
        # GN is always desirable as an option though because it can be baked & many other reasons
        use_shape_keys = self.evaluate_input("Use Shape Key")
        if use_shape_keys: # check and see if we can.
            if self.inputs.get("Deformer"): # I guess this isn't available in some node group contexts... bad. FIXME
                if (links := self.inputs["Deformer"].links):
                    if not links[0].from_node.parameters.get("Use Shape Key"):
                        use_shape_keys = False
                    elif links[0].from_node.parameters.get("Use Shape Key") == False:
                        use_shape_keys = False
        self.parameters["Use Shape Key"] = use_shape_keys
        if use_shape_keys:
            self.gen_shape_key(bContext)
        else:
            self.gen_morph_target_modifier(bContext)