bone_widget/operators.py

import bpy
from bpy.types import Operator
from bpy.props import BoolProperty, IntProperty, FloatProperty, FloatVectorProperty, \
StringProperty

from bone_widget import ctx
from .transform_utils import transform_matrix, apply_mat_to_verts, get_bone_size_factor, \
get_bone_matrix
from .icon_utils import render_widget
from .shape_utils import get_clean_shape, symmetrize_bone_shape, link_to_col, get_bone, custom_shape_matrix

import re
from pathlib import Path
import os
from tempfile import gettempdir
from mathutils import Matrix
import json


class BW_OT_copy_widgets(Operator):
    bl_idname = 'bonewidget.copy_widgets'
    bl_label = "Copy Widgets"

    @classmethod
    def poll(cls, context):
        return context.object and context.object.type == 'ARMATURE'

    def execute(self, context):
        from tempfile import gettempdir

        blend_file = Path(gettempdir())/'bone_widgets.blend'

        shapes = []

        ob = context.object
        if ob.mode == 'POSE':
            bones = ctx.selected_bones
        else:
            bones = ctx.bones

        for b in bones:
            if b.custom_shape:
                s = b.custom_shape.copy()
                s.data = s.data.copy()

                if not b.use_custom_shape_bone_size:
                    mat = transform_matrix(scale=(1/b.bone.length,)*3)
                    s.data.transform(mat)

                #s.data.transform(s.matrix_world.inverted())
                #s.matrix_world = Matrix()
                
                s['.bw_bone'] = b.name

                if bpy.app.version_string < '3.0.0':
                    s['.bw_custom_shape_scale'] = b.custom_shape_scale
                else:
                    s['.bw_custom_shape_translation'] = b.custom_shape_translation
                    s['.bw_custom_shape_rotation_euler'] = b.custom_shape_rotation_euler
                    s['.bw_custom_shape_scale_xyz'] = b.custom_shape_scale_xyz
                
                s['.bw_custom_shape_transform'] = b.custom_shape_transform
                s['.bw_use_custom_shape_bone_size'] = b.use_custom_shape_bone_size
                #s['.bw_size_factor'] = get_bone_size_factor(b, s, b.use_custom_shape_bone_size)

                shapes.append(s)

        bpy.data.libraries.write(str(blend_file), set(shapes))

        for s in shapes:
            data = s.data
            data_type = s.type

            bpy.data.objects.remove(s)

            if data_type == 'MESH':
                bpy.data.meshes.remove(data)
            elif data_type == 'CURVE':
                bpy.data.curves.remove(data)


        # for b in ctx.bones:
        #     if b.custom_shape:
        #         for k in b.custom_shape.keys():
        #             if k.startswith('.bw'):
        #                 del b.custom_shape[k]

        return {'FINISHED'}


