rig_picker/core/picker.py

import bpy
import gpu
from gpu_extras.batch import batch_for_shader
import blf
from mathutils import bvhtree, Vector
from mathutils.geometry import (intersect_point_quad_2d, intersect_point_tri_2d, 
        intersect_tri_tri_2d)
from ..constants import PICKERS
from .addon_utils import get_operator_from_id
from .geometry_utils import bounding_rect

from pathlib import Path
import re
import json
import os

import threading 


class Shape:
    def __init__(self, picker, points, polygons=None, edges=None, tooltip='', color=None, source_name=''):

        self.type = 'display'
        self.picker = picker
        self.rig = picker.rig
        self.source_name = source_name

        self.hover = False
        self.press = False

        self.shader = gpu.shader.from_builtin('UNIFORM_COLOR')

        #self.hover_shader = gpu.shader.from_builtin('UNIFORM_COLOR')
        #self.hover_shader.uniform_float("color", [1, 1, 1, 0.1])
        self.hover_color = [1, 1, 1, 0.1]

        self.tooltip = tooltip
        self.color = color
        self.points = points
        self.polygons = polygons or []
        self.edges = edges or []

        self.p_batch = batch_for_shader(self.shader, 'TRIS', {"pos": self.points}, indices=self.polygons)
        self.e_batch = batch_for_shader(self.shader, 'LINES', {"pos": self.points}, indices=self.edges)

        self.rect = bounding_rect(points)

    @property
    def color(self):
        return self._color
    
    @color.setter
    def color(self, color=None):
        if not color:
            color = [0.5, 0.5, 0.5, 1]
        elif isinstance(color, (tuple, list)) and len(color) in (3, 4):
            if len(color) == 3:
                color = [*color, 1]
            elif len(color) == 4:
                color = list(color)
            else:
                raise Exception('color must have a len of 3 or 4')
        else:
            raise Exception(f'color is {type(color)} must be None or (tuple, list)')
        
        #self.shader.uniform_float("color", color)
        self._color = color

    def draw(self):

        self.shader.bind()
        self.shader.uniform_float("color", self.color)

        if self.polygons:
            self.p_batch.draw(self.shader)
        if self.edges:
            self.e_batch.draw(self.shader)

    def move_event(self, location):
        if not intersect_point_quad_2d(location, *self.rect):
            self.hover = False
            return False
        
        for p in self.polygons:
            if intersect_point_tri_2d(location, *[self.points[i] for i in p]):
                self.hover = True
                return True
        
        self.hover = False
        return False
            
    def press_event(self, mode='SET'):
        self.press = True

    def release_event(self, mode='SET'):
        self.press = False


