rig_picker/operators/picker.py

import bpy
from bpy.props import EnumProperty
from ..constants import PICKERS
from ..core.bl_utils import get_view_3d_override
from ..core.picker import get_picker_path, pack_picker, unpack_picker
#from .func_bgl import draw_callback_px
#from .func_bgl import select_bone
#from core.picker import *
#from .utils import is_over_region
import gpu
from mathutils import Vector
from gpu_extras.batch import batch_for_shader
from pathlib import Path
import json
import os


vertex_shader = '''
    layout (location = 0) in vec2 pos;

    flat out vec2 startPos;
    out vec2 vertPos;

    uniform mat4 viewMatrix;

    void main()
    {
        vec4 outPos = viewMatrix * vec4(pos.x, pos.y, 0.0, 1.0);
        gl_Position = outPos;
        vertPos = pos.xy / outPos.w;
        startPos = vertPos;
    }

'''

fragment_shader = '''
    flat in vec2 startPos;
    in vec2 vertPos;

    out vec4 fragColor;

    uniform vec4 color;
    uniform float dashSize;
    uniform float gapSize;

    void main()
    {
        vec2  dir  = (vertPos.xy - startPos.xy);
        float dist = length(dir);

        if (fract(dist / (dashSize + gapSize)) > dashSize/(dashSize + gapSize))
            discard; 
        fragColor = color;
    }
'''




def draw_callback(self):
    #print('draw callback border')
    if not self.draw_border:
        return

    gpu.state.blend_set('ALPHA')

    #print('DRAW BORDER')
    
    self.color_shader.bind()
    self.color_shader.uniform_float("color", self.bg_color)
    self.bg_batch.draw(self.color_shader)

    self.dash_shader.bind()
    matrix = gpu.matrix.get_projection_matrix()

    self.dash_shader.uniform_float("color", self.border_color)
    self.dash_shader.uniform_float("viewMatrix", matrix)
    self.dash_shader.uniform_float("dashSize", 5)
    self.dash_shader.uniform_float("gapSize", 4)
    
    self.contour_batch.draw(self.dash_shader)

    gpu.state.blend_set('NONE')


def is_picker_space(context):
    sp = context.space_data
    if sp and (sp.type == 'NODE_EDITOR' and sp.tree_type == 'RigPickerTree'):
        return True
    
    return False


class RP_OT_box_select(bpy.types.Operator):
    """Tooltip"""
    bl_idname = "node.rp_box_select"
    bl_label = "Picker Box Select"

    mode: EnumProperty(items=[(i, i.title(), '') for i in ('SET', 'EXTEND', 'SUBSTRACT')])

    color_shader = gpu.shader.from_builtin('2D_UNIFORM_COLOR')
    dash_shader = gpu.types.GPUShader(vertex_shader, fragment_shader)

    @classmethod
    def poll(cls, context):
        if not is_picker_space(context):
            return
            
        ob = context.object
        return ob and ob in PICKERS

    '''
    def mode_from_event(self, event):
        if event.alt:
            return 'SUBSTRACT'
        elif event.ctrl or event.shift:
            return 'EXTEND'
        else:
            return 'SET'
    '''

    def invoke(self, context, event):
        #print(f'invoke: type={event.type}, value={event.value}, ctrl={event.ctrl}, shift={event.shift}, alt={event.alt}')

        if context.object.mode != 'POSE':
            bpy.ops.object.posemode_toggle()

        self.timer = None
        #self.mode = self.mode_from_event(event)
        #self.invoke_event = event.copy()
        self.region = context.region
        self.draw_border = False

        self.picker = PICKERS[context.object]

        self.start_mouse = event.mouse_region_x, event.mouse_region_y
        #self.shader = line_strip_shader
        self.border_color = [1, 1, 1, 1]
        self.bg_color = [1, 1, 1, 0.05]

        #args = (self, context)
        self._handle = bpy.types.SpaceNodeEditor.draw_handler_add(draw_callback, (self,), 'WINDOW', 'POST_PIXEL')
        
        context.window_manager.modal_handler_add(self)

        self.picker.press_event(self.mode)
        self.region.tag_redraw()

        return {'RUNNING_MODAL'}
    
    def modal(self, context, event):

        self.mouse = event.mouse_region_x, event.mouse_region_y
        points_x = [v[0] for v in (self.start_mouse, self.mouse)]
        points_y = [v[1] for v in (self.start_mouse, self.mouse)]

        self.border = [
            (min(points_x), max(points_y)), 
            (max(points_x), max(points_y)), 
            (max(points_x), min(points_y)), 
            (min(points_x), min(points_y))
        ]

        self.bg_batch = batch_for_shader(self.color_shader, 'TRI_FAN', {"pos": self.border})
        self.contour_batch = batch_for_shader(self.dash_shader, 'LINE_LOOP', {"pos": self.border})

        self.draw_border = True

        self.region.tag_redraw()

        if event.value == 'RELEASE':
            return self.release_event()
                
        return {'RUNNING_MODAL'}

    def release_event(self):
        scn = bpy.context.scene

        if scn.rig_picker.use_pick_bone:
            self.picker.assign_bone_event()

        elif (self.start_mouse[0] != self.mouse[0] and self.start_mouse[1] != self.mouse[1]):
            self.picker.border_select(self.border, self.mode)
        else:
            self.picker.move_event(self.mouse)
            self.picker.release_event(self.mode)
        
        bpy.ops.ed.undo_push(message="Box Select")
        
        return self.exit()

    def exit(self):
        #print('Border Select Finished')

        bpy.types.SpaceNodeEditor.draw_handler_remove(self._handle, 'WINDOW')
        self.region.tag_redraw()
        return {'FINISHED'}