class BW_OT_paste_widgets(Operator):
    bl_idname = 'bonewidget.paste_widgets'
    bl_label = "Paste Widgets"

    path: StringProperty(subtype='FILE_PATH')

    @classmethod
    def poll(cls, context):
        return context.object and context.object.type == 'ARMATURE'

    def execute(self, context):
        from tempfile import gettempdir

        rig = context.object

        #for b in rig.pose.bones:
        #    if b.custom_shape:
        #        bpy.data.objects.remove(b.custom_shape)
                #b.custom_shape = get_clean_shape(b, b.custom_shape, separate=b.custom_shape.users>2)

        if self.path:
            blend_file = Path(self.path)
        else:
            blend_file = Path(gettempdir())/'bone_widgets.blend'

        with bpy.data.libraries.load(str(blend_file), link=False) as (data_src, data_dst):
            data_dst.objects = data_src.objects
        shapes = data_src.objects

        #old_shapes = set()
        for s in shapes:
            bone = rig.pose.bones.get(s['.bw_bone'])
            if not bone:
                print('No bone found for', s)
                continue

            if bpy.app.version_string < '3.0.0':
                bone.custom_shape_scale = s['.bw_custom_shape_scale']
            else:
                bone.custom_shape_scale_xyz = s['.bw_custom_shape_scale_xyz']

            link_to_col(s, ctx.widget_col)
            #if bone.custom_shape:
            #size_factor = max(s.dimensions)#get_bone_size_factor(bone, s, relative=False)
            #size_factor /= s['.bw_size_factor']
            #size_factor = s['.bw_size_factor']
            #mat = transform_matrix(scale=(size_factor,)*3)
            #s.data.transform(Matrix.Scale(1/size_factor, 4))
            #s.scale = 1,1,1


            bone.custom_shape = get_clean_shape(bone, s, col=ctx.widget_col, prefix=ctx.prefs.prefix, 
            separate=False, apply_transforms=False)

            if bpy.app.version_string < '3.0.0':
                #bone.custom_shape_scale = s['.bw_custom_shape_scale']
                bone.custom_shape_scale = 1
            else:
                bone.custom_shape_translation = s['.bw_custom_shape_translation']
                bone.custom_shape_rotation_euler = s['.bw_custom_shape_rotation_euler']
                #bone.custom_shape_scale_xyz = s['.bw_custom_shape_scale_xyz']
                bone.custom_shape_scale_xyz = 1,1,1
                #use_custom_shape_bone_size

            bone.custom_shape_transform = s['.bw_custom_shape_transform']
            bone.use_custom_shape_bone_size = True#s['.bw_use_custom_shape_bone_size']

            bone.bone.show_wire = not bool(s.data.polygons)
        

            #mat = transform_matrix(scale=(s['.bw_size_factor'],)*3)
            #s.data.transform(mat)

        for b in rig.pose.bones:
            if b.custom_shape:
                for k in list(b.custom_shape.keys()):
                    if k.startswith('.bw'):
                        del b.custom_shape[k]
        


        return {'FINISHED'}


class BW_OT_remove_unused_shape(Operator):
    bl_idname = 'bonewidget.remove_unused_shape'
    bl_label = "Remove Unused Shape"

    def execute(self, context):
        objects = list(ctx.widget_col.all_objects)
        for ob in objects:
            if not get_bone(ob) and ob in bpy.data.objects[:]:
                bpy.data.objects.remove(ob)

        return {'FINISHED'}


class BW_OT_auto_color(Operator):
    bl_idname = 'bonewidget.auto_color'
    bl_label = "Auto Color"
    bl_options = {'REGISTER', 'UNDO'}

    def is_bone_protected(self, bone):
        rig = bone.id_data
        return any([i==j==True for i, j in zip(rig.data.layers_protected, bone.bone.layers)])

    def execute(self, context):
        ob = context.object
        bones = context.selected_pose_bones or ob.pose.bones

        for b in bones:
            if self.is_bone_protected(b):
                continue
            for group in ob.pose.bone_groups:
                if any(i in b.name.lower() for i in group.name.lower().split(' ')):
                    b.bone_group = group

        return {'FINISHED'}
        

class BW_OT_load_default_color(Operator):
    bl_idname = 'bonewidget.load_default_color'
    bl_label = "Load Default Color"
    bl_options = {'REGISTER', 'UNDO'}

    select_color: FloatVectorProperty(name='Color', subtype='COLOR', default=[0.0, 1.0, 1.0])
    active_color: FloatVectorProperty(name='Color', subtype='COLOR', default=[1.0, 1.0, 1.0])

    def execute(self, context):
        ob = context.object

        colors = { 
            'Black': [0, 0, 0],
            'Yellow': [1, 1, 0],
            'Red': [0, 0.035, 0.95],
            'Blue': [1, 0, 0],
            'Green': [0.733333, 0.937255, 0.356863],
            'Pink': [1, 0.1, 0.85],
            'Purple': [0.67, 0, 0.87],
            'Brown': [0.75, 0.65, 0],
            'Orange': [1, 0.6, 0],
            }

        for k, v in colors.items():
            bone_group = ob.pose.bone_groups.get(k)

            if not bone_group:
                bone_group = ob.pose.bone_groups.new(name=k)
            
            bone_group.color_set = 'CUSTOM'

            bone_group.colors.normal = v
            bone_group.colors.select = self.select_color
            bone_group.colors.active = self.active_color

        return {'FINISHED'}