class BoneShape(Shape):
    def __init__(self, picker, points, polygons, edges, bone, tooltip='', color=None, source_name=''):
        super().__init__(picker, points=points, polygons=polygons, edges=edges, 
        tooltip=tooltip, color=color, source_name=source_name)

        self.type = 'bone'
        self.bone = bone
        self.active_color = [1, 1, 1, 0.1]

        self.bone_colors = self.get_bone_colors()

    @property
    def select(self):
        if not self.bone:
            return False

        return self.bone in (bpy.context.selected_pose_bones or []) #self.bone.bone.select

    @property
    def active(self):
        if not self.bone:
            return False

        return self.bone == bpy.context.active_pose_bone #self.rig.data.bones.active == self.bone.bone
    
    @property
    def hide(self):
        if not self.bone:
            return False

        #return self.bone not in (bpy.context.visible_pose_bones or [])
        
        return self.bone.bone.hide or not any(l.is_visible for l in self.bone.bone.collections)

    @property
    def bone_color(self):
        if not self.bone:
            return [0, 0, 0, 1]

        if self.select and self.active:
            return self.bone_colors['active']
        elif self.select:
            return self.bone_colors['select']
        elif self.hide:
            return self.bone_colors['hide']
        else:
            return self.bone_colors['normal']

    def get_bone_colors(self):

        theme = bpy.context.preferences.themes['Default']
        bone_colors = {
            'select': [*theme.view_3d.bone_pose, 1],
            'normal': [0.05, 0.05, 0.05, 1],
            'active': [*theme.view_3d.bone_pose_active, 1],
            'hide': [0.85, 0.85, 0.85, 0.2],
        }

        if not self.bone:
            return bone_colors

        if self.bone.color.palette == 'CUSTOM':
            bone_color = self.bone

        elif self.bone.color.palette == 'DEFAULT':
            bone_color = self.bone.bone

        normal_color = bone_color.color.custom.normal.copy()
        normal_color.s *= 0.75

        bone_colors['normal'] = [*normal_color, 1]
        bone_colors['select'] = [*bone_color.color.custom.select, 1]
        bone_colors['active'] = [*bone_color.color.custom.active, 1]
        bone_colors['hide'] = [*normal_color, 0.1]

        return bone_colors

    def draw(self):
        gpu.state.blend_set('ALPHA')

        if self.hide:
            self.shader.uniform_float("color", (*self.color[:3], 0.4))
            self.p_batch.draw(self.shader)
        #elif self.select:
        #    self.shader.uniform_float("color", self.bone_color)
        #    self.p_batch.draw(self.shader)
        else:
            super().draw()

            # Overlay the fill slightly with the bone color
            #self.shader.uniform_float("color", (*self.bone_color[:3], 0.1))
            #self.p_batch.draw(self.shader)

        if self.select:
            color = self.hover_color

        if self.select or self.hover:
            color = self.hover_color
            
            if self.select and self.hover:
                color = self.active_color

            self.shader.uniform_float("color", color)
            self.p_batch.draw(self.shader)
            
        #Overlay the fill slightly with the bone color
        self.shader.uniform_float("color", (*self.bone_colors['normal'][:3], 0.1))
        self.p_batch.draw(self.shader)
        gpu.state.blend_set('NONE')


        #self.contour_shader.bind()

        #print(self.bone_color)
        if self.select or self.active:
            gpu.state.line_width_set(2.0)

        #if not self.hide:
        self.shader.uniform_float("color", self.bone_color)
        #for b in self.contour_batches:
        self.e_batch.draw(self.shader)

        gpu.state.line_width_set(1.0)
    
    def assign_bone_event(self):

        #print('assign_bone_event', self)
        scn = bpy.context.scene
        rig = scn.rig_picker.rig
        source_object = scn.objects.get(self.source_name)
        if not source_object:
            print(f'Source object {self.source_name} not found')
            return
        
        active_bone = rig.data.bones.active
        if not active_bone:
            print('You need to have an active bone')
            return

        source_object.rig_picker.name = rig.data.bones.active.name

    def release_event(self, mode='SET'):
        super().release_event(mode)

        if self.hide or not self.bone:
            return

        select = True
        if mode == 'SUBSTRACT':
            select = False

        self.bone.bone.select = select
        if self.hover:

            if mode != 'SUBSTRACT':
                self.rig.data.bones.active = self.bone.bone


    def border_select(self, border, mode='SET'):
        if not self.bone:
            return

        if self.hide:
            self.bone.bone.select = False
            return

        rect_tri1 = self.rect[0], self.rect[1], self.rect[2]
        rect_tri2 = self.rect[2], self.rect[3], self.rect[0]

        border_tri1 = border[0], border[1], border[2]
        border_tri2 = border[2], border[3], border[0]

        if (not intersect_tri_tri_2d(*border_tri1, *rect_tri1) and
            not intersect_tri_tri_2d(*border_tri1, *rect_tri2) and 
            not intersect_tri_tri_2d(*border_tri2, *rect_tri1) and 
            not intersect_tri_tri_2d(*border_tri2, *rect_tri2)):
            return

        select = True
        if mode == 'SUBSTRACT':
            select = False
        
        for polygon in self.polygons:
            points = [self.points[i] for i in polygon]

            if intersect_tri_tri_2d(*border_tri1, *points):
                self.bone.bone.select = select
                return
            
            if intersect_tri_tri_2d(*border_tri2, *points):
                self.bone.bone.select = select
                return

            '''
            for b in border:
                if intersect_point_tri_2d(b, *points):
                    self.bone.bone.select = select
                    return
                
            for p in points:
                if intersect_point_quad_2d(p, *border):
                    self.bone.bone.select = select
                    return
            '''

