from pathlib import Path
import bpy
 
foldertoimport = r"  "
objroot = None
def RestoreMissingWeight(obj):
    if obj.type == "MESH":
 
        #restore the deleted vertex group
        obj.vertex_groups.new(name = objroot)
        #if obj.vertex_groups[objroot] == None:
 
        #for every vertex group in the mesh lock it, except the root         
        for a in obj.vertex_groups:
            if a != obj.vertex_groups[objroot]:
                a.lock_weight = True
 
        #reference for the root vertex group
        group = obj.vertex_groups[objroot]
 
        #for every vertex in the mesh add a weight of 1 in the root vertex group
        for vert in obj.data.vertices:
 
            group.add( [vert.index], 1, 'ADD' )
        #normalize all the mesh vertex groups so every one have a value and no vertex is left with 0 weight
        #basicly it restore missing the weight data for the root bone
        bpy.ops.object.vertex_group_normalize_all(lock_active = False)
 
        #unlock all vertex groups
        for a in obj.vertex_groups:
            a.lock_weight = False
 
 
def createnewfcurves(rootname = objroot ,rootdata_path = ''):
 
    ConvertObjAnimDataToBone(rootname = objroot ,data_path = rootdata_path , rootindex = 0)
    ConvertObjAnimDataToBone(rootname = objroot ,data_path = rootdata_path , rootindex = 1)
    ConvertObjAnimDataToBone(rootname = objroot ,data_path = rootdata_path , rootindex = 2)
    if rootdata_path == 'rotation_quaternion':
 
        ConvertObjAnimDataToBone(rootname = objroot ,data_path = rootdata_path , rootindex = 3)
 
def ConvertObjAnimDataToBone(rootname = objroot,data_path = '' , rootindex = 0):
 
    obj = bpy.context.object
    if obj:
 
        action = obj.animation_data.action
        Root_bone = obj.pose.bones.get(rootname)
        if Root_bone != None: 
 
            Root_bone.keyframe_insert(data_path, index = rootindex ,frame = 1,)#group = 'root', options = {'INSERTKEY_NEEDED' })
            fcurve = action.fcurves.find(data_path,index = rootindex )
            bfcurve = action.fcurves.find('pose.bones["{}"].{}'.format(rootname,data_path) , index = rootindex )
 
            for ky in fcurve.keyframe_points:
 
                bfcurve.keyframe_points.insert(frame = ky.co.x , value = ky.co.y)
 
 
 
 
def ResetActionTransforms(action_name):
    action = bpy.data.actions.get()
    if action_name == None:
        return
    action_name = bpy.context.active_object.animation_data.action.name
    data_path = 'location'
    if action:
        # From this action, retrieve the appropriate F-Curve
                                        #index of the desired Coordinates 
                                        # in this case it means the x location
        fcurve = action.fcurves.find(data_path, index = 0 )
        if fcurve:
            #save the value of the first frame of the object action for later use
            kframe = fcurve.keyframe_points[0].co.y
            # Iterate over all keyframes
            for a in fcurve.keyframe_points:
                a.co.y -= kframe
                                            #index of the desired coordinates 
                                        # in this case it means the y location
        fcurve = action.fcurves.find(data_path, index = 1 )
        if fcurve:
            #save the value of the first frame of the object action for later use
            kframe = fcurve.keyframe_points[0].co.y
            # Iterate over all keyframes
            for a in fcurve.keyframe_points:
                a.co.y -= kframe
                                        #index of the desired coordinates 
                                        # in this case it means the z location
        fcurve = action.fcurves.find(data_path, index = 2 )
        if fcurve:
            #save the value of the first frame of the object action for later use
            kframe = fcurve.keyframe_points[0].co.y
            # Iterate over all keyframes
            for a in fcurve.keyframe_points:
 
                #[1] is the index of the y coordinate in f-curve terms
                #it's basicly the same thing as using ".y"
                #leaving it here just to remember it's possible
 
                a.co[1] -= kframe
                #a.co[1] -= bpy.context.active_object.location.z
def get_or_create_fcurve(action, data_path, array_index=-1, group=None):
    for fc in action.fcurves:
        if fc.data_path == data_path and (array_index < 0 or fc.array_index == array_index):
            return fc
 
    fc = action.fcurves.new(data_path, index=array_index)
    fc.group = group
    return fc 
def add_keyframe_quat(action, quat, frame, bone_prefix, group):
    for i in range(len(quat)):
        fc = get_or_create_fcurve(action, bone_prefix + "rotation_quaternion", i, group)
        pos = len(fc.keyframe_points)
        fc.keyframe_points.add(1)
        fc.keyframe_points[pos].co = [frame, quat[i]]
        fc.update()   
