mantis/xForm_definitions.py

import bpy
from .base_definitions import xFormNode
from bpy.types import Node
from .utilities import (prRed, prGreen, prPurple, prWhite,
                              prOrange,
                              wrapRed, wrapGreen, wrapPurple, wrapWhite,
                              wrapOrange,)
from .base_definitions import get_signature_from_edited_tree


def TellClasses():
    return [
        # xFormNullNode,
        xFormBoneNode,
        xFormRootNode,
        xFormArmatureNode,
        xFormGeometryObjectNode,
        ]

def default_traverse(self, socket):
    if (socket == self.outputs["xForm Out"]):
        return self.inputs["Relationship"]
    if (socket == self.inputs["Relationship"]):
        return self.outputs["xForm Out"]
    return None

# Representing an Empty or non-armature-Object
# class xFormNullNode(Node, xFormNode):
#     '''A node representing a Null node'''
#     bl_idname = 'xFormNullNode'
#     bl_label = "Null"
#     bl_icon = 'EMPTY_AXIS'

#     # === Optional Functions ===
#     def init(self, context):
#         self.inputs.new('StringSocket', "Name")
#         self.inputs.new('RelationshipSocket', "Relationship")
#         self.inputs.new('RotationOrderSocket', "Rotation Order")
#         self.inputs.new('MatrixSocket', "Matrix")
#         self.outputs.new('xFormSocket', "xForm Out")


def check_if_connected(start, end, line):
    started=False
    for path_nc in line:
        prWhite("    ", path_nc.signature)
        if path_nc.signature == start.signature:
            started = True
        elif path_nc.signature == end.signature:
            break
        if started:
            if path_nc.inputs.get("Connected"):
                if path_nc.evaluate_input("Connected") == False:
                    return False
    else:
        return False
    return True

class xFormRootNode(Node, xFormNode):
    '''A node representing the world node'''
    bl_idname = 'xFormRootNode'
    bl_label = "World Root"
    bl_icon = 'WORLD'
    initialized : bpy.props.BoolProperty(default = False)

    def init(self, context):
        self.outputs.new('RelationshipSocket', "World Out")
        self.initialized=True



# I had chat gpt flip these so they may be a little innacurate
# always visible
main_names = {
"Name":'StringSocket',
"Rotation Order":'RotationOrderSocket',
"Relationship":'RelationshipSocket',
"Matrix":'MatrixSocket',}

# IK SETTINGS
ik_names = {
"IK Stretch":'FloatFactorSocket',
"Lock IK":'BooleanThreeTupleSocket',
"IK Stiffness":'NodeSocketVector',
"Limit IK":'BooleanThreeTupleSocket',
"X Min":'NodeSocketFloatAngle',
"X Max":'NodeSocketFloatAngle',
"Y Min":'NodeSocketFloatAngle',
"Y Max":'NodeSocketFloatAngle',
"Z Min":'NodeSocketFloatAngle',
"Z Max":'NodeSocketFloatAngle',
}

#display settings
display_names = {
"Bone Collection":'BoneCollectionSocket',
"Custom Object":'xFormSocket',
"Custom Object xForm Override":'xFormSocket',
"Custom Object Scale to Bone Length":'BooleanSocket',
"Custom Object Wireframe":'BooleanSocket',
"Custom Object Scale":'VectorScaleSocket',
"Custom Object Translation":'VectorSocket',
"Custom Object Rotation":'VectorEulerSocket',
}

# deform_names
deform_names = {
"Deform":'BooleanSocket',
"Envelope Distance":'FloatPositiveSocket',
"Envelope Weight":'FloatFactorSocket',
"Envelope Multiply":'BooleanSocket',
"Envelope Head Radius":'FloatPositiveSocket',
"Envelope Tail Radius":'FloatPositiveSocket',
}

