auto_walk/OP_expand_cycle_step.py

import bpy, re
from . import fn
from time import time

## step 3
# - Bake cycle modifier keys -chained with- step the animation path
# - Pin the feet (separated ops)

def bake_cycle(on_selection=True):
    print(fn.helper())
    debug = fn.get_addon_prefs().debug
    obj = bpy.context.object
    if obj.type != 'ARMATURE':
        print('ERROR', 'active is not an armature type')
        return

    act = fn.set_expanded_action(obj)
    if not act:
        return
    
    if debug: print('action', act.name)

    # obj.animation_data.action = act

    ct_fcu = len(act.fcurves)
    ct = 0
    ct_no_cycle = 0

    for fcu in act.fcurves:
        
        ## if a curve is not cycled don't touch
        if not [m for m in fcu.modifiers if m.type == 'CYCLES']:
            ct_no_cycle += 1
            continue
        
        if debug: print(fcu.data_path, 'has cycle')
        #-# only on location :
        # if not fcu.data_path.endswith('.location'):
        #     continue

        # prop = fcu.data_path.split('.')[-1]

        b_name = fcu.data_path.split('"')[1]
        if debug: print(b_name, 'has cycle')
        pb = obj.pose.bones.get(b_name)
        if not pb:
            print(f'{b_name} is invalid')
            continue
        
        #-# limit on selection if passed
        if on_selection and not pb.bone.select:
            continue
        
        if debug: print(b_name, 'seems ok')
        
        #-# only on selected and visible curve
        # if not fcu.select or fcu.hide:
        #     continue


        fcu_kfs = []
        for k in fcu.keyframe_points:
            k_dic = {}
            # k_dic['k'] = k
            k_dic['co'] = k.co
            k_dic['interpolation'] = k.interpolation
            k_dic['type'] = k.type
            fcu_kfs.append(k_dic)

        first = fcu_kfs[0]['co'][0]
        # second = fcu_kfs[1]['co'][0] 
        # before_last= fcu_kfs[-2]['co'][0]
        last = fcu_kfs[-1]['co'][0]
        # first_offset = second - first
        
        current_offset = offset = last - first

        keys_num = len(fcu_kfs)
        if debug: print(keys_num)

        if keys_num <= 1:
            if debug: print(b_name, f'{keys_num} key')
            continue

        ## ! important: delete last after computing offset IF cycle have first frame repeatead as last !
        fcu_kfs.pop()
        # print('offset: ', offset)

        if debug: print('keys', len(fcu_kfs))
        ## expand to end frame
        
        end = bpy.context.scene.frame_end # maybe add possibility define target manually
        iterations = ((end - last) // offset) + 1
        if debug: print('iterations: ', iterations)
        for _i in range(int(iterations)):
            for kf in fcu_kfs:
                # create a new key, adding offset to keys  
                fcu.keyframe_points.add(1)
                new = fcu.keyframe_points[-1]
                for att, val in kf.items():
                    if att == 'co':
                        new.co = (val[0] + current_offset, val[1])
                    else:
                        setattr(new, att, val)
            current_offset += offset

        ct += 1
    
    if ct_fcu == ct_no_cycle: # skipped because no cycle exists
        rexpand = re.compile(r'_expanded\.?\d{0,3}$')
        if rexpand.search(act.name):
            # is an autogenerated one
            org_action_name = rexpand.sub('', act.name)
            org_action = bpy.data.actions.get(org_action_name)
            if not org_action:
                return ('ERROR', 'No fcurve with anim cycle found (on expanded action)')
            obj.animation_data.action = org_action
        return ('ERROR', 'No fcurve with anim cycle found (back to unexpanded)')

    if not ct:
        return ('ERROR', 'No fcurve treated (! action duplicated to _expand !)')

    # cleaning update 
    fn.update_action(act)
    print('end of anim cycle keys baking')
    # C.scene.frame_current = org_frame
        # detect last key in contact


def step_path():
    '''Step the path anim of the curve to constant'''
    print(fn.helper())

    obj = bpy.context.object
    if obj.type != 'ARMATURE':
        return ('ERROR', 'active is not an armature type')

    # found curve through constraint
    curve, const = fn.get_follow_curve_from_armature(obj)
    if not const:
        return ('ERROR', 'No constraints found')

    act = fn.get_obj_action(obj)
    if not act: return
    
    # CHANGE - removed int from frame
    # keyframes = [int(k.co[0]) for fcu in act.fcurves for k in fcu.keyframe_points]
    keyframes = [k.co[0] for fcu in act.fcurves for k in fcu.keyframe_points]
    keyframes = list(set(keyframes))
    
    curve = const.target
    if not curve:
        return ('ERROR', f'no target set for {curve.name}')

    # get a new generated action for the curve
    # Follow path animation is on the DATA of the fcurve
    fact = fn.set_generated_action(curve.data)
    if not fact:
        return
    
    t_fcu = False
    for fcu in fact.fcurves:
        ## fcu data_path is just a string
        if fcu.data_path == 'eval_time':
            t_fcu = fcu
    
    if not t_fcu:
        return ('ERROR', f'no eval_time animation in {curve.name}')
    
    timevalues = [t_fcu.evaluate(kf) for kf in keyframes]
    for kf, value in zip(keyframes, timevalues):
        ## or use t_fcu.keyframe_points.add(len(kf))
        curve.data.eval_time = value
        curve.data.keyframe_insert('eval_time', frame=kf, options={'INSERTKEY_AVAILABLE'})
        # ``INSERTKEY_NEEDED````INSERTKEY_AVAILABLE`` (only available channels)

    ## set all to constant
    for k in t_fcu.keyframe_points:
        k.interpolation = 'CONSTANT'
    
    # cleaning update (might not be needed here)
    fn.update_action(act)
    print('end of step_anim')

class UAC_OT_bake_cycle_and_step(bpy.types.Operator):
    bl_idname = "anim.bake_cycle_and_step"
    bl_label = "Bake key and step path "
    bl_description = "Bake the key and step the animation path according to those key"
    bl_options = {"REGISTER", "UNDO"}

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

    def execute(self, context):
        
        err = bake_cycle(context.scene.anim_cycle_settings.expand_on_selected_bones)
        if err:
            self.report({err[0]}, err[1])
            if err[0] == 'ERROR':
                return {"CANCELLED"}
        
        if not context.scene.anim_cycle_settings.linear:
            # CHAINED ACTION : step the path of the curve path
            err = step_path()
            if err:
                self.report({err[0]}, err[1])
                if err[0] == 'ERROR':
                    return {"CANCELLED"}
        else:
            # Delete points in curve action between first and last and go LINEAR
            curve = context.scene.anim_cycle_settings.path_to_follow
            if curve:
                act = fn.get_obj_action(curve.data)
                if act:
                    timef = next((fcu for fcu in act.fcurves if fcu.data_path == 'eval_time'), None)
                    if timef:
                        keys_ct = len(timef.keyframe_points)
                        if keys_ct > 2:
                            for k in reversed(timef.keyframe_points[1:-2]):
                                timef.keyframe_points.remove(k)
                            for k in timef.keyframe_points:
                                k.interpolation = 'LINEAR'
                            print(f'Anim path to linear : Deleted all keys ({keys_ct - 2}) on anim path except first and last')
        
        # CHAINED ACTION pin feet ?? : Step the path of the curve path

        return {"FINISHED"}


# detect contact key
# [ 'array_index', 'auto_smoothing', 'bl_rna', 'color', 'color_mode', 'convert_to_keyframes', 'convert_to_samples', 'data_path',
#  'driver', 'evaluate', 'extrapolation', 'group', 'hide', 'is_empty', 'is_valid', 'keyframe_points', 'lock', 'modifiers', 'mute',
#  'range', 'rna_type', 'sampled_points', 'select', 'update', 'update_autoflags']


def pin_down_feets():
    print(fn.helper())
    obj = bpy.context.object
    if obj.type != 'ARMATURE':
        print('ERROR', 'active is not an armature type')
        return

    # Delete current action if its not the main one
    # create a new '_autogen' one 
    act = fn.set_generated_action(obj)
    if not act:
        return ('ERROR', f'No action on {obj.name}')
    print('action', act.name)

    act = obj.animation_data.action

    to_change_list = [
        (bpy.context.scene, 'frame_current', '_undefined'), # use '_undefined' when no value to assign for now
        (bpy.context.scene.render, 'use_simplify', True),
        (bpy.context.scene.render, 'simplify_subdivision', 0),
    ]
    ## Link armature in a new collection and exclude all the others

    ## STORE for context manager store/restore  /-
    tmp_col = bpy.data.collections.get('TMP_COLLECTION_PINNING')
    if not tmp_col:
        tmp_col = bpy.data.collections.new('TMP_COLLECTION_PINNING')
    if tmp_col.name not in bpy.context.scene.collection.children:
        bpy.context.scene.collection.children.link(tmp_col)
    if obj not in tmp_col.objects[:]:
        tmp_col.objects.link(obj)

    for vlc in bpy.context.view_layer.layer_collection.children:
        if vlc.collection == tmp_col:
            continue
        to_change_list.append((vlc, 'exclude', True))
    
    #-/

    with fn.attr_set(to_change_list):
        t0 = time()
        ct = 0
        done = {}
        for fcu in act.fcurves:

            # check only location
            if not fcu.data_path.endswith('.location'):
                continue

            # prop = fcu.data_path.split('.')[-1]

            b_name = fcu.data_path.split('"')[1]
            # print('b_name: ', b_name, fcu.is_valid)

            pb = obj.pose.bones.get(b_name)
            if not pb:
                print(f'{b_name} is invalid')
                continue

            start_contact = None
            for k in fcu.keyframe_points:

                if k.type == 'EXTREME':
                    bpy.context.scene.frame_set(int(k.co[0]))
                    if start_contact is None:
                        start_contact=k.co[0]
                        # record coordinate relative to referent object (or world coord)
                        bone_mat = pb.matrix.copy()
                        # bone_mat = obj.matrix_world @ pb.matrix.copy()
                        continue
                    
                    if b_name in done.keys():
                        if k.co[0] in done[b_name]:
                            continue
                    else:
                        # mark as treated (all curve of this bone at this time)
                        done[b_name] = [k.co[0]]

                    #-# Insert keyframe to match Hold position
                    # print(f'Apply on {b_name} at {k.co[0]}')
                    
                    #-# assign previous matrix
                    pbl = pb.location.copy()
                    # l, _r, _s = bone_mat.decompose()
                    
                    pb.matrix = bone_mat # Exact same position
                    # pb.location.x = pbl.x # dont touch x either
                    pb.location.z = pbl.z
                    # pb.location.y = l.y (weirdly not working)

                    # bpy.context.view_layer.update()

                    #-# moyenne des 2 ?
                    # pb.location, pb.rotation_euler, pb.scale = average_two_matrix(pb.matrix, bone_mat) ## marche pas du tout !
                    
                    ## insert keyframe
                    pb.keyframe_insert('location')
                    # only touched Y location
                    # pb.keyframe_insert('rotation_euler')

                    k.type = 'JITTER' # 'BREAKDOWN' 'MOVING_HOLD' 'JITTER'
                    ct += 1

                else:
                    if start_contact is not None:
                        # print('fcu.data_path: ', fcu.data_path, fcu.array_index)
                        # print(f'{b_name} contact range {start_contact} - {k.co[0]}')
                        start_contact = None
                        
                    # print(i, fcu.data_path, fcu.array_index)
                # print('time', k.co[0], '- value', k.co[1])
                
                #k.handle_left
                #k.handle_right
                ##change handler type ([‘FREE’, ‘VECTOR’, ‘ALIGNED’, ‘AUTO’, ‘AUTO_CLAMPED’], default ‘FREE’)
                #k.handle_left_type = 'AUTO_CLAMPED'
                #k.handle_right_type = 'AUTO_CLAMPED'

        print(f'--\n{ct} keys changed in {time()-t0:.2f}s\n--') # fcurves treated in 

    ## RESTORE
    # without >> 433 keys changed in 29.15s
    # with all collection excluded >> 433 keys changed in 25.00s
    # with simplify >> 9.57s

    tmp_col.objects.unlink(obj)
    bpy.data.collections.remove(tmp_col)

  


class UAC_OT_pin_feets(bpy.types.Operator):
    bl_idname = "anim.pin_feets"
    bl_label = "Pin Feets"
    bl_description = "Pin feets on keys marked as extreme"
    bl_options = {"REGISTER", "UNDO"}

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

    def execute(self, context):
        # context.scene.anim_cycle_settings.expand_on_selected_bones
        err = pin_down_feets()
        if err:
            self.report({err[0]}, err[1])
            if err[0] == 'ERROR':
                return {"CANCELLED"}

        return {"FINISHED"}

classes=(
UAC_OT_bake_cycle_and_step,
UAC_OT_pin_feets,
)

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

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