class RP_OT_picker_transform(bpy.types.Operator):
    """Tooltip"""
    bl_idname = "node.picker_transform"
    bl_label = "Move Bone in Picker View"

    mode : EnumProperty(items=[(m, m.title(), '') for m in ('TRANSLATE', 'ROTATE', 'SCALE')])

    @classmethod
    def poll(cls, context):
        return is_picker_space(context)

    def execute(self, context):
        with context.temp_override(**get_view_3d_override()):
            if self.mode == 'TRANSLATE':
                bpy.ops.transform.translate("INVOKE_DEFAULT")
            elif self.mode == 'ROTATE':
                bpy.ops.transform.rotate("INVOKE_DEFAULT")
            elif self.mode == 'SCALE':
                bpy.ops.transform.resize("INVOKE_DEFAULT")

        return {"FINISHED"}

    # def invoke(self, context, event):
    #     self.mouse = event.mouse_region_x, event.mouse_region_y

    #     self.bone_data = {b: b.matrix.copy() for b in context.selected_pose_bones}

    #     context.window_manager.modal_handler_add(self)
    #     return {'RUNNING_MODAL'}

    # def modal(self, context, event):

    #     delta_x = (event.mouse_region_x - self.mouse[0]) / 1000
    #     delta_y = (event.mouse_region_y - self.mouse[1]) / 1000

    #     for bone, matrix in self.bone_data.items():
    #         bone.matrix.translation = matrix.translation + Vector((delta_x, 0, delta_y))


    #     #print(delta_x, delta_y)

    #     if event.type=="LEFTMOUSE" and event.value == 'RELEASE':
    #         return {'FINISHED'}

    #     if event.type=="RIGHTMOUSE":
    #         for bone, matrix in self.bone_data.items():
    #             bone.matrix = matrix

    #         return {'CANCELLED'}

    #     return {'RUNNING_MODAL'}


class RP_OT_function_execute(bpy.types.Operator):
    bl_idname = "rigpicker.function_execute"
    bl_label = 'Function Execute'

    shape_index = bpy.props.IntProperty()

    @classmethod
    def poll(cls, context):
        if not is_picker_space(context):
            return

    def execute(self,context):
        event = self.event
        ob = context.object
        shape = ob.data.rig_picker['shapes'][self.shape_index]

        function = shape['function']
        if shape.get('variables'):
            variables=shape['variables'].to_dict()
        else:
            variables={}

        variables['event']=event
        globals()[function](variables)

        return {'FINISHED'}

    def invoke(self,context,event):
        self.event = event
        return self.execute(context)


class RP_OT_reload_picker(bpy.types.Operator):
    bl_idname = "rigpicker.reload_picker"
    bl_label = 'Reload Picker'

    #@classmethod
    #def poll(cls, context):
    #    if not is_picker_space(context):
    #        return

    def execute(self,context):
        PICKERS.clear()

        for a in context.screen.areas: 
            a.tag_redraw()

        return {"FINISHED"}


class RP_OT_toogle_bone_layer(bpy.types.Operator):
    bl_idname = "rigpicker.toogle_bone_layer"
    bl_label = 'Toogle Bone Layer'

    @classmethod
    def poll(cls, context):
        if not is_picker_space(context):
            return

        ob = context.object
        picker = PICKERS.get(ob)

        if picker.hover_shape and picker.hover_shape.type == 'bone':
            return True

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

        bone = picker.hover_shape.bone
        hide = picker.hover_shape.hide

        if bone:
            for i, l in enumerate(bone.bone.layers):
                if l:
                    ob.data.layers[i] = hide

        context.region.tag_redraw()

        return {"FINISHED"}


