## Wrap an expanded animation to an on place loop
## Define a bone that will be offseted to stay fix above root

import bpy
import math
from mathutils import Vector, Matrix, Euler

def get_bone_transform_at_frame(b, act, frame):
    '''Find every loc, rot, scale values at given frame'''
    transform = {}
    for channel in ('location', 'rotation_euler', 'scale'):
        chan_list = []
        for i in range(3):
            f = act.fcurves.find(f'pose.bones["{b.name}"].{channel}', index=i)
            if not f:
                # print(frame, channel, 'not animated ! using current value') # Dbg
                chan_list.append(getattr(b, channel)) # get current value since not animated
                continue
            chan_list.append(f.evaluate(frame))
        
        # print(frame, b.name, channel, chan_list) # Dbg
        if channel == 'rotation_euler':
            transform[channel] = Euler(chan_list)
        else:
            transform[channel] = Vector(chan_list)

    return transform # loc, rot, scale

def dopesheet_summary(obj_list=None):
    if obj_list is None:
        obj_list = bpy.context.selected_objects
    elif isinstance(obj_list, bpy.types.Object):
        obj_list = [obj_list]

    start = bpy.context.scene.frame_start
    end = bpy.context.scene.frame_end
    frames = []

    for obj in obj_list:
        for fcurve in obj.animation_data.action.fcurves:
            for keyframe_point in fcurve.keyframe_points:
                x, y = keyframe_point.co
                if x >= start and x <= end and x not in frames:
                    frames.append(x)
                    ## for returning an int (import math)
                    #frames.append((math.ceil(x)))
    return sorted(frames)

def pose_bone_frame_summary(bone_list=None, filter_channel=None) -> list[int]:
    '''Return frame numbers where bone(s) have keys
    :bone_list: a single pose bone or a list of pose bone, if not porvided, context.selected_pose_bone
    :filter_channel: str or tuple of str to filter channel (ex: 'location')
    '''
    if bone_list is None:
        bone_list = bpy.context.selected_pose_bone
    elif isinstance(bone_list, bpy.types.PoseBone):
        bone_list = [bone_list]


    start, end = bpy.context.scene.frame_start, bpy.context.scene.frame_end
    frames = []
    for bone in bone_list:
        for fcurve in bone.id_data.animation_data.action.fcurves:
            if not fcurve.data_path.startswith(f'pose.bones["{bone.name}"]'):
                continue

            if filter_channel:
                if not fcurve.data_path.endswith(filter_channel):
                    continue
            # print('fcurve.data_path: ', fcurve.data_path)
                    
            for keyframe_point in fcurve.keyframe_points:
                x, _y = keyframe_point.co
                if x >= start and x <= end and x not in frames:
                    frames.append(x)
                    ## for returning an int (import math)
                    #frames.append((math.ceil(x)))
    return sorted(frames)

def get_all_keyframe(use_only = True):
    sum = set()
    for action in D.actions:
        if use_only and action.use_fake_user and action.users == 1:
            #avoid saved (fake user) but unused actions
            pass
        elif use_only and action.users == 0:
            #avoid 0 user actions
            pass

        else:
            for fcurve in action.fcurves:
                for key in fcurve.keyframe_points:
                    sum.add(key.co[0])
    return sum


def get_keyframes(obj_list):
    keyframes = []
    for obj in obj_list:
        anim = obj.animation_data
        if anim is not None and anim.action is not None:
            for fcu in anim.action.fcurves:
                for keyframe in fcu.keyframe_points:
                    x, y = keyframe.co
                    if x not in keyframes:
                        keyframes.append((math.ceil(x)))
    return keyframes


def scale_matrix_from_vector(scale):
    '''recreate a neutral mat scale'''
    matscale_x = Matrix.Scale(scale[0], 4,(1,0,0))
    matscale_y = Matrix.Scale(scale[1], 4,(0,1,0))
    matscale_z = Matrix.Scale(scale[2], 4,(0,0,1))
    matscale = matscale_x @ matscale_y @ matscale_z
    return matscale