class OperatorShape(Shape):
    def __init__(self, picker, points, polygons, operator, tooltip='', color=None, source_name=''):
        super().__init__(picker, points=points, polygons=polygons, tooltip=tooltip, 
            color=color, source_name=source_name)

        self.type = 'operator'
        self.active_color = [1, 1, 1, 0.15]
        self.press_color = [0, 0, 0, 0.25]
        self.operator = operator
        #self.arguments = arguments#{k:eval(v)}
        #self.operator = get_operator_from_id(self.operator)

        if not tooltip:
            self.tooltip = self.operator.replace('bpy.ops.', '').replace("'INVOKE_DEFAULT', ", '')

        #self.reg_args = re.compile(r'(\w+)=')
    '''
    def parse_args(self):
        
        args = self.reg_args.split(self.arguments)[1:]
        #print(args, zip(args[::2], args[1::2]))
        return {k: eval(v) for k, v in zip(args[::2], args[1::2])}
        #return {k:eval(v) for k, v in self.reg_args.split(self.arguments)}
        '''

    def release_event(self, mode='SET'):
        super().release_event(mode)

        #args = self.parse_args()
        if not self.operator:
            return

        exec(self.operator)
        #f'bpy.ops;{idname}'

        #print(self.idname)
        #print(self.arguments)
        #else:
        #    self.bone.bone.select = False

    def draw(self):
        super().draw()

        if self.press:
            color = self.press_color
        elif self.hover:
            color = self.hover_color
        else:
            return

        gpu.state.blend_set('ALPHA')
        self.shader.uniform_float("color", color)
        self.p_batch.draw(self.shader)
        gpu.state.blend_set('NONE')


class Picker:
    def __init__(self, rig, shapes):

        self.region = bpy.context.region
        self.rig = rig

        self.translation = Vector((0, 0))
        self.shapes = []
        self.box_select = None
        self.hover_shape = None
        self.shader = gpu.shader.from_builtin('UNIFORM_COLOR')

        self.mouse = None

        for shape_data in shapes:
            if not shape_data['points']:
                continue
                
            if shape_data['type'] in ('CANVAS', 'DISPLAY'):
                shape = Shape(
                    self,
                    points=shape_data['points'], 
                    polygons=shape_data['polygons'], 
                    edges=shape_data['edges'], 
                    color=shape_data['color']
                )

            elif shape_data['type'] == 'BONE':
                bone = rig.pose.bones.get(shape_data['bone'])
                #if not bone:
                #    print(f'Bone {shape_data["bone"]} not exist')
                #    continue

                shape = BoneShape(
                    self,
                    source_name=shape_data['source_name'],
                    points=shape_data['points'], 
                    polygons=shape_data['polygons'], 
                    edges=shape_data['edges'], 
                    bone=bone, 
                    color=shape_data['color']
                )

            elif shape_data['type'] == 'OPERATOR':
                shape = OperatorShape(
                    self,
                    source_name=shape_data['source_name'],
                    points=shape_data['points'], 
                    polygons=shape_data['polygons'],
                    operator=shape_data['operator'],
                    color=shape_data['color'],
                    tooltip=shape_data['tooltip'], 
                )

            self.shapes.append(shape)

        self.rect = bounding_rect([p for s in self.shapes for p in s.points])

    def assign_bone_event(self):
        for shape in self.shapes:
            if shape.type=='bone' and shape.hover:
                shape.assign_bone_event()
        
        bpy.ops.rigpicker.save_picker()

    def press_event(self, mode='SET'):
        for shape in self.shapes:
            #print(s)
            if shape.hover:
                shape.press_event(mode)
            else:
                shape.press = False

    def release_event(self, mode='SET'):


            #bpy.ops.pose.select_all(action='DESELECT')

        #print('PICKER release event', mode)
        #print(f'type={event.type}, value={event.value}, ctrl={event.ctrl}, shift={event.shift}, alt={event.alt}')

        for shape in self.shapes:
            if shape.hover:
                shape.release_event(mode)
            
            shape.press = False
        
        #bpy.context.area.tag_redraw()

        '''
        picker.tooltip_event(event='SHOW')
        region.tag_redraw()



            #context.region.tag_redraw()

        picker.tooltip_event(event='HIDE')
        bpy.app.timers.register(partial(tooltip, context.region), first_interval=1)
        '''
    '''
    def tooltip_event(self, event):



        self.tooltip = ''

        if event == 'SHOW':
            if self.hover_shape.type == 'bone':
                self.tooltip = self.hover_shape.bone.name
   
        #bpy.context.region.tag_redraw()
        '''

    def border_select(self, border, mode):
        border = [Vector(b) - Vector(self.translation) for b in border]



        for shape in self.shapes:
            if shape.type != 'bone':
                continue

            shape.border_select(border, mode)

    def move_event(self, location):
        self.mouse = Vector(location) - Vector(self.translation)

        self.hover_shape = None
        for shape in reversed(self.shapes):
            if self.hover_shape:
                shape.hover = False
            elif shape.move_event(self.mouse):
                self.hover_shape = shape

        #if point_inside_rect(self.end, rect):
        #    over = point_over_shape(self.end,points, edges)

        #if bpy.app.timers.is_registered(self.tooltip_event):
        #try:
        #    bpy.app.timers.unregister(self.tooltip_event)
        #except:
        #    pass

        #bpy.app.timers.register(self.tooltip_event, first_interval=1)

    def draw(self):
        with gpu.matrix.push_pop():
            gpu.matrix.translate(self.translation)
            for shape in self.shapes:
                shape.draw()

    def under_mouse(self, location):
        location = Vector(location) - Vector(self.translation)
        return intersect_point_quad_2d(location, *self.rect)

    @property
    def center(self):
        relative_center = (self.rect[1] + self.rect[3]) * 0.5
        return relative_center + self.translation


