bl_info = {
    "name": "PSA additite animation restore",
    "author": "Noterstone",
    "version": (0, 0, 1),
    "blender": (3, 0, 0),
    "description": "PSA additite animation restore Unreal Engine",
    "category": "Import-Export",
}

import bpy
import math
import mathutils
from bpy.types import Operator, Panel, PropertyGroup, PoseBone
from bpy.utils import register_class, unregister_class
from bpy.props import StringProperty, EnumProperty
from mathutils import Vector, Matrix, Quaternion


def update_source_armature(self, context):   
    scn = context.scene
    if scn.source_armature != "":
        src_armature = get_object(scn.source_armature)
        if scn.target_armature != "":
            set_global_scale(context)


def update_target_armature(self,context):
    scn = context.scene
    if scn.source_armature != "" and scn.target_armature != "":
        set_global_scale(context)


class OBJECT_OT_pick_object(Operator):
    bl_idname = "object.pick_object"
    bl_label = "pick_object"
    bl_options = {'UNDO'}

    action : EnumProperty(
        items=(
                ('pick_source_armature', 'pick_source_armature', ''),
                ('pick_target_armature', 'pick_target_armature', '')
            )
        )

    @classmethod
    def poll(cls, context):
        return (context.active_object != None)

    def execute(self, context):
        use_global_undo = context.preferences.edit.use_global_undo
        context.preferences.edit.use_global_undo = False

        try:
            _pick_object(self.action)

        finally:
            context.preferences.edit.use_global_undo = use_global_undo

        return {'FINISHED'}

def _pick_object(action):
    obj = bpy.context.object
    scene = bpy.context.scene

    if action == "pick_source_armature":
        scene.source_armature = obj.name
    if action == "pick_target_armature":
        scene.target_armature = obj.name

class OBJECT_OT_restore_additive_anim(Operator):
    bl_idname = "object.restore_additive_anim"
    bl_label = ""
    bl_options = {'UNDO'}

    def draw(self, context):
        layout = self.layout
        scn = context.scene
                
        row = layout.row(align=True)
        
    def execute(self, context):
        sce = bpy.context.scene
        source_armature = bpy.data.objects.get(context.scene.source_armature)
        target_armature = bpy.data.objects.get(context.scene.target_armature)

        for f in range(sce.frame_start, sce.frame_end+1):
            sce.frame_set(f)
            for target_pose_bone in target_armature.pose.bones:
                source_pose_bone = source_armature.pose.bones[target_pose_bone.name]
                
                # is not correct, is calculations in POSE space
                # must be calculated in local space for each of the bones
                target_pose_bone.location += source_pose_bone.location
                target_pose_bone.rotation_quaternion *= source_pose_bone.rotation_quaternion
                target_pose_bone.scale *= source_pose_bone.scale
                target_pose_bone.keyframe_insert("location", frame=f)
                target_pose_bone.keyframe_insert("rotation_quaternion", frame=f)
                target_pose_bone.keyframe_insert("scale", frame=f)

        sce.frame_set(0)
        return {'FINISHED'}

class OBJECT_PT_AdditiveAnimUERestore(Panel):
    bl_label = "Additive anim UE restore"
    bl_space_type = "VIEW_3D"
    bl_region_type = "UI"
    bl_category = "AdditiveAnimUERestore"

    def draw(self, context):
        layout = self.layout
        object = context.object
        scn = context.scene
        
        layout.label(text="Base pose animated armature:")
        row = layout.row(align=True)
        row.prop_search(scn, "source_armature", bpy.data, "objects", text="")
        row.operator("object.pick_object", text="", icon='EYEDROPPER').action = 'pick_source_armature'

        layout.label(text="Armature with additive anim:")
        row = layout.row(align=True)
        row.prop_search(scn, "target_armature", bpy.data, "objects", text="")
        row.operator("object.pick_object", text="", icon='EYEDROPPER').action = 'pick_target_armature'

        row = layout.row(align=True)
        col = layout.column(align=True)
        col.operator("object.restore_additive_anim", text="Restore anim", icon="PLAY")
        layout.separator()


classes = [
    OBJECT_OT_restore_additive_anim,
    OBJECT_OT_pick_object,
    OBJECT_PT_AdditiveAnimUERestore
]

def register():
    for cl in classes:
        register_class(cl)

    bpy.types.Scene.target_armature = StringProperty(name = "Target Armature anim", default="", description="Destination armature", update=update_target_armature)
    bpy.types.Scene.source_armature = StringProperty(name = "Source Armature anim", default="", description="Source armature", update=update_source_armature)

def unregister():
    for cl in classes:
        unregister_class(cl)

    del bpy.types.Scene.target_armature
    del bpy.types.Scene.source_armature

if __name__ == '__main__':
    register()