class BW_OT_copy_bone_groups(Operator):
    bl_idname = 'bonewidget.copy_bone_groups'
    bl_label = "Copy Bone Groups"
    bl_options = {'REGISTER', 'UNDO'}

    @classmethod
    def poll(self, context):
        return context.object and context.object.type == 'ARMATURE'

    def execute(self, context):
        print('Copy Bone Group')

        ob = context.object

        bone_groups = {}
        for bg in ob.pose.bone_groups:
            bone_groups[bg.name] = {
                'colors': [bg.colors.normal[:], bg.colors.select[:], bg.colors.active[:]],
                'bones': [b.name for b in ob.pose.bones if b.bone_group == bg]
                }
        
        blend_file = Path(gettempdir()) / 'bw_copy_bone_groups.json'
        
        blend_file.write_text(json.dumps(bone_groups))

        return {'FINISHED'}


class BW_OT_paste_bone_groups(Operator):
    bl_idname = 'bonewidget.paste_bone_groups'
    bl_label = "Paste Bone Groups"
    bl_options = {'REGISTER', 'UNDO'}

    clear: BoolProperty(default=False)

    @classmethod
    def poll(self, context):
        return context.object and context.object.type == 'ARMATURE'

    def execute(self, context):
        print('Paste Bone Group')

        ob = context.object

        blend_file = Path(gettempdir()) / 'bw_copy_bone_groups.json'
        
        bone_groups = json.loads(blend_file.read_text())

        if self.clear:
            for bg in reversed(ob.pose.bone_groups):
                ob.pose.bone_groups.remove(bg)
        
        for bg_name, bg_data in bone_groups.items():
            bg = ob.pose.bone_groups.get(bg_name)
            if not bg:
                bg = ob.pose.bone_groups.new(name=bg_name)

            bg.color_set = 'CUSTOM'
            
            bg.colors.normal,bg.colors.select, bg.colors.active = bg_data['colors']

            for b in bg_data['bones']:
                bone = ob.pose.bones.get(b)
                if bone:
                    bone.bone_group = bg

        return {'FINISHED'}


class BW_OT_add_folder(Operator):
    bl_idname = 'bonewidget.add_folder'
    bl_label = "Add Folder"

    def execute(self, context):
        folder = ctx.prefs.folders.add()
        folder.path = ''

        return {'FINISHED'}


class BW_OT_remove_folder(Operator):
    bl_idname = 'bonewidget.remove_folder'
    bl_label = "Remove Folder"
    bl_options = {'REGISTER', 'UNDO'}

    index: IntProperty()

    def execute(self, context):

        ctx.prefs.folders.remove(self.index)

        if self.index == 0:
            ctx.prefs.folder_enum = str(ctx.default_folder_path)

        #update_folder_items()

        bpy.context.area.tag_redraw()

        return {'FINISHED'}


class BW_OT_refresh_folders(Operator):
    bl_idname = 'bonewidget.refresh_folders'
    bl_label = "Refresh"

    def execute(self, context):
        for f in ctx.folders:
            f.widgets.clear()
            f.load_widgets()
        return {'FINISHED'}


class BW_OT_rename_folder(Operator):
    bl_idname = 'bonewidget.rename_folder'
    bl_label = "Rename Folder"
    bl_options = {'REGISTER', 'UNDO'}

    name: StringProperty(name='Name')

    @classmethod
    def poll(cls, context):
        return ctx.active_folder

    def execute(self, context):
        folder = ctx.active_folder
        folder.rename(self.name)

        return {'FINISHED'}
    
    def invoke(self, context, event):
        wm = context.window_manager
        self.name = Path(ctx.prefs.folder_enum).name

        return wm.invoke_props_dialog(self)


