gp_interpolate/interpolate_strokes/operators.py

import bpy
import numpy as np
from time import perf_counter
from mathutils import Vector, Matrix

from gp_interpolate.utils import (matrix_transform,
                                  plane_on_bone,
                                  ray_cast_point,
                                  intersect_with_tesselated_plane,
                                  triangle_normal,
                                  search_square,
                                  get_gp_draw_plane,
                                  create_plane,
                                  attr_set)

from mathutils.geometry import (barycentric_transform,
                                intersect_point_tri,
                                intersect_point_line,
                                intersect_line_plane,
                                tessellate_polygon)


def following_key(forward=True):
    direction = 1 if forward else -1
    cur_frame = bpy.context.scene.frame_current
    settings = bpy.context.scene.gp_interpo_settings
    
    if settings.mode == 'FRAME':
        return cur_frame + (settings.padding * direction)
    
    elif settings.mode == 'GPKEY':
        layers = bpy.context.object.data.layers
        frames = [f.frame_number for l in layers for f in l.frames]
    
    elif settings.mode == 'RIGKEY':
        col = settings.target_collection
        if not col:
            col = bpy.context.scene.collection            
        for arm in [o for o in col.all_objects if o.type == 'ARMATURE']:
            if not arm.animation_data or not arm.animation_data.action:
                continue
            frames = [k.co.x for fc in arm.animation_data.action.fcurves for k in fc.keyframe_points]
    
    if not frames:
        return

    frames.sort()
    if forward:
        new = next((f for f in frames if f > cur_frame), None)
    else:
        below = [f for f in frames if f < cur_frame]
        if not below:
            return
        new = below[-1]

    return int(new)


## TODO: add bake animation to empty for later GP layer parenting