bbone_names = {
    "BBone Segments":"IntSocket", # BONE
    "BBone X Size":"FloatSocket", # BONE
    "BBone Z Size":"FloatSocket", # BONE
    # "bbone_mapping_mode":"StringSocket", <== BONE
    "BBone HQ Deformation":"BooleanSocket", # BONE bbone_mapping_mode
    "BBone X Curve-In":"FloatSocket", # BONE AND POSE
    "BBone Z Curve-In":"FloatSocket", # BONE AND POSE
    "BBone X Curve-Out":"FloatSocket", # BONE AND POSE
    "BBone Z Curve-Out":"FloatSocket", # BONE AND POSE
    "BBone Roll-In":"FloatSocket", # BONE AND POSE
    "BBone Roll-Out":"FloatSocket", # BONE AND POSE
    "BBone Inherit End Roll":"BooleanSocket", # BONE
    "BBone Scale-In":"VectorSocket", # BONE AND POSE
    "BBone Scale-Out":"VectorSocket", # BONE AND POSE
    "BBone Ease-In":"FloatSocket", # BONE AND POSE
    "BBone Ease-Out":"FloatSocket", # BONE AND POSE
    "BBone Easing":"BooleanSocket", # BONE
    "BBone Start Handle Type":"EnumBBoneHandleType", # BONE
    "BBone Custom Start Handle":"StringSocket", # BONE
    "BBone Start Handle Scale":"BooleanThreeTupleSocket", # BONE
    "BBone Start Handle Ease":"BooleanSocket", # BONE
    "BBone End Handle Type":"EnumBBoneHandleType", # BONE
    "BBone Custom End Handle":"StringSocket", # BONE
    "BBone End Handle Scale":"BooleanThreeTupleSocket", # BONE
    "BBone End Handle Ease":"BooleanSocket", # BONE

}

other_names = {
    "Lock Location":'BooleanThreeTupleSocket',
    "Lock Rotation":'BooleanThreeTupleSocket',
    "Lock Scale":'BooleanThreeTupleSocket',
    "Hide":'HideSocket',
}

from mathutils import Color
xFormColor = Color((0.093172, 0.047735, 0.028036)).from_scene_linear_to_srgb()



class xFormBoneNode(Node, xFormNode):
    '''A node representing a Bone'''
    bl_idname = 'xFormBoneNode'
    bl_label = "Bone"
    bl_icon = 'BONE_DATA'
    
    display_ik_settings : bpy.props.BoolProperty(default=False)
    display_vp_settings : bpy.props.BoolProperty(default=False)
    display_def_settings : bpy.props.BoolProperty(default=False)
    display_bb_settings : bpy.props.BoolProperty(default=False)
    socket_count : bpy.props.IntProperty()
    initialized : bpy.props.BoolProperty(default = False)
    
    def init(self, context):

        for name, sock_type in main_names.items():
            self.inputs.new(sock_type, name)

        for name, sock_type in ik_names.items():
            s = self.inputs.new(sock_type, name)
            s.hide = True

        for name, sock_type in display_names.items():
            s = self.inputs.new(sock_type, name)
            if s.name in ['Custom Object', 'Bone Collection']:
                continue
            s.hide = True
        
        for name, sock_type in deform_names.items():
            s = self.inputs.new(sock_type, name)
            if s.name == 'Deform':
                continue
            s.hide = True
        
        for name, sock_type in bbone_names.items():
            s = self.inputs.new(sock_type, name)
            if s.name == "BBone Segments":
                continue
            s.hide = True

        for name, sock_type in other_names.items():
            self.inputs.new(sock_type, name)
        # could probably simplify this further with iter_tools.chain() but meh

        self.socket_count = len(self.inputs)
        #
        self.outputs.new('xFormSocket', "xForm Out")


        # set up some defaults...
        self.inputs['Rotation Order'].default_value = "XYZ"
        self.inputs['Lock Location'].default_value[0] = True
        self.inputs['Lock Location'].default_value[1] = True
        self.inputs['Lock Location'].default_value[2] = True
        self.inputs['Lock Rotation'].default_value[0] = True
        self.inputs['Lock Rotation'].default_value[1] = True
        self.inputs['Lock Rotation'].default_value[2] = True
        self.inputs['Lock Scale'].default_value[0] = True
        self.inputs['Lock Scale'].default_value[1] = True
        self.inputs['Lock Scale'].default_value[2] = True

        # color
        self.use_custom_color = True
        self.color = xFormColor
        #
        self.initialized=True
    
    def draw_buttons(self, context, layout):
        # return
        layout.operator("mantis.add_custom_property", text='+Add Custom Parameter')
        # layout.label(text="Edit Parameter ... not implemented")
        if (len(self.inputs) >= self.socket_count):
            layout.operator("mantis.remove_custom_property", text='-Remove Custom Parameter')
        else:
            layout.label(text="")
        
    def display_update(self, parsed_tree, context):
        if context.space_data:
            nc = parsed_tree.get(get_signature_from_edited_tree(self, context))
            other_nc = None
            if len(self.inputs.get("Relationship").links)>0:
                prev_node = self.inputs.get("Relationship").links[0].from_node
                if prev_node:
                    other_nc = parsed_tree.get(get_signature_from_edited_tree(prev_node, context))
            
            if nc and other_nc:
                self.display_vp_settings = nc.inputs["Custom Object"].is_connected
                self.display_def_settings = nc.evaluate_input("Deform")
                self.display_bb_settings = nc.evaluate_input("BBone Segments") > 1
                self.display_ik_settings = False
                #
                from .node_container_common import ( trace_all_lines_up,
                                                     trace_single_line)
                #TODO find a much faster way to do this
                if False and self.id_data.do_live_update: #TODO this is extremely freaking slow
                    trace = trace_all_lines_up(nc, "xForm Out")

                    for key in trace.keys():
                        if (ik_nc:= parsed_tree.get(key)):
                            if ik_nc.__class__.__name__ in ["LinkInverseKinematics"]:
                                # if the tree is invalid? This shouldn't be necessary.
                                if ik_nc.inputs["Input Relationship"].is_connected:
                                    chain_count = ik_nc.evaluate_input("Chain Length")
                                    # if chain_count == 0: # this is wrong
                                    #     self.display_ik_settings = True
                                    # else:
                                    if ik_nc.evaluate_input("Use Tail") == False:
                                        chain_count+=1
                                    for line in trace[key]:
                                        # preprocess it to get rid of non-xForms:
                                        xForm_line=[]
                                        for path_nc in line:
                                            if path_nc == ik_nc:
                                                if ik_nc.inputs["Input Relationship"].links[0].from_node != prev_path_nc:
                                                    break # not a constraint connection
                                            if path_nc.node_type == 'XFORM':
                                                xForm_line.append(path_nc)
                                            prev_path_nc = path_nc
                                        else: # if it can finish cleaning the line of non-xForms and still has a connection
                                            if (len(xForm_line) < chain_count) or (chain_count == 0):
                                                self.display_ik_settings = True
                
                inp = nc.inputs["Relationship"]
                link = None
                if inp.is_connected:
                    link = inp.links[0]
                while(link):
                    if link.from_node.__class__.__name__ in ["LinkInverseKinematics"]:
                        self.display_ik_settings = link.from_node.evaluate_input("Use Tail")
                        break
                    inp = link.from_node.outputs[link.from_socket]
                    inp = inp.traverse_target
                    if not inp:
                        break
                    if inp.links:
                        link = inp.links[0]
                    else:
                        link = None
        
            
            for name in ik_names.keys():
                self.inputs[name].hide = not self.display_ik_settings
            
            for name in display_names.keys():
                if name in ['Custom Object', 'Bone Collection']: continue
                self.inputs[name].hide = not self.display_vp_settings

            for name in deform_names.keys():
                if name in ['Deform']: continue
                self.inputs[name].hide = not self.display_def_settings

            for name in bbone_names.keys():
                if name in ['BBone Segments']: continue
                self.inputs[name].hide = not self.display_bb_settings
            

    # def update(self):
    #     self.display_update()

        
    
    # def copy(ectype, archtype):
        # # TODO: automatically avoid duplicating names
        # ectype.inputs["Name"].default_value = ""


