auto_walk/OP_expand_cycle_step.py

from doctest import SKIP
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, end=None):
    print(fn.helper())
    end = end or bpy.context.scene.frame_end

    print('end: ', end)
    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_baked_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

    # all_keys = [k.co.x for fc in act.fcurves if not '.offset' in fc.data_path for k in fc.keyframe_points]
    all_keys = fn.get_x_pos_of_visible_keys(obj, act)
    first = min(all_keys) # int(min(all_keys))
    last = max(all_keys) # int(max(all_keys))
    offset = last - first

    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)

        current_offset = offset

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

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

        ## delete last after computing offset IF cycle have first frame repeatead as last !
        
        # fcu_kfs_without_last = fcu_kfs.copy()
        # fcu_kfs_without_last.pop()

        # last_kf = fcu_kfs.pop() (or just iterate with slicing [:-1])

        if debug >= 2: print('keys', len(fcu_kfs))

        iterations = int( ((end - last) // offset) + 1 )
        if debug >= 2: print('iterations: ', iterations)
        
        for i in range(iterations):
            # if i == iterations - 1: # last
            #     kfs = fcu_kfs
            # else:
            #     kfs = fcu_kfs_without_last
            # for kf in kfs:

            for count, kf in enumerate(fcu_kfs):
                if count == keys_num and i < iterations - 1:
                   # last key of fcurves, to use only if on last iteration
                   continue

                # 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
        re_baked = re.compile(r'_baked\.?\d{0,3}$')
        if re_baked.search(act.name):
            # is an autogenerated one
            org_action_name = re_baked.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 baked action)')
            obj.animation_data.action = org_action
        return ('ERROR', 'No fcurve with cyclic modifier found (used to determine what to bake)')

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

    # 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())

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

    act = fn.get_obj_action(ob)
    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))

    offset_fc = None
    for fc in act.fcurves:
        if all(x in fc.data_path for x in ('pose.bones', 'constraints', 'offset')):
            offset_fc = fc
    
    if not offset_fc:
        return ('ERROR', f'no offset animation in action {act.name}')
    
    data_path = offset_fc.data_path
    const = ob.path_resolve(data_path.rsplit('.', 1)[0])

    timevalues = [offset_fc.evaluate(kf) for kf in keyframes]
    for kf, value in zip(keyframes, timevalues):
        ## or use t_fcu.keyframe_points.add(len(kf))


        const.offset = value
        const.keyframe_insert('offset', frame=kf, options={'INSERTKEY_AVAILABLE'})
        # ``INSERTKEY_NEEDED````INSERTKEY_AVAILABLE`` (only available channels)

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