class RP_OT_context_menu_picker(bpy.types.Operator):
    bl_idname = "node.context_menu_picker"
    bl_label = 'Context Menu Picker'

    @classmethod
    def poll(cls, context):
        return is_picker_space(context)

    def execute(self, context):
        bpy.ops.wm.call_menu(name='RP_MT_context_menu')

        return {"FINISHED"}


class RP_OT_pack_picker(bpy.types.Operator):
    """Pack Unpack the picker on the rig"""
    bl_idname = "rigpicker.pack_picker"
    bl_label = 'Toogle Bone Layer'

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

    def execute(self, context):
        print('Pack Picker')
        
        rig = context.object
        picker_path = get_picker_path(rig)

        if 'picker' in rig.data.rig_picker.keys():
            unpack_picker(rig)
            self.report({"INFO"}, f'The picker is unpacked')
            return {"FINISHED"}

        if not picker_path.exists():
            self.report({"ERROR"}, f'The path of the picker not exist: {picker_path}')
            return {"CANCELLED"}

        pack_picker(rig)
        self.report({"INFO"}, f'The picker is packed')

        return {"FINISHED"}


class RP_OT_call_operator(bpy.types.Operator):
    bl_idname = "rigpicker.call_operator"
    bl_label = 'Toogle Bone Layer'

    operator: bpy.props.StringProperty()

    @classmethod
    def poll(cls, context):
        return is_picker_space(context)

    def execute(self, context):

        print('CALL OPERATOR', self.operator)

        try:
            exec(self.operator)
        except Exception as e:
            self.report({"ERROR"}, e)
        
        context.region.tag_redraw()

        return {"FINISHED"}


class RP_MT_context_menu(bpy.types.Menu):
    bl_label = "Context Menu"
    #bl_idname = "RP_MT_context_menu"

    # Set the menu operators and draw functions
    def draw(self, context):
        layout = self.layout
        col = layout.column()
        col.use_property_split = True

        ob = context.object
        picker = PICKERS.get(ob)

        if picker.hover_shape and picker.hover_shape.type == 'bone':
            bone = picker.hover_shape.bone
            
            for key in bone.keys():
                layout.prop(bone,f'["{key}"]', slider=True)

        #layout.operator("rigpicker.show_bone_layer", text="Show Bone Layer", ).type = 'ACTIVE'
        #layout.operator("rigidbody.objects_add", text="B Add Passive").type = 'PASSIVE'

'''
class RP_OT_ui_draw(bpy.types.Operator):
    bl_idname = "rigpicker.ui_draw"
    bl_label = "Rig UI Draw"

    _handle = None
    tmp_ob = None
    tmp_bones = []
    start = (0,0)
    end = (0,0)
    border = ((0,0),(0,0),(0,0),(0,0))
    is_border = False
    press = False
    scale = 1
    offset = 0
    outside_point = (-1,-1)
    addon_keymaps = []

    def set_shorcut(self,context):
        ob = context.object
        addon = bpy.context.window_manager.keyconfigs.addon

        if ob and ob.type =='ARMATURE' and ob.data.rig_picker and addon:
            for i,shape in [(i,s) for i,s in enumerate(ob.data.rig_picker['shapes']) if s.get('function') and s.get('shortcut')]:
                km = addon.keymaps.new(name = 'Image Generic', space_type = 'IMAGE_EDITOR',region_type = 'WINDOW')

                split = shape["shortcut"].split(' ')
                if len(split)==1:
                    shortcut = shape["shortcut"].upper()
                    modifier = None
                else:
                    shortcut = split[1].upper()
                    modifier = split[0].lower()

                kmi = km.keymap_items.new("rigpicker.function_execute", type = shortcut, value = "CLICK")
                kmi.properties.shape_index = i

                if modifier:
                    setattr(kmi,modifier,True)

                self.addon_keymaps.append(km)

    def remove_shorcut(self,context):
        # Remove Shortcut
        wm = bpy.context.window_manager
        for km in self.addon_keymaps:
            for kmi in km.keymap_items:
                km.keymap_items.remove(kmi)

        self.addon_keymaps.clear()

    def modal(self, context, event):
        inside = is_over_region(self,context,event)

        if context.object and context.object.type == 'ARMATURE' and context.area:
            if not context.screen.is_animation_playing:
                if self.tmp_ob != context.object:
                    context.area.tag_redraw()
                    self.remove_shorcut(context)
                    self.set_shorcut(context)
                    self.tmp_ob = context.object

                if self.tmp_bones != context.selected_pose_bones:
                    context.area.tag_redraw()
                    self.tmp_bones = context.selected_pose_bones

                if inside:
                    context.area.tag_redraw()

            if event.type == 'LEFTMOUSE':
                if event.value == 'PRESS': # start selection
                    if inside:
                        self.start = (event.mouse_region_x,event.mouse_region_y)
                        self.press = True

                elif event.value == 'RELEASE'  and self.press:
                    self.end = (event.mouse_region_x, event.mouse_region_y)

                    select_bone(self, context, event)
                    bpy.ops.ed.undo_push()

                    self.is_border= False
                    self.press = False

            if event.type == 'MOUSEMOVE':
                self.end = (event.mouse_region_x, event.mouse_region_y)

                if self.press:
                    b_x = (min(self.start[0], self.end[0]), max(self.start[0], self.end[0]))
                    b_y = (min(self.start[1], self.end[1]), max(self.start[1], self.end[1]))
                    self.border = ((b_x[0], b_y[1]), (b_x[1], b_y[1]), (b_x[1], b_y[0]), (b_x[0], b_y[0]))
                    self.is_border = True if (b_x[1]-b_x[0])+(b_y[1]-b_y[0]) > 4 else False

                if self.is_border:
                    select_bone(self, context, event)

            elif event.type in {'ESC',} and inside:
                bpy.types.SpaceImageEditor.draw_handler_remove(self._handle, 'WINDOW')
                self.remove_shorcut(context)


                return {'CANCELLED'}

        return {'PASS_THROUGH'}

    def invoke(self, context, event):
        #shortcut Creation



        context.space_data.image = None
        self.adress = context.area.as_pointer()
        args = (self, context)

        self._handle = bpy.types.SpaceImageEditor.draw_handler_add(draw_callback_px, args, 'WINDOW', 'POST_PIXEL')


        context.window_manager.modal_handler_add(self)

        return {'RUNNING_MODAL'}
'''

