auto_walk/OP_setup.py

import bpy
import numpy as np
from mathutils import Vector, Quaternion
from math import sin, cos, radians
from . import fn

## all action needed to setup the walk

## Need to know what is the forward axis


## https://www.meccanismocomplesso.org/en/3d-rotations-and-euler-angles-in-python/

class AW_OT_autoset_axis(bpy.types.Operator):
    bl_idname = "autowalk.autoset_axis"
    bl_label = "Auto Set Axis"
    bl_description = "Define forward axis from armature"
    bl_options = {"REGISTER", "UNDO"}

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

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

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

        ## TODO root need to be user defined to assume corner cases
        root = ob.pose.bones.get('world')
        if not root:
            print('No root found')
            return {"CANCELLED"}

        root_pos = ob.matrix_world @ root.matrix.to_translation()

        # gather all .R suffixed bones

        all_right_pos = [ob.matrix_world @ b.matrix.to_translation() for b in ob.pose.bones if b.bone.name.lower().endswith('.r')]

        if not all_right_pos:
            print('No .r suffix found, need to adjust manually')
            return {"CANCELLED"}

        # get median position 
        median_r_pos = (sum(all_right_pos, Vector()) / len(all_right_pos)) - root_pos

        median_r_vec = Vector((median_r_pos[0], median_r_pos[1], 0))
        median_r_vec.normalize()

        # Determine forward direction by rotating 90 degrees CCW on up axis
        guessed_forward = Quaternion(Vector((0,0,1)), radians(90)) @ median_r_vec

        ## TODO determine what axis of root align bests with guessed_forward (minimum 3D angle)

        axises = {'FORWARD_Y':(0,1,0), 'TRACK_NEGATIVE_Y':(0,-1,0), 'FORWARD_Z':(0,0,1), 'TRACK_NEGATIVE_Z':(0,0,-1)} # 'FORWARD_X':(1,0,0), 'TRACK_NEGATIVE_X':(-1,0,0)
        ref_angle = 10
        best_axis = False
        for axis, vec in axises.items(): # (1,0,0), (-1,0,0)
            angle = guessed_forward.angle(ob.matrix_world @ root.matrix @ Vector(vec))
            print(axis, 'angle: ', angle)
            if angle < ref_angle:
                ref_angle = angle
                best_axis = axis
        
        if not best_axis:
            print('No axis found')
            return {"CANCELLED"}

        context.scene.anim_cycle_settings.forward_axis = best_axis
        return {"FINISHED"}

class AW_OT_create_cycles_modifiers(bpy.types.Operator):
    bl_idname = "autowalk.create_cycles_modifiers"
    bl_label = "Add Cycles Modifiers"
    bl_description = "Add cycles modifier on all bones not prefixed [mch, org, def, vis,fdl]\
        \nCtrl + Click : remove ALL Cycles modifiers on all fcurves"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context):
        return context.object and context.object.type == 'ARMATURE'
    
    def invoke(self, context, event):
        self.ctrl = event.ctrl
        return self.execute(context)

    def execute(self, context):
        if self.ctrl:
            fn.remove_all_cycles_modifier(context.object)
        else:
            fn.create_cycle_modifiers(context.object)

        return {"FINISHED"}


classes=(
AW_OT_autoset_axis,
AW_OT_create_cycles_modifiers,
)

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

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