def has_key_at_frame(item, act=None, frame=None, channel='location', verbose=False):
    '''Return True if pose bone has a key at passed frame'''

    if frame is None:
        frame = bpy.context.scene.frame_current
    
    if isinstance(item, bpy.types.Object):
        ## Object
        if act is None:
            act = item.animation_data.action
        data_path = channel
    else:
        ## Consider it's a Pose bone
        if act is None:
            act = item.id_data.animation_data.action
        data_path = f'pose.bones["{item.name}"].{channel}'

    for i in range(0,3):
        f = act.fcurves.find(data_path, index=i)
        if not f:
            if verbose:
                print(f'{item.name} has not {data_path}')
            continue
        
        if f.is_empty:
            if verbose:
                print(f'fcurve has not keyframes: {f.data_path} {i}')
            continue
        
        if not f.is_valid:
            if verbose:
                print(f'fcurve is invalid {f.data_path} {i}')
            continue
        
        ## ? int frame ?
        if next((k for k in f.keyframe_points if k.co.x == frame), None) is not None:
            if verbose:
                print(f'{item.name} {channel} is keyframed')
            return True
    return False

"""
def has_channel_key_at_frame(pb, act, frame, channel='location'):
    '''Return Rrue if pose bone has a key at passed frame'''
    for i in range(0,3):
        f = act.fcurves.find(f'pose.bones["{pb.name}"].{channel}', index=i)
        if not f:
            continue
        if f.is_empty:
            print(f'fcurve has not keyframes: {f.data_path} {i}')
            continue
        if not f.is_valid:
            print(f'fcurve is invalid {f.data_path} {i}')
            continue
        
        ## ? int frame ?
        if next((k for k in f.keyframe_points if k.co.x == frame), None) is not None:
            return True
    return False
"""

def assign_world_location_to_pose_bone(pb, newloc):
    # mat = scn.cursor.matrix
    
    # get world space matrix
    mat = pb.id_data.matrix_world @ pb.matrix
    loc, rot, _scale = mat.decompose()

    # compose
    # loc_mat = Matrix.Translation(loc)
    loc_mat = Matrix.Translation(newloc)
    rot_mat = rot.to_matrix().to_4x4()
    scale_mat = scale_matrix_from_vector(pb.scale)
    new_mat = loc_mat @ rot_mat @ scale_mat

    # assign
    pb.matrix = pb.id_data.matrix_world.inverted() @ new_mat


def sort_bones_by_hierarchy_depth(bones):
    hierarchy_depth = {}
    
    def get_depth(bone):
        if bone.parent is None:
            return 0
        if bone.name in hierarchy_depth:
            return hierarchy_depth[bone.name]
        else:
            depth = get_depth(bone.parent) + 1
            hierarchy_depth[bone.name] = depth
            return depth
    
    for bone in bones:
        get_depth(bone)
    
    # print('hierarchy_depth: ', hierarchy_depth) # Dbg
    # Ex: hierarchy_depth -> {'world': 1, 'walk': 2, 'scale-all': 3, 'rotate-hip': 4, 'spine01': 5, 'spine02':6, 'foot.R': 1, 'foot.L': 1}
    sorted_bones = sorted(bones, key=lambda b: hierarchy_depth[b.name])
    return sorted_bones