# class BW_OT_show_preferences(Operator):
#     """Display the preferences to the tab panel"""
#     bl_idname = 'bonewidget.show_preferences'
#     bl_label = "Show Preferences"
#     bl_options = {'REGISTER'} 

#     def execute(self, context):
#         #ctx.init_transforms()
#         return {'FINISHED'}


class BW_OT_transform_widget(Operator):
    """Transform the Rotation Location or Scale of the selected shapes"""
    bl_idname = 'bonewidget.transform_widget'
    bl_label = "Transfom"
    bl_options = {'REGISTER', 'UNDO'}

    symmetrize: BoolProperty(name='Symmetrize', default=True)

    @classmethod
    def poll(cls, context):
        return context.object and context.object.mode == 'POSE'

    def execute(self, context):
        ctx.init_transforms()
        return {'FINISHED'}


class BW_OT_create_widget(Operator):
    """Create the widget shape and assign it to the bone"""
    bl_idname = 'bonewidget.create_widget'
    bl_label = "Create Widget"
    bl_options = {'REGISTER', 'UNDO'}

    symmetrize: BoolProperty(name='Symmetrize', default=True)

    @classmethod
    def poll(cls, context):
        return ctx.selected_bones and ctx.active_widget

    def execute(self, context):
        folder = ctx.active_folder        
        blend = folder.get_widget_path(folder.active_widget)

        with bpy.data.libraries.load(str(blend), link=False) as (data_src, data_dst):
            data_dst.objects = data_src.objects
        shape = data_src.objects[0]

        for bone in ctx.selected_bones:
            #shape_copy = shape.copy()

            if bpy.app.version_string < '3.0.0':
                bone.custom_shape_scale = 1.0 
            else:
                bone.custom_shape_scale_xyz = [1.0]*3 # for blender 3.0

            bone.bone.show_wire = not bool(shape.data.polygons)
        
            #if bone.custom_shape and bone.custom_shape.users == 2:
            #    bpy.data.objects.remove(bone.custom_shape)
            
            #copy_shape = shape.copy()
            #copy_shape.data = copy_shape.data.copy()
            bone.custom_shape = get_clean_shape(bone, shape, col=ctx.widget_col, prefix=ctx.prefs.prefix)

            if self.symmetrize:
                symmetrize_bone_shape(bone, prefix=ctx.prefs.prefix)

        bpy.data.objects.remove(shape)
        ctx.init_transforms()


        return {'FINISHED'}


class BW_OT_match_transform(Operator):
    bl_idname = 'bonewidget.match_transform'
    bl_label = "Match Transforms"
    bl_options = {'REGISTER', 'UNDO'}

    relative: BoolProperty(default=False)

    def execute(self, context):
        for bone in ctx.selected_bones:
            shape = bone.custom_shape
            if not shape:
                continue

            size_factor = get_bone_size_factor(bone, shape, self.relative)
            mat = transform_matrix(scale=(size_factor,)*3)
            shape.data.transform(mat)
            #apply_mat_to_verts(shape.data, mat)

        ctx.init_transforms()

        return {'FINISHED'}


class BW_OT_edit_widget(Operator):
    bl_idname = 'bonewidget.edit_widget'
    bl_label = "Edit Widget"
    bl_options = {'REGISTER', 'UNDO'}

    clean: BoolProperty(name='Clean', default=True)

    @classmethod
    def poll(cls, context):
        return ctx.bone and ctx.bone.custom_shape

    def execute(self, context):
        #ctx._rig = ctx.rig
        widgets = ctx.selected_widgets
        bones = ctx.selected_bones
        rig = ctx.rig
        prefs = ctx.prefs

        if self.clean:
            bpy.ops.bonewidget.clean_widget()

        bpy.ops.object.mode_set(mode='OBJECT')
        bpy.ops.object.select_all(action='DESELECT')

        rig.hide_set(True)
        
        ctx.widget_col.hide_viewport = False
        ctx.widget_layer_col.exclude = False
        ctx.widget_layer_col.hide_viewport = False

        for bone in bones:
            if not bone.custom_shape:
                continue
            #link_to_col(w, ctx.widget_col)
            bone.custom_shape.select_set(True)
            context.view_layer.objects.active = bone.custom_shape

        bpy.ops.object.mode_set(mode='EDIT')


        return {'FINISHED'}