class GP_OT_interpolate_stroke(bpy.types.Operator):
    bl_idname = "gp.interpolate_stroke"
    bl_label = "Interpolate Stroke"
    bl_description = 'Interpolate Stroke'
    bl_options = {'REGISTER', 'UNDO'} 

    @classmethod
    def poll(cls, context):
        if context.active_object and context.object.type == 'GPENCIL'\
            and context.mode in ('EDIT_GPENCIL', 'SCULPT_GPENCIL', 'PAINT_GPENCIL'):
            return True
        cls.poll_message_set("Need a Grease pencil object in Edit or Sculpt mode")
        return False

    @classmethod
    def description(cls, context, properties):
        if properties.next:
            return f"Interpolate Stroke Forward"
        else:
            return f"Interpolate Stroke Backward"

    next : bpy.props.BoolProperty(name='Next', default=True, options={'SKIP_SAVE'})

    def execute(self, context):
        settings = context.scene.gp_interpo_settings
        scn = bpy.context.scene

        # auto_key_status = context.tool_settings.use_keyframe_insert_auto
        # context.tool_settings.use_keyframe_insert_auto = True

        ## Determine on what key to jump
        frame_to_jump = following_key(forward=self.next)
        if frame_to_jump is None:
            self.report({'WARNING'}, 'No keyframe available in this direction')
            return {'CANCELLED'}
        
        gp = context.object

        matrix = np.array(gp.matrix_world, dtype='float64')#.inverted()
        
        col = settings.target_collection
        if not col:
            col = scn.collection

        origin = np.array(scn.camera.matrix_world.to_translation(), 'float64')
        # print('----')

        tgt_strokes = [s for s in gp.data.layers.active.active_frame.strokes if s.select]
        
        ## If nothing selected in sculpt/paint, Select all before triggering
        if not tgt_strokes and context.mode in ('SCULPT_GPENCIL', 'PAINT_GPENCIL'):
            for s in gp.data.layers.active.active_frame.strokes:
                s.select = True
            tgt_strokes = gp.data.layers.active.active_frame.strokes

        if not tgt_strokes:            
            self.report({'ERROR'}, 'No stroke selected !')
            return {'CANCELLED'}

        included_cols = [c.name for c in gp.users_collection]

        if settings.method == 'BONE':
            ## Follow Bone method (WIP)
            if not settings.target_rig or not settings.target_bone:
                self.report({'ERROR'}, 'No Bone selected')
                return {'CANCELLED'}

            included_cols.append('interpolation_tool')

            ## ensure collection and plane exists
            # get/create collection
            col = bpy.data.collections.get('interpolation_tool')
            if not col:
                col = bpy.data.collections.new('interpolation_tool')

            if col.name not in bpy.context.scene.collection.children:
                bpy.context.scene.collection.children.link(col)
                col.hide_viewport = True
            
            # get/create meshplane
            plane = bpy.data.objects.get('interpolation_plane')
            if not plane:
                plane = create_plane(name='interpolation_plane')
                plane.select_set(False)

            if plane.name not in col.objects:
                col.objects.link(plane)
            ## TODO: Ensure the plane is not animated!

        else:
            # Geometry mode
            if col != context.scene.collection:
                included_cols.append(col.name)
            ## Maybe include a plane just behing geo ? probably bad idea

        ## Prepare context manager
        store_list = [
            # (context.view_layer.objects, 'active', gp),
            (context.tool_settings, 'use_keyframe_insert_auto', True),
            # (bpy.context.scene.render, 'simplify_subdivision', 0),
            ]
        
        for vlc in context.view_layer.layer_collection.children:
            store_list.append(
                (vlc, 'exclude', vlc.name not in included_cols),
                # (vlc, 'hide_viewport', vlc.name not in included_cols), # viewport viz
            )

        with attr_set(store_list):
            if settings.method == 'BONE':
                ## replace plane
                bone_plane = plane_on_bone(settings.target_rig.pose.bones.get(settings.target_bone),
                                            arm=settings.target_rig,
                                            set_rotation=settings.use_bone_rotation,
                                            mesh=True)
                ## Set collection visibility
                intercol = bpy.data.collections.get('interpolation_tool')
                vl_col = bpy.context.view_layer.layer_collection.children.get(intercol.name)
                intercol.hide_viewport = vl_col.exclude = vl_col.hide_viewport = False

                # Override collection
                col = intercol
                ## TODO: Hide all other collections
            
            dg = bpy.context.evaluated_depsgraph_get()
            strokes_data = []

            for stroke in tgt_strokes:
                nb_points = len(stroke.points)

                local_co = np.empty(nb_points * 3, dtype='float64')
                stroke.points.foreach_get('co', local_co)
                # local_co_3d = local_co.reshape((nb_points, 3))
                world_co_3d = matrix_transform(local_co.reshape((nb_points, 3)), matrix)

                stroke_data = []
                for i, point in enumerate(stroke.points):
                    point_co_world = world_co_3d[i]
                    
                    object_hit, hit_location, tri, tri_indices = ray_cast_point(point_co_world, origin, dg)
                    if not object_hit or object_hit not in col.all_objects[:]:
                        for square_co in search_square(point_co_world, factor=settings.search_range):
                            object_hit, hit_location, tri, tri_indices = ray_cast_point(square_co, origin, dg)
                            if object_hit and object_hit in col.all_objects[:]:
                                
                                hit_location = intersect_line_plane(origin, point_co_world, tri[0], triangle_normal(*tri))
                                
                                break
                    
                    stroke_data.append((stroke, point_co_world, object_hit, hit_location, tri, tri_indices))

                strokes_data.append(stroke_data)

            # Copy stroke selection, jump frame and paste

            bpy.ops.gpencil.copy()
            scn.frame_set(frame_to_jump)
            plan_co, plane_no = get_gp_draw_plane(gp)
            bpy.ops.gpencil.paste()

            if settings.method == 'BONE':
                bone_plane = plane_on_bone(settings.target_rig.pose.bones.get(settings.target_bone),
                                            arm=settings.target_rig,
                                            set_rotation=settings.use_bone_rotation,
                                            mesh=True)
                
            dg = bpy.context.evaluated_depsgraph_get()
            matrix_inv = np.array(gp.matrix_world.inverted(), dtype='float64')#.inverted()
            new_strokes = gp.data.layers.active.active_frame.strokes[-len(strokes_data):]

            for new_stroke, stroke_data in zip(new_strokes, strokes_data):
                world_co_3d = [] # np.array(len()dtype='float64')#np.
                for stroke, point_co, object_hit, hit_location, tri_a, tri_indices in stroke_data:
                    eval_ob = object_hit.evaluated_get(dg)
                    tri_b = [eval_ob.data.vertices[i].co for i in tri_indices]
                    tri_b = matrix_transform(tri_b, eval_ob.matrix_world)
                    
                    new_loc = barycentric_transform(hit_location, *tri_a, *tri_b)
                    world_co_3d.append(new_loc)

                # Reproject on plane
                new_world_co_3d = [intersect_line_plane(origin, p, plan_co, plane_no) for p in world_co_3d]        
                new_local_co_3d = matrix_transform(new_world_co_3d, matrix_inv)

                nb_points = len(new_stroke.points)
                new_stroke.points.foreach_set('co', new_local_co_3d.reshape(nb_points*3))
                new_stroke.points.update()

        ## Reset autokey status
        # context.tool_settings.use_keyframe_insert_auto = auto_key_status # (Done in context manager)

        ## TODO: Remove plane on the fly
        
        return {'FINISHED'}