def wrap_animation(ref, root=None, verbose=False):
    scn = bpy.context.scene
    org_frame = scn.frame_current
    
    ob = bpy.context.object
    ac = ob.animation_data.action

    ## Auto set root with classic names
    if not root:
        root = ob.pose.bones.get('root')
    if not root:
        root = ob.pose.bones.get('world')
    # if not root:
    #     root = ob.pose.bones.get('walk')
    
    target_bones = bpy.context.selected_pose_bones

    ## During the loop, move upper hiercharchy bones first
    ## (Not 100% sure that's needed, as it's usually not targetting dependent bones)
    target_bones = sort_bones_by_hierarchy_depth(target_bones)
    
    for i in range(len(target_bones))[::-1]:
        ## Remove the root:
        if target_bones[i] == root:
            target_bones.pop(i)
        
        ## Always put the ref bone at the end:
        if target_bones[i] == ref:
            target_bones.append(target_bones.pop(i))
    
    print('Targeted bones', [b.name for b in target_bones]) #Dbg
    
    '''
    ### Iterate in whole summary:
    ### -> Frame -> Bones
    frames = dopesheet_summary(ob)

    ## remove duplicates
    if verbose:
        print('-Keyframes-')
        print(frames)

    ### Note: /!\ Still some weird error, work when select and place bones one by one !
    ### Not optimized but need to iterate in bone first then frame for each...
    for f in reversed(frames):
        frame = int(f)
        if not verbose:
            print(f'{frame:04d}', end='\r')
        scn.frame_set(frame)
        
        ## calc distance
        ref_loc = (ref.id_data.matrix_world @ ref.matrix).translation.copy()    
        ref_loc =  Vector((ref_loc.x, ref_loc.y, 0)) # Remove Z
        
        ## On a moving root, recalculate root_loc on each iteration
        root_loc = (root.id_data.matrix_world @ root.matrix).translation.copy()
        root_loc = Vector((root_loc.x, root_loc.y, 0)) # Remove Z
        
        # get reset_vector
        reset_vec = root_loc - ref_loc
        
        for pb in sorted_selected_pose_bones:
            # if not has_channel_key_at_frame(pb, ac, frame, channel='location'):
            if not has_key_at_frame(pb, ac, frame, channel='location', verbose=verbose):
                if verbose:
                    print(f'{frame:03d}: {pb.name} no location keys')
                continue

            #### assign the matrix from world
            ## bone world position
            bone_world_loc = (pb.id_data.matrix_world @ pb.matrix).translation.copy()
            new_pos = bone_world_loc + reset_vec
            # scn.cursor.location = new_pos #Dbg
            
            assign_world_location_to_pose_bone(pb, new_pos)
                        
            ## create the keyframe
            pb.keyframe_insert('location', frame=frame)
    '''

    ### Iterate in individual bone summary:
    ### -> bone -> frames
    
    for pb in target_bones:
        # Calculate single bone frame summary
        frames = pose_bone_frame_summary(pb, filter_channel='location')
        print(f'{pb.name} : {len(frames)} location frames')

        for f in reversed(frames):
            frame = int(f)
            if not verbose:
                print(f'{pb.name} : frame {frame:04d}', end='\r')
            scn.frame_set(frame)

            ## calc distance
            ref_loc = (ref.id_data.matrix_world @ ref.matrix).translation.copy()    
            ref_loc =  Vector((ref_loc.x, ref_loc.y, 0)) # Remove Z

            ## On a moving root, recalculate root_loc on each iteration
            root_loc = (root.id_data.matrix_world @ root.matrix).translation.copy()
            root_loc = Vector((root_loc.x, root_loc.y, 0)) # Remove Z

            # get reset_vector
            reset_vec = root_loc - ref_loc

            if not has_key_at_frame(pb, ac, frame, channel='location', verbose=verbose):
                if verbose:
                    print(f'{frame:03d}: {pb.name} no location keys')
                continue

            #### assign the matrix from world
            ## bone world position
            bone_world_loc = (pb.id_data.matrix_world @ pb.matrix).translation.copy()
            new_pos = bone_world_loc + reset_vec

            assign_world_location_to_pose_bone(pb, new_pos)

            ## create the keyframe
            pb.keyframe_insert('location', frame=frame)

    if not verbose:
        print()
    scn.frame_current = org_frame


class AW_OT_wrap_animation(bpy.types.Operator):
    bl_idname = "autowalk.wrap_animation"
    bl_label = "Wrap Animation"
    bl_description = "Wrap the current animation on selected bones.\
        \nApply Reference -> root offset to selected bones on each frames"
    bl_options = {"REGISTER", "UNDO"}

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

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

        settings = context.scene.anim_cycle_settings
        if not settings.wrap_ref_bone or not settings.wrap_root_bone:
            self.report({'ERROR'}, 'Need to specify both reference and root bone fields')
            return {"CANCELLED"}

        if not ob.animation_data or not ob.animation_data.action:
            self.report({'ERROR'}, f'No animation data on {ob.name}')
            return {"CANCELLED"}
        
        ref = settings.wrap_ref_bone
        root = settings.wrap_root_bone

        ## as strings
        ref = ob.pose.bones.get(ref)
        root = ob.pose.bones.get(root)

        if not ref:
            self.report({'ERROR'}, f'Reference bone not found in object: {context.object.name}')
            return {"CANCELLED"}
        if not root:
            self.report({'ERROR'}, f'Root bone not found in object: {context.object.name}')
            return {"CANCELLED"}

        wrap_animation(ref, root)
        return {"FINISHED"}
    
classes=(
AW_OT_wrap_animation,
)

def register():
    # bpy.types.Object.pose_bone = bpy.props.PointerProperty(type=bpy.types.PoseBone)
    
    for cls in classes:
        bpy.utils.register_class(cls)

def unregister():
    for cls in reversed(classes):
        bpy.utils.unregister_class(cls)
    
    # del bpy.types.Object.pose_bone