class xFormArmatureNode(Node, xFormNode):
    '''A node representing an Armature object node'''
    bl_idname = 'xFormArmatureNode'
    bl_label = "Armature"
    bl_icon = 'OUTLINER_OB_ARMATURE'
    initialized : bpy.props.BoolProperty(default = False)

    def init(self, context):
        self.inputs.new('StringSocket', "Name")
        self.inputs.new('RelationshipSocket', "Relationship")
        self.inputs.new('RotationOrderSocket', "Rotation Order")
        self.inputs.new('MatrixSocket', "Matrix")
        self.outputs.new('xFormSocket', "xForm Out")

        # color
        self.use_custom_color = True
        self.color = xFormColor

        self.initialized=True


class xFormGeometryObjectNode(Node, xFormNode):
    """Represents a curve or mesh object."""
    bl_idname = "xFormGeometryObject"
    bl_label = "Geometry Object"
    bl_icon = "EMPTY_AXIS"
    initialized : bpy.props.BoolProperty(default = False)
    
    def init(self, context):
        self.inputs.new('StringSocket', "Name")
        self.inputs.new('GeometrySocket', "Geometry")
        self.inputs.new('MatrixSocket', "Matrix")
        self.inputs.new('RelationshipSocket', "Relationship")
        self.inputs.new('DeformerSocket', "Deformer")
        self.inputs.new ("HideSocket", "Hide in Viewport")
        self.inputs.new ("BooleanSocket", "Hide in Render")
        self.outputs.new('xFormSocket', "xForm Out")

        # color
        self.use_custom_color = True
        self.color = xFormColor

        self.initialized=True