# copy_anim.py
"""
Blender Python code to copy animation between armatures or proxy armatures.

The purpose of the 'Copy Animation' feature is to allow for animation to be
copied from one armature to another, en masse, rather than having to individual
push and move action objects.

The main use for this is to repair files in which animated proxy rigs have
become incompatible or broken for some reason. Common examples include a name
change in the rig or armature object in a character asset file, extra bones
added, and so on. There is no simple way in Blender to update these proxies.

It is possible to create a new proxy, though, and with this tool to speed up
the process, the animation can be transferred to it all at once.

The tool also allows for the animation to be correctly copied and scaled by
a scale factor, so that animation can be copied from a proxy defined at one
scale to one defined at another.

This comes up when an animation file was built incorrectly at the wrong scale
and needs to be corrected, after animating has already begun.

The scaling feature has been tested on Rigify-based rigs, and resets the
bone constraints as needed, during the process.
"""

import bpy, bpy.types, bpy.utils, bpy.props

#----------------------------------------
## TOOLS
# This might be moved into another module later

def copy_object_animation(sourceObj, targetObjs,
        dopesheet=False, nla=False, rescale=False, scale_factor=1.0,
        report=print):
    """
    Copy Dope Sheet & NLA editor animation from active object to selected objects.
    Most useful with armatures. Assumes bones match. Can be rescaled in the process.
    
    From StackExchange post:
    https://blender.stackexchange.com/questions/74183/how-can-i-copy-nla-tracks-from-one-armature-to-another
    """
    for targetObj in targetObjs:
        if targetObj.animation_data is not None:
            targetObj.animation_data_clear()

        targetObj.animation_data_create()   

        source_animation_data = sourceObj.animation_data
        target_animation_data = targetObj.animation_data
        
        # copy the dopesheet animation (active animation)
        if dopesheet:
            report({'INFO'}, 'Copying Dopesheet animation')
            if source_animation_data.action is None:
                report({'WARNING'}, 
                    "CLEARING target dope sheet - old animation saved with 'fake user'")
                if target_animation_data.action is not None:
                    target_animation_data.action.use_fake_user = True
                target_animation_data.action = None
            else:
                if rescale:
                    target_animation_data.action = copy_animation_action_with_rescale(
                        source_animation_data.action, scale_factor)
                else:
                    target_animation_data.action = copy_animation_action_with_rescale(
                        source_animation_data.action, scale_factor)
                
                target_animation_data.action.name = targetObj.name + 'Action'
        
        if nla:
            report({'INFO'}, 'Copying NLA strips')
            if source_animation_data:
                # Create new NLA tracks based on the source
                for source_nla_track in source_animation_data.nla_tracks:
                    target_nla_track = target_animation_data.nla_tracks.new()
                    target_nla_track.name = source_nla_track.name
                    # In each track, create action strips base on the source
                    for source_action_strip in source_nla_track.strips:
                        
                        if rescale:
                            new_action = copy_animation_action_with_rescale(
                                    source_action_strip.action, scale_factor)
                        else:
                            new_action = source_action_strip.action
                        
                        target_action_strip = target_nla_track.strips.new(
                            new_action.name,
                            source_action_strip.frame_start,
                            new_action)   
                        
                        # For each strip, copy the properties -- EXCEPT the ones we
                        # need to protect or can't copy
                        # introspect property names (is there a better way to do this?)
                        props = [p for p in dir(source_action_strip) if
                                    not p in ('action',)
                                    and not p.startswith('__') and not p.startswith('bl_')
                                    and source_action_strip.is_property_set(p)
                                    and not source_action_strip.is_property_readonly(p)
                                    and not source_action_strip.is_property_hidden(p)]
                        for prop in props:
                            setattr(target_action_strip, prop, getattr(source_action_strip, prop))
                        

# Adapted from reference:
# https://www.reddit.com/r/blender/comments/eu3w6m/guide_how_to_scale_a_rigify_rig/
#

def reset_armature_stretch_constraints(rig_object):
    """
    Reset stretch-to constraints on an armature object - necessary after rescaling.
    """
    bone_count = 0
    for bone in rig_object.pose.bones:
        for constraint in bone.constraints:
            if constraint.type == "STRETCH_TO":
                constraint.rest_length = 0
                bone_count += 1
    return bone_count


def rescale_animation_action_in_place(action, scale_factor):
    """
    Rescale a list of animation actions by a scale factor (in-place).
    """
    #for fcurve in bpy.data.actions[action].fcurves:
    for fcurve in action.fcurves:
        data_path = fcurve.data_path
        if data_path.startswith('pose.bones[') and data_path.endswith('].location'):
            for p in fcurve.keyframe_points:
                p.co[1] *= scale_factor
                p.handle_left[1] *= scale_factor
                p.handle_right[1] *= scale_factor
    return action

def copy_animation_action_with_rescale(action, scale_factor):
    """
    Copy an animation action, rescaled.
    """
    new_action = action.copy()
    new_action.name = new_action.name[:-4]+'.rescale'
    return rescale_animation_action_in_place(new_action, scale_factor)


    

#----------------------------------------