"""
class GP_OT_interpolate_stroke(bpy.types.Operator):
    bl_idname = "gp.interpolate_stroke"
    bl_label = "Interpolate Stroke"
    bl_description = 'Interpolate Stroke'
    bl_options = {'REGISTER', 'UNDO'} 

    @classmethod
    def poll(cls, context):
        if context.active_object and context.object.type == 'GPENCIL'\
            and context.mode in ('EDIT_GPENCIL', 'SCULPT_GPENCIL', 'PAINT_GPENCIL'):
            return True
        cls.poll_message_set("Need a Grease pencil object in Edit or Sculpt mode")
        return False

    @classmethod
    def description(cls, context, properties):
        if properties.next:
            return f"Interpolate Stroke Forward"
        else:
            return f"Interpolate Stroke Backward"

    next : bpy.props.BoolProperty(name='Next', default=True, options={'SKIP_SAVE'})

    def execute(self, context):
        settings = context.scene.gp_interpo_settings

        auto_key_status = context.tool_settings.use_keyframe_insert_auto
        context.tool_settings.use_keyframe_insert_auto = True

        ## Determine on what key to jump
        frame_to_jump = following_key(forward=self.next)
        if frame_to_jump is None:
            self.report({'WARNING'}, 'No keyframe available in this direction')
            return {'CANCELLED'}
        
        gp = context.object

        scn = bpy.context.scene
        dg = bpy.context.evaluated_depsgraph_get()
        matrix = np.array(gp.matrix_world, dtype='float64')#.inverted()
        col = settings.target_collection
        if not col:
            col = scn.collection

        origin = np.array(scn.camera.matrix_world.to_translation(), 'float64')
        # print('----')

        tgt_strokes = [s for s in gp.data.layers.active.active_frame.strokes if s.select]
        
        ## If nothing selected in sculpt/paint, Select all before triggering
        if not tgt_strokes and context.mode in ('SCULPT_GPENCIL', 'PAINT_GPENCIL'):
            for s in gp.data.layers.active.active_frame.strokes:
                s.select = True
            tgt_strokes = gp.data.layers.active.active_frame.strokes

        if not tgt_strokes:            
            self.report({'ERROR'}, 'No stroke selected !')
            return {'CANCELLED'}
        
        strokes_data = []
        if settings.method == 'BONE':
            ## Follow Bone method (Full WIP)
            if not settings.target_rig or not settings.target_bone:
                self.report({'ERROR'}, 'No Bone selected')
                return {'CANCELLED'}

            bone_plane = plane_on_bone(settings.target_rig.pose.bones.get(settings.target_bone),
                                        arm=settings.target_rig,
                                        set_rotation=settings.use_bone_rotation)

            strokes_data = []
            for stroke in tgt_strokes:
                nb_points = len(stroke.points)
                local_co = np.empty(nb_points * 3, dtype='float64')
                stroke.points.foreach_get('co', local_co)
                # local_co_3d = local_co.reshape((nb_points, 3))
                world_co_3d = matrix_transform(local_co.reshape((nb_points, 3)), matrix)

                stroke_data = []

                for i, point in enumerate(stroke.points):
                    point_co_world = world_co_3d[i]
                    hit_location, tri, tri_indices = intersect_with_tesselated_plane(point_co_world, origin, bone_plane)
                    ## probably easier to just generate a single vast triangle and use it
                    
                    ## Store same as other method (without object hit)
                    stroke_data.append((stroke, point_co_world, hit_location, tri, tri_indices))
                
                strokes_data.append(stroke_data)

        else:
            ## Geometry method
            for stroke in tgt_strokes:
                nb_points = len(stroke.points)

                local_co = np.empty(nb_points * 3, dtype='float64')
                stroke.points.foreach_get('co', local_co)
                # local_co_3d = local_co.reshape((nb_points, 3))
                world_co_3d = matrix_transform(local_co.reshape((nb_points, 3)), matrix)

                stroke_data = []
                for i, point in enumerate(stroke.points):
                    point_co_world = world_co_3d[i]
                    
                    object_hit, hit_location, tri, tri_indices = ray_cast_point(point_co_world, origin, dg)
                    if not object_hit or object_hit not in col.all_objects[:]:
                        for square_co in search_square(point_co_world, factor=settings.search_range):
                            object_hit, hit_location, tri, tri_indices = ray_cast_point(square_co, origin, dg)
                            if object_hit and object_hit in col.all_objects[:]:
                                
                                hit_location = intersect_line_plane(origin, point_co_world, tri[0], triangle_normal(*tri))
                                
                                break
                    
                    stroke_data.append((stroke, point_co_world, object_hit, hit_location, tri, tri_indices))
                
                strokes_data.append(stroke_data)

        # Copy stroke selection, jump frame and paste

        bpy.ops.gpencil.copy()

        scn.frame_set(frame_to_jump)

        plan_co, plane_no = get_gp_draw_plane(gp)
        
        bpy.ops.gpencil.paste()



        if settings.method == 'BONE':
            matrix_inv = np.array(gp.matrix_world.inverted(), dtype='float64')#.inverted()
            new_strokes = gp.data.layers.active.active_frame.strokes[-len(strokes_data):]

            bone_plane = plane_on_bone(settings.target_rig.pose.bones.get(settings.target_bone),
                                        arm=settings.target_rig,
                                        set_rotation=settings.use_bone_rotation)

            for new_stroke, stroke_data in zip(new_strokes, strokes_data):
                world_co_3d = [] # np.array(len()dtype='float64')#np.
                for stroke, point_co, hit_location, tri_a, tri_indices in stroke_data:
                    tri_b = [bone_plane[i] for i in tri_indices]
                    # tri_b = matrix_transform(tri_b, settings.target_rig.matrix_world)
                    ## rotate tri_b by bone differential angle camera's aim axis ?

                    new_loc = barycentric_transform(hit_location, *tri_a, *tri_b)
                    world_co_3d.append(new_loc)

    
                # Reproject on plane
                new_world_co_3d = [intersect_line_plane(origin, p, plan_co, plane_no) for p in world_co_3d]        
                new_local_co_3d = matrix_transform(new_world_co_3d, matrix_inv)
                
                nb_points = len(new_stroke.points)
                new_stroke.points.foreach_set('co', new_local_co_3d.reshape(nb_points*3))
                new_stroke.points.update()

        else:
            dg = bpy.context.evaluated_depsgraph_get()
            
            matrix_inv = np.array(gp.matrix_world.inverted(), dtype='float64')#.inverted()
            new_strokes = gp.data.layers.active.active_frame.strokes[-len(strokes_data):]

            for new_stroke, stroke_data in zip(new_strokes, strokes_data):
                world_co_3d = [] # np.array(len()dtype='float64')#np.
                for stroke, point_co, object_hit, hit_location, tri_a, tri_indices in stroke_data:
                    eval_ob = object_hit.evaluated_get(dg)
                    tri_b = [eval_ob.data.vertices[i].co for i in tri_indices]
                    tri_b = matrix_transform(tri_b, eval_ob.matrix_world)
                    
                    new_loc = barycentric_transform(hit_location, *tri_a, *tri_b)
                    world_co_3d.append(new_loc)

                # Reproject on plane
                new_world_co_3d = [intersect_line_plane(origin, p, plan_co, plane_no) for p in world_co_3d]        
                new_local_co_3d = matrix_transform(new_world_co_3d, matrix_inv)
                
                nb_points = len(new_stroke.points)
                new_stroke.points.foreach_set('co', new_local_co_3d.reshape(nb_points*3))
                new_stroke.points.update()

        ## Reset autokey status
        context.tool_settings.use_keyframe_insert_auto = auto_key_status

        return {'FINISHED'}
"""

classes = (
    GP_OT_interpolate_stroke,
)

def register():
    for c in classes:
        bpy.utils.register_class(c)
    
        
def unregister():
    for c in reversed(classes):
        bpy.utils.unregister_class(c)