class BW_OT_return_to_rig(Operator):
    bl_idname = 'bonewidget.return_to_rig'
    bl_label = "Return to rig"
    bl_options = {'REGISTER', 'UNDO'}

    clean: BoolProperty(name='Clean', default=True)

    def invoke(self, context, event):
        self.symmetrize = ctx.prefs.auto_symmetrize
        return self.execute(context)

    def execute(self, context):
        rig = ctx.rig

        #print('Rig', rig)
        bone = ctx.bone
        bones = ctx.selected_bones


        ctx.widget_col.hide_viewport = False
        bpy.ops.object.mode_set(mode='OBJECT')

        bpy.ops.object.select_all(action='DESELECT')

        rig.hide_set(False)
        ctx.widget_col.hide_viewport = True

        context.view_layer.objects.active = rig            

        bpy.ops.object.mode_set(mode='POSE')
        rig.data.bones.active = bone.bone

        for b in rig.pose.bones:
            b.bone.select = b in bones+[bone]

        if self.clean:
            bpy.ops.bonewidget.clean_widget()

        return {'FINISHED'}

'''
class BW_OT_symmetrize_widget(Operator):
    bl_idname = 'bonewidget.symmetrize_widget'
    bl_label = "Symmetrize"
    bl_options = {'REGISTER', 'UNDO'}

    match_transform: BoolProperty(True)

    def get_name_side(self, name, fallback=None):
        if name.lower().endswith('.l'):
            return 'LEFT'
        elif name.lower().endswith('.r'):
            return 'RIGHT'
        return fallback

    def mirror_name(self, name):
        mirror = None

        match = {
        'R': 'L',
        'r': 'l',
        'L': 'R',
        'l': 'r',
        }

        separator = ['.', '_']

        if name.startswith(tuple(match.keys())):
            if name[1] in separator:
                mirror = match[name[0]] + name[1:]

        if name.endswith(tuple(match.keys())):
            if name[-2] in separator:
                mirror = name[:-1] + match[name[-1]]

        return mirror

    def execute(self, context):
        bones = ctx.selected_bones
        if ctx.bone:
            active_side = self.get_name_side(ctx.bone.name, 'LEFT')
            if active_side:
                bones = [b for b in ctx.selected_bones if self.get_name_side(b) == active_side]

        for bone in bones:
            flip_name = self.mirror_name(bone.name)
            flip_bone = ctx.rig.pose.bones.get(flip_name)

            if flip_bone:
                shape = flip_bone.custom_shape
                if shape.users <= 2:
                    bpy.data.objects.remove(shape)

                shape = bone.custom_shape.copy()
                if shape

                flip_bone.custom_shape = shape
                shape.matrix_world = get_bone_matrix(bone)
                shape.data.transform(transform_matrix(scale=(-1, 1, 1)))
                ctx.rename_shape(shape, flip_bone)

        return {'FINISHED'}
'''


class BW_OT_add_widget(Operator):
    bl_idname = 'bonewidget.add_widget'
    bl_label = "Add Widget"
    bl_options = {'REGISTER', 'UNDO'}

    replace: BoolProperty(default=False, name='Replace')

    @classmethod
    def poll(cls, context):
        return ctx.widget

    def invoke(self, context, event):
        if event.ctrl:
            self.replace = True
        
        return self.execute(context)

    def execute(self, context):
        folder = ctx.active_folder
        shape = ctx.widget
        bone = ctx.bone

        if self.replace and ctx.active_widget:
            name = ctx.active_widget.name
        else: #Find a unique name
            name = folder.get_widget_display_name(bone.name if bone else shape.name)

            i = 0
            org_name = name
            while name in folder.widgets:
                name = f'{org_name} {i:02d}'
                i += 1
       
        widget_path = folder.get_widget_path(name)
        icon_path = folder.get_icon_path(name)

        widget_path.parent.mkdir(exist_ok=True, parents=True)
        bpy.data.libraries.write(str(widget_path), {shape})

        # Copy the shape to apply the bone matrix to the mesh
        shape_copy = shape.copy()
        shape_copy.data = shape_copy.data.copy()
        
        if bone:
            shape_copy.matrix_world = get_bone_matrix(bone)

            mat = custom_shape_matrix(bone)
            shape_copy.data.transform(mat)

        render_widget(shape_copy, icon_path)
        folder.add_widget(name)

        bpy.data.objects.remove(shape_copy)

        bpy.context.area.tag_redraw()

        return {'FINISHED'}


