from bpy import data

INTERPOLATION = "LINEAR"

location_kf_counter = 0
rotation_kf_counter = 0
rotation_kf_inc_counter = 0
rotation_kf_fix_counter = 0
rotation_kf_nofix_counter = 0

for action in data.actions:
    print("AGR_INTERP: Processing action {}".format(action.name))
    rot_fcurve_lists = {}
    for fcurve in action.fcurves:
        fc_split = fcurve.data_path.rsplit(".", 1)
        fc_name = (action.name + "." + fc_split[0]) if len(fc_split) == 2 else action.name
        fc_type = fc_split[-1]
        # Interpolate location normally
        if fc_type == "location":
            for kf in fcurve.keyframe_points:
                if kf.interpolation != INTERPOLATION:
                    kf.interpolation = INTERPOLATION
                    location_kf_counter += 1
        # Collect all 4 quaternion components to fix them, group them by their "names" extracted from data_path
        elif fc_type == "rotation_quaternion":
            if not fc_name in rot_fcurve_lists:
                rot_fcurve_lists[fc_name] = []
            rot_fcurve_lists[fc_name].append(fcurve)
    for rot_fcurves_name, rot_fcurves in rot_fcurve_lists.items():
        print("AGR_INTERP: Processing rotation for {}".format(rot_fcurves_name))
        rot_fcurves_n = len(rot_fcurves)
        if rot_fcurves_n > 0:
            if rot_fcurves_n == 4:
                # Ensure rotation interpolates correctly
                # Loop through each keyframe of the first component
                other_fcurves = rot_fcurves[1:]
                kf_is = [0 for fcurve in other_fcurves] # is = indexes
                prev_components = [None for fcurve in rot_fcurves]
                incomplete_kfs = set()
                for comp0 in rot_fcurves[0].keyframe_points:
                    if comp0.interpolation != INTERPOLATION:
                        comp0.interpolation = INTERPOLATION
                        rotation_kf_counter += 1
                    # Ensure other components exist at this frame (referred to as frame is "complete" from now on)
                    comp0_x = comp0.co[0]
                    components = [comp0]
                    incomplete = False
                    # Loop through the rest of the components
                    for i, fcurve in enumerate(other_fcurves):
                        kf_i = kf_is[i]
                        kf_i_max = len(fcurve.keyframe_points) - 1
                        if kf_i > kf_i_max:
                            continue
                        kf = fcurve.keyframe_points[kf_i]
                        if kf.interpolation != INTERPOLATION:
                            kf.interpolation = INTERPOLATION
                            rotation_kf_counter += 1
                        kf_x = kf.co[0]
                        # Initially populate component with the previous frame's value
                        # to keep track of most recent values even when there are incomplete frames
                        components.append(prev_components[len(components)-1])
                        # If the component's current keyframe's frame is less than the current frame
                        # go forward the keyframes to see if we can find a matching one
                        while kf_i <= kf_i_max and kf_x < comp0_x:
                            # If we need to skip a keyframe, it must be incomplete
                            incomplete_kfs.add(str(kf_x))
                            components[-1] = kf
                            if kf.interpolation != INTERPOLATION:
                                kf.interpolation = INTERPOLATION
                                rotation_kf_counter += 1
                            kf_i += 1
                            kf = fcurve.keyframe_points[kf_i]
                            kf_x = kf.co[0]
                        if kf_x == comp0_x:
                            # Found a valid keyframe of the component at this frame
                            components[-1] = kf
                            # Move index by one since we found a frame for the current index's keyframe
                            kf_i += 1
                        else:
                            # kf_x > comp_x or reached the end of fcurve kf_points
                            # The current keyframe is incomplete, try next one
                            # Index is not moved, since this keyframe can be part of another complete frame
                            incomplete = True
                            incomplete_kfs.add(str(comp0_x))
                        # Update index to avoid checking keyframes more than once
                        kf_is[i] = kf_i
                    if prev_components.count(None) == 0:
                        # Compare to previous frame, determine if negated is closer
                        curr_y_values = [component.co[1] for component in components]
                        neg_y_values = [-component.co[1] for component in components]
                        last_y_values = [component.co[1] for component in prev_components]
                        curr_diff = sum([abs(last_y_values[i] - curr_y_values[i]) for i in range(len(curr_y_values))])
                        neg_diff = sum([abs(last_y_values[i] - neg_y_values[i]) for i in range(len(neg_y_values))])
                        if neg_diff < curr_diff:
                            if incomplete:
                                print("AGR_INTERP: Cannot negate keyframe at {}: quaternion is incomplete".format(components[0].co.x))
                                rotation_kf_nofix_counter += 1
                            else:
                                # VVV Too much spam VVV
                                # print("AGR_INTERP: Negating keyframe at {}".format(components[0].co.x))
                                for i, component in enumerate(components):
                                    component.co[1] = neg_y_values[i]
                                rotation_kf_fix_counter += 1
                    # Store components of the frame for future comparisons
                    prev_components = components
                if len(incomplete_kfs) > 0:
                    print("AGR_INTERP: Skipped incomplete keyframes at {}".format(", ".join(incomplete_kfs)))
                    rotation_kf_inc_counter += len(incomplete_kfs)
            else:
                print("AGR_INTERP: Invalid amount of rotation quaternion components: {}".format(rot_fcurves_n))

print("AGR_INTERP: Interpolated {} location keyframes".format(location_kf_counter))
print("AGR_INTERP: Interpolated {} rotation keyframes".format(rotation_kf_counter))
print("AGR_INTERP: Found {} incomplete quaternions".format(rotation_kf_inc_counter))
print("AGR_INTERP: Negated {} rotation keyframes".format(rotation_kf_fix_counter))
print("AGR_INTERP: Could not negate {} rotation keyframes due to incomplete quaternions".format(rotation_kf_nofix_counter))