def frames_matching(action, data_path):
    frames = set()
    for fc in action.fcurves:
        if fc.data_path == data_path:
            fri = [kp.co[0] for kp in fc.keyframe_points]
            frames.update(fri)
    return frames
def group_eq(_obj, action, bone, bone_prefix, order):
    """Converts only one group/bone in one action - Euler to Quaternion."""
    # pose_bone = bone
    data_path = bone_prefix + "rotation_euler"
    frames = frames_matching(action, data_path)
    group = action.groups[bone.name]
 
    for fr in frames:
        euler = bone.rotation_euler.copy()
        for fc in action.fcurves:
            if fc.data_path == data_path:
                euler[fc.array_index] = fc.evaluate(fr)
        quat = euler.to_quaternion()
 
        add_keyframe_quat(action, quat, fr, bone_prefix, group)
        bone.rotation_mode = order
 
 
def removefcurves(data_path = ''):
 
    obj = bpy.context.object
    action = obj.animation_data.action
    if action:
        fcurve = action.fcurves.find(data_path, index = 0 )
        if fcurve:
                obj.animation_data.action.fcurves.remove(fcurve)
        fcurve = action.fcurves.find(data_path, index = 1 )
        if fcurve:
                obj.animation_data.action.fcurves.remove(fcurve)
        fcurve = action.fcurves.find(data_path, index = 2 )
        if fcurve:
                obj.animation_data.action.fcurves.remove(fcurve)
        if data_path == 'rotation_quaternion':
 
            fcurve = action.fcurves.find(data_path, index = 3 )
            if fcurve:
                obj.animation_data.action.fcurves.remove(fcurve)
 
 
 
 
folder = Path(foldertoimport)
 
fbx_files = [f for f in folder.glob("**/*.fbx") if f.is_file()]
for fbx_file in fbx_files: #i
 
    #bpy.ops.object.select_all(action='DESELECT')
    #bpy.ops.object.select_by_type(type='ARMATURE')
    bpy.ops.import_scene.fbx(filepath=str(fbx_file))
 
 
    object_to_delete = bpy.data.objects['BaseNode']
    bpy.data.objects.remove(object_to_delete, do_unlink=True)
 
    for obj in bpy.context.selected_objects:
 
 
 
        if obj.type == "ARMATURE":
            objroot = obj.name
            print(objroot)
            #rename the object name, action and armature properties
            obj.name = fbx_file.stem
            obj.animation_data.action.name = fbx_file.stem
            obj.data.name = fbx_file.stem
 
            #set scene to edit mode
            bpy.ops.object.mode_set(mode='EDIT', toggle=False)
            obArm = bpy.context.active_object #get the armature object
            ebs = obArm.data.edit_bones
            eb = ebs.new(objroot)
            eb.head = (0, 0, 0) # if the head and tail are the same, the bone is deleted
            eb.tail = (0, 7.02337 , 0)
 
            #make the root bone the parent to all bones that don't have a parent
            for bones in obArm.data.bones: 
                if bones.parent == None:
                    bonename = bones.name
                    ebs.data.edit_bones[bonename].parent = ebs.data.edit_bones[objroot]
 
            #set scene mode pose mode
            bpy.ops.object.mode_set(mode='POSE', toggle=False)
 
            #bpy.ops.object.select_all(action='DESELECT')
            Root_bone = obj.pose.bones.get(objroot)
 
            #convert armature animation from euler to quaternion
            group_eq(obArm,obj.animation_data.action, Root_bone,'','QUATERNION')
            if Root_bone != None:
                bpy.data.objects[obj.name].data.bones[objroot].select = True
 
                #move armature animation to the bone
                createnewfcurves(rootdata_path= 'location')
                createnewfcurves(rootdata_path= 'rotation_quaternion')
                #delete unecessary animations from the armature
                removefcurves('rotation_euler')
                removefcurves('scale')
                removefcurves('location')
                removefcurves('rotation_quaternion')
                #reset the armature location and rotation
                obj.location = 0, 0, 0
                obj.rotation_euler = 0,0,0
 
 
 
 
 
 
 
            bpy.ops.object.mode_set(mode='OBJECT', toggle=False)
        view_layer = bpy.context.view_layer
 
 
        if obj.type == 'ARMATURE':
            obj.select_set (False)
 
        if obj.type == "MESH":
            view_layer.objects.active = obj
            RestoreMissingWeight(obj)