class BW_OT_remove_widget(Operator):
    bl_idname = 'bonewidget.remove_widget'
    bl_label = "Remove Widget"
    bl_options = {'REGISTER', 'UNDO'}

    @classmethod
    def poll(cls, context):
        return ctx.active_widget

    def execute(self, context):
        folder = ctx.active_folder
        widget = ctx.active_widget

        folder.remove_widget(widget)

        return {'FINISHED'}


class BW_OT_clean_widget(Operator):
    bl_idname = 'bonewidget.clean_widget'
    bl_label = "Clean"
    bl_options = {'REGISTER', 'UNDO'}

    all: BoolProperty(name='All', default=False)
    symmetrize: BoolProperty(name='Symmetrize', default=True)

    @classmethod
    def poll(cls, context):
        return context.object and context.object.mode == 'POSE'

    def invoke(self, context, event):
        self.symmetrize = ctx.prefs.auto_symmetrize
        return self.execute(context)

    def execute(self, context):
        scene = context.scene
        prefs = ctx.prefs

        if self.all:
            bones = ctx.bones
        else:
            bones = ctx.selected_bones

        for bone in bones:
            shape = bone.custom_shape
            if not shape:
                continue

            bone.custom_shape = get_clean_shape(bone, shape, separate=prefs.auto_separate, 
            rename=prefs.auto_rename, col=ctx.widget_col, prefix=prefs.prefix, match=prefs.auto_match_transform)

            if self.symmetrize:
                symmetrize_bone_shape(bone, prefix=ctx.prefs.prefix)

            if shape in bpy.data.objects[:] and shape.users <= 1:
                bpy.data.objects.remove(shape)

        return {'FINISHED'}


classes = (
    BW_OT_remove_unused_shape,
    BW_OT_auto_color,
    BW_OT_load_default_color,
    BW_OT_add_folder,
    BW_OT_remove_folder,
    BW_OT_refresh_folders,
    BW_OT_rename_folder,
    BW_OT_transform_widget,
    BW_OT_match_transform,
    BW_OT_create_widget,
    BW_OT_edit_widget,
    BW_OT_return_to_rig,
    #BW_OT_symmetrize_widget,
    BW_OT_add_widget,
    BW_OT_remove_widget,
    BW_OT_clean_widget,
    # BW_OT_show_preferences,
    BW_OT_copy_widgets,
    BW_OT_paste_widgets,
    BW_OT_copy_bone_groups,
    BW_OT_paste_bone_groups
)

def bw_bone_group_menu(self, context):
    layout = self.layout
    layout.operator("bonewidget.load_default_color", icon='IMPORT')
    layout.operator("bonewidget.auto_color", icon='COLOR')
    layout.operator("bonewidget.copy_bone_groups", icon='COPYDOWN')
    layout.operator("bonewidget.paste_bone_groups", icon='PASTEDOWN')

def register():
    for cls in classes:
        bpy.utils.register_class(cls)

    bpy.types.DATA_MT_bone_group_context_menu.prepend(bw_bone_group_menu)

def unregister():
    bpy.types.DATA_MT_bone_group_context_menu.remove(bw_bone_group_menu)

    for cls in reversed(classes):
        bpy.utils.unregister_class(cls)