keymaps = []
def register_keymaps():
    wm = bpy.context.window_manager

    km = wm.keyconfigs.addon.keymaps.new(name="Node Editor", space_type="NODE_EDITOR")

    kmi = km.keymap_items.new("rigpicker.call_operator", type="X", value="PRESS")
    kmi.properties.operator = "bpy.ops.animtoolbox.reset_bone()"
    keymaps.append((km, kmi))
    
    kmi = km.keymap_items.new("rigpicker.call_operator", type="A", value="PRESS")
    kmi.properties.operator = "bpy.ops.pose.select_all(action='SELECT')"
    keymaps.append((km, kmi))

    kmi = km.keymap_items.new("rigpicker.call_operator", type="A", value="PRESS", alt=True)
    kmi.properties.operator = "bpy.ops.pose.select_all(action='DESELECT')"
    keymaps.append((km, kmi))

    kmi = km.keymap_items.new("rigpicker.toogle_bone_layer", type="LEFTMOUSE", value="DOUBLE_CLICK")
    keymaps.append((km, kmi))

    kmi = km.keymap_items.new("node.picker_transform", type="G", value="PRESS")
    kmi.properties.mode = 'TRANSLATE'
    keymaps.append((km, kmi))

    kmi = km.keymap_items.new("node.picker_transform", type="R", value="PRESS")
    kmi.properties.mode = 'ROTATE'
    keymaps.append((km, kmi))

    kmi = km.keymap_items.new("node.picker_transform", type="S", value="PRESS")
    kmi.properties.mode = 'SCALE'
    keymaps.append((km, kmi))

    kmi = km.keymap_items.new("node.context_menu_picker", type="RIGHTMOUSE", value="PRESS")
    keymaps.append((km, kmi))

    kmi = km.keymap_items.new("node.rp_box_select", type="LEFTMOUSE", value="PRESS")
    kmi.properties.mode = 'SET'
    keymaps.append((km, kmi))

    kmi = km.keymap_items.new("node.rp_box_select", type="LEFTMOUSE", value="PRESS", shift=True)
    kmi.properties.mode = 'EXTEND'
    keymaps.append((km, kmi))

    #kmi = km.keymap_items.new("node.rp_box_select", type="LEFTMOUSE", value="PRESS", ctrl=True)
    kmi = km.keymap_items.new("node.rp_box_select", type="LEFTMOUSE", value="PRESS", alt=True)
    kmi.properties.mode = 'SUBSTRACT'

    keymaps.append((km, kmi))

def unregister_keymaps():
    for km, kmi in keymaps:
        km.keymap_items.remove(kmi)
    keymaps.clear()

classes = (
    RP_OT_box_select,
    RP_OT_function_execute,
    RP_OT_reload_picker,
    RP_OT_toogle_bone_layer,
    RP_OT_call_operator,
    RP_MT_context_menu,
    RP_OT_picker_transform,
    RP_OT_context_menu_picker,
    RP_OT_pack_picker
    #RP_OT_ui_draw
)


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

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