class PickerGroup:
    def __init__(self, rig, pickers):

        self.view_location = Vector((0, 0))
        self.region = bpy.context.region
        self.rig = rig
        self.pickers = []

        self.mouse = Vector((0, 0))

        self.hover_shape = None

        self.tooltip_shape = None
        self.tooltip_mouse = Vector((0, 0))
        self.tooltip = ''

        self.timer = None

        for picker in pickers:
            self.add_picker(picker)
        
    def add_picker(self, picker):
        self.pickers.append(Picker(self.rig, picker))

    def draw(self):
        y = 0
        for picker in self.pickers:
            height = picker.rect[1][1] - picker.rect[-1][1]

            picker.translation = Vector((0, -y-height*0.5))
            picker.draw()

            y += height + 50

            #break #TODO for now only draw first picker

    def move_event(self, mouse):
        self.mouse = mouse
        location = self.region.view2d.region_to_view(*mouse)

        view_location = self.region.view2d.region_to_view(0, 0)
        if view_location != self.view_location:
            self.view_location = view_location
            self.tooltip = ''
            if self.timer:
                self.timer.cancel()
            return

        hover_shape = None
        for picker in self.pickers:
            if not picker.under_mouse(location):
                continue

            picker.move_event(location)
            
            if picker.hover_shape:
                hover_shape = picker.hover_shape

        self.hover_shape = hover_shape

        if self.tooltip_shape is not self.hover_shape:
            self.tooltip = ''

        if self.timer:
            self.timer.cancel()

        self.timer = threading.Timer(0.5, self.tooltip_event)
        self.timer.start()

    def press_event(self, mode='SET'):
        self.clear_tooltip()

        for picker in self.pickers:
            picker.press_event(mode)

    def release_event(self, mode='SET'):

        if mode == 'SET':
            for bone in self.rig.data.bones:
                bone.select = False

        for picker in self.pickers:
            picker.release_event(mode)

    def border_select(self, border, mode):
        border = [self.region.view2d.region_to_view(*b) for b in border]

        if mode == 'SET':
            for bone in self.rig.data.bones:
                bone.select = False

        for picker in self.pickers:
            picker.border_select(border, mode)

    def clear_tooltip(self):
        self.tooltip = ''
        self.tooltip_shape = None
        self.timer.cancel()
        self.region.tag_redraw()
    
    def tooltip_event(self):
        #print('Tooltip Event', self)

        #print(self.hover_shape, self.hover_shape.type)

        if self.hover_shape and self.hover_shape.type != 'display':
            if self.hover_shape.type == 'bone':
                self.tooltip = self.hover_shape.bone.name
            else:
                self.tooltip = self.hover_shape.tooltip
            self.tooltip_shape = self.hover_shape
        else:
            return self.clear_tooltip()
            

        self.tooltip_mouse = self.mouse

        self.timer.cancel()

        #print(self.tooltip)

        self.region.tag_redraw()

    @property
    def rect(self):
        return bounding_rect([co - Vector(p.translation) for p in self.pickers for co in p.rect])

    @property
    def center(self):
        center = sum(self.rect, Vector((0, 0))) / len(self.rect)
        center[1] = -center[1]
        return center


def get_picker_path(rig, start=None):
    picker_path = rig.data.get('rig_picker', {}).get('source')
    if not picker_path:
        return

    picker_path = bpy.path.abspath(picker_path, library=rig.data.library, start=start)

    return Path(os.path.abspath(picker_path))


def pack_picker(rig, start=None):
    picker_path = get_picker_path(rig, start=start)
    if picker_path and picker_path.exists():
        if 'rig_picker' not in rig.data.keys():
            rig.data['rig_picker'] = {}

        rig.data['rig_picker']['picker'] = json.loads(picker_path.read_text(encoding='utf-8'))


def unpack_picker(rig):
    if 'rig_picker' not in rig.data.keys():
        return

    if 'picker' in rig.data['rig_picker'].keys():
        del rig.data['rig_picker']['picker']