class AW_OT_bake_cycle_and_step(bpy.types.Operator):
    bl_idname = "autowalk.bake_cycle_and_step"
    bl_label = "Bake keys"
    bl_description = "Bake the keys to a new baked animation\
        \nStep path animation according to those key (if not in Linear)\
        \n(duplicate to a new 'baked' action)"
    bl_options = {"REGISTER", "UNDO"}

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

    end_frame : bpy.props.IntProperty(name='End Frame',options={'SKIP_SAVE'})
    
    def invoke(self,context,event):
        self.end_frame = context.scene.frame_end
        act = fn.get_obj_action(context.object)
        if not act:
            self.report({'ERROR'}, 'No Animation set on active object')
            return {"CANCELLED"}

        act = fn.get_origin_action(act)
        # all_keys = [k.co.x for fc in act.fcurves if not '.offset' in fc.data_path for k in fc.keyframe_points]
        all_keys = fn.get_x_pos_of_visible_keys(context.object, act) # no offset and only visible bone layers
        self.first = min(all_keys) # int(min(all_keys))
        self.last = max(all_keys) # int(max(all_keys))
        self.offset = self.last - self.first
        # return self.execute(context) # uncomment only this to skip pop-up and keep scene.end
        return context.window_manager.invoke_props_dialog(self) # width=400

    def draw(self, context):
        layout = self.layout
        layout.use_property_split = True
        layout.label(text='End of cycle duplication')
        layout.prop(self, 'end_frame', text='End Frame')
        iteration = ((self.end_frame - self.first) // self.offset) + 1
        real_end_cycle = iteration * self.offset + self.first
        layout.label(text=f'Cycle will stop at frame: {real_end_cycle}')

    def execute(self, context):
        
        err = bake_cycle(context.scene.anim_cycle_settings.expand_on_selected_bones, end=self.end_frame)
        if err:
            self.report({err[0]}, err[1])
            if err[0] == 'ERROR':
                return {"CANCELLED"}
        
        ## all followup is not needed when animating on one

        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)
            act = fn.get_obj_action(context.object)
            if act:
                timef = next((fc for fc in act.fcurves if fc.data_path.endswith('.offset') and 'constraint' in fc.data_path), 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
    debug = fn.get_addon_prefs().debug
    scn = bpy.context.scene
    # Delete current action if its not the main one
    # create a new '_pinned' one 
    act = fn.set_generated_action(obj)
    if not act:
        return ('ERROR', f'No action on {obj.name}')

    if debug:
        print('action', act.name)

    act = obj.animation_data.action

    to_change_list = [
        (bpy.context.scene, 'frame_current'),
        (bpy.context.scene, 'frame_start', bpy.context.scene.frame_start-100),
        (bpy.context.scene, 'frame_end', bpy.context.scene.frame_end+100),
        (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 = {}
        viewed_data_paths = []
        for fcu in act.fcurves:
            
            # check only location
            if not fcu.data_path.endswith('.location'):
                continue
            
            all_extremes = [k for k in fcu.keyframe_points if k.type == 'EXTREME']
            if not all_extremes:
                if debug: print(f'skip (no extremes marks): {fcu.data_path}')
                continue

            if len(all_extremes) == len(fcu.keyframe_points):
                if debug: print(f'skip (only extremes): {fcu.data_path}')
                continue

            # skip same data path with different array index to avoid multiple iteration
            if fcu.data_path in viewed_data_paths:
                if debug: print(f'skip (already evaluated): {fcu.data_path}')
                continue
            
            # TODO might need to filter the other array index (using a find)
            # to do the action only on the one that has the best set of markers
            # maybe check the wider values ? (meanning it's in direction of move)

            viewed_data_paths.append(fcu.data_path)

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

            b_name = fcu.data_path.split('"')[1] # print('b_name: ', b_name, fcu.is_valid)
            
            ## TEST : only foot bones
            if not 'foot' in b_name:
                continue

            pb = obj.pose.bones.get(b_name)
            if not pb:
                print(f'{b_name} is invalid')
                continue
            
            contact_ranges = []
            start_contact = None
            prev = None
            for k in fcu.keyframe_points:
                if k.type == 'EXTREME':
                    prev = k
                    if start_contact is None:
                        start_contact = int(k.co.x)
                        continue
                else:
                    if start_contact is not None:
                        # print(f'contact range {start_contact} - {k.co.x:.0f}')
                        if prev:
                            contact_ranges.append((start_contact, int(prev.co.x)))
                        start_contact = None
                        prev = None
            
            if not contact_ranges:
                if debug >= 2: print(f'SKIP (no extreme keys): {b_name} > {prop}')
                continue

            if debug: print(f'fcurve: {b_name} > {prop}')

             # iterate in reverse ranges (not really necessary)
            for r in reversed(contact_ranges):
                print(f'range: {r}')
                first = True
                for i in range(r[0], r[1]+1)[::-1]: # start from the end of the range
                # for i in range(r[0], r[1]+1):
                    bpy.context.scene.frame_set(i)
                    if first:
                        # record coordinate relative to referent object (or world coord)
                        bone_mat = pb.matrix.copy()
                        # bone_mat = obj.matrix_world @ pb.matrix.copy()
                        first = False
                        continue

                    # print(f'Apply on {b_name} at {i}')

                    #-# assign previous matrix
                    # pbl = pb.location.copy()
                    
                    pb.matrix = bone_mat # Exact same position
                    
                    ## maybe align on a specific axis
                    # pb.location.x = pbl.x # dont touch x either

                    #pb.location.z = pbl.z # Z is not necessarily up in local axis, need to check first
                    
                    # 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')

                    # if i == r[1]+1: # (last key) in normal
                    # if i == r[0]: # (last key) in reverse
                    #     continue

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

        print(f'--\n{ct} keys changed/added 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 set to 0 >> 9.57s

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


class AW_OT_pin_feets(bpy.types.Operator):
    bl_idname = "autowalk.pin_feets"
    bl_label = "Pin Feets"
    bl_description = "Pin feets on keys marked as extreme\n(duplicate to a new 'pinned' action)"
    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"}


# --- Quick action management

class AW_OT_set_action(bpy.types.Operator):
    bl_idname = "autowalk.set_action"
    bl_label = "Set action by name"
    bl_description = "Set action on active object using passed name"
    bl_options = {"REGISTER", "INTERNAL"}
    
    act_name : bpy.props.StringProperty(options={'SKIP_SAVE'})

    def execute(self, context):
        act = bpy.data.actions.get(self.act_name)
        if not act:
            self.report({'ERROR'}, f'Could not find action {self.act_name} in bpy.data.actions')
            return {"CANCELLED"}
    
        context.object.animation_data.action = act
        return {"FINISHED"}

class AW_OT_step_back_actions(bpy.types.Operator):
    bl_idname = "autowalk.step_back_actions"
    bl_label = "Actions Step Back"
    bl_description = "Step back to a previous action if 'baked' or 'pinned' action are not ok"
    bl_options = {"REGISTER", "INTERNAL", "UNDO"}

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

    def invoke(self, context, event):
        if context.object.animation_data.use_tweak_mode:
            self.report({'ERROR'}, f'Cannot access animation in NLA')
            return {"CANCELLED"}

        act = context.object.animation_data.action
        base_name = act.name.replace('_baked', '').replace('_pinned', '')
        base_name = re.sub(r'\.\d{3}', '', base_name) # remove duplicate to search everything that has the same base

        # self.actions = [a for a in bpy.data.actions if a.name.startswith(base_name) and not a.name == act.name] # skip current action
        self.actions = [a for a in bpy.data.actions if a.name.startswith(base_name)]

        if not len(self.actions):
            self.report({'ERROR'}, f'no other action found for {act.name}\nUsing basename{base_name}')
            return {'CANCELLED'}
        
        if len(self.actions) == 1:
            context.object.animation_data.action = self.actions[0]
            return self.execute(context)

        self.actions.sort(key=lambda x: len(x.name))
        return context.window_manager.invoke_props_popup(self, event)


    def draw(self, context):
        layout = self.layout
        # layout.label(text=f'Current Action: {context.object.animation_data.action.name}')
        layout.label(text='Actions with same name base:')
        for a in self.actions:
            if a == context.object.animation_data.action:
                layout.label(text=f'(current) >> {a.name}', icon='ACTION')
                continue
            layout.operator('autowalk.set_action', text=a.name, icon='ACTION').act_name = a.name

    def execute(self, context):
        return {"FINISHED"}


classes=(
AW_OT_bake_cycle_and_step,
AW_OT_pin_feets,
AW_OT_set_action,
AW_OT_step_back_actions,
)

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

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