animation/tools/loop_fix.py

"""loop_fix.py — transform a mocap GLB so it loops fluidly under Godot LoopMode.Linear.

The game wraps each clip straight from its last keyframe back to frame 0. Raw mocap
takes don't end where they began, so the wrap pops. This tool rewrites a clip so that,
at the seam, both POSE and MOTION (and root position) are continuous, which is exactly
what the sibling gate tools/loop_qc.py measures (read both loop_qc.py and
cc_retarget.py before editing this).

Three stages, applied per clip:

  A. de-drift (root)   — subtract a linear horizontal ramp from the pelvis location so
                         the clip ends where it starts. Vertical untouched (crouches are
                         choreography). Brings `root drift` under the 3 cm limit.
  B. pick loop cut     — search the last ~40% of the clip for the frame F whose core-bone
                         pose+velocity best matches frame 0. Trim the action to [W, F].
  C. seam blend        — crossfade the final W frames of the kept region into the first W
                         frames (motion-graph seam blend, quaternion slerp w/ hemisphere
                         correction, smoothstep weight). Guarantees pose + velocity
                         continuity when the game wraps the last frame back to the first.

Run headless (same pattern as the other tools in tools/):
  blender --background --python tools/loop_fix.py -- --src <in.glb> --out <out.glb> \
          [--name <Clip>] [--blend-frames 24] [--min-keep 0.6] \
          [--no-trim] [--no-blend] [--no-dedrift]

Exit 0 = ok, 2 = error. [loop_fix]-prefixed prints are greppable through Blender noise.
"""
import bpy, sys, math, argparse
from mathutils import Quaternion

# Bones used to SCORE the loop cut (Stage B) — the structurally important chain.
# Fingers are cosmetic and below the visual noise floor, so they are ignored while
# scoring (but still blended in Stage C — blending is cheap).
CORE_BONES = (
    "pelvis", "spine_01", "spine_02", "spine_03", "neck_01", "Head",
    "clavicle_l", "upperarm_l", "lowerarm_l", "hand_l",
    "clavicle_r", "upperarm_r", "lowerarm_r", "hand_r",
    "thigh_l", "calf_l", "foot_l",
    "thigh_r", "calf_r", "foot_r",
)


def parse():
    argv = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
    p = argparse.ArgumentParser()
    p.add_argument("--src", required=True)
    p.add_argument("--out", required=True)
    p.add_argument("--name", default=None, help="output clip name (default: source action name)")
    p.add_argument("--blend-frames", type=int, default=24, help="crossfade window W")
    p.add_argument("--min-keep", type=float, default=0.6, help="minimum fraction of clip kept")
    p.add_argument("--no-trim", action="store_true", help="skip Stage B (use last frame as cut)")
    p.add_argument("--no-blend", action="store_true", help="skip Stage C (no seam crossfade)")
    p.add_argument("--no-dedrift", action="store_true", help="skip Stage A (no root de-drift)")
    return p.parse_args(argv)


def smoothstep(t):
    return 3.0 * t * t - 2.0 * t * t * t


def ang(q1, q2):
    """Shortest angular distance (deg) between two quaternions — same as loop_qc.ang."""
    d = abs(max(-1.0, min(1.0, q1.dot(q2))))
    return math.degrees(2.0 * math.acos(d))


def slerp(q1, q2, t):
    """Hemisphere-corrected spherical lerp of two quaternions. Pure-math (no API guessing)."""
    a = list(q1)
    b = list(q2)
    dot = sum(a[i] * b[i] for i in range(4))
    if dot < 0.0:           # hemisphere correction — skip this and short arcs become long spins
        b = [-x for x in b]
        dot = -dot
    if dot > 0.9995:        # near-parallel → linear interpolation is stable, slerp divides by ~0
        r = [a[i] + (b[i] - a[i]) * t for i in range(4)]
    else:
        th0 = math.acos(min(1.0, max(-1.0, dot)))
        s0 = math.sin(th0)
        th = th0 * t
        s1 = math.cos(th) - dot * math.sin(th) / s0
        s2 = math.sin(th) / s0
        r = [a[i] * s1 + b[i] * s2 for i in range(4)]
    n = math.sqrt(sum(x * x for x in r)) or 1.0
    return Quaternion([x / n for x in r])


def lerp(v1, v2, t):
    return tuple(v1[i] + (v2[i] - v1[i]) * t for i in range(3))


def load(src):
    bpy.ops.wm.read_factory_settings(use_empty=True)
    bpy.ops.import_scene.gltf(filepath=src)
    arm = next(o for o in bpy.data.objects if o.type == "ARMATURE")
    act = arm.animation_data.action
    # strip any meshes (these clips are meshless, but match cc_retarget defensively)
    for o in [o for o in bpy.data.objects if o.type == "MESH"]:
        bpy.data.objects.remove(o, do_unlink=True)
    for pb in arm.pose.bones:
        pb.rotation_mode = "QUATERNION"
    bpy.context.view_layer.update()
    return arm, act


def sample_frames(arm, act):
    """Read per-frame pose (quaternion for every bone, pelvis location) into memory.

    Reads matrix_basis the same way loop_qc.local_quats does, so what we measure here is
    exactly what the gate measures. Blender 5.1 stores keys in slotted channelbags, so we
    avoid touching fcurves directly and just evaluate the posed rig frame by frame.
    """
    f0, f1 = int(act.frame_range[0]), int(act.frame_range[1])
    scene = bpy.context.scene
    dg = bpy.context.evaluated_depsgraph_get()
    bones = [pb.name for pb in arm.pose.bones]
    root = "pelvis" if "pelvis" in bones else bones[0]

    rot = {f: {} for f in range(f0, f1 + 1)}        # frame -> bone -> Quaternion
    loc = {f: None for f in range(f0, f1 + 1)}      # frame -> pelvis (x,y,z)
    for f in range(f0, f1 + 1):
        scene.frame_set(f)
        dg.update()
        for pb in arm.pose.bones:
            rot[f][pb.name] = pb.matrix_basis.to_quaternion().normalized()
        loc[f] = tuple(arm.pose.bones[root].matrix_basis.translation)
    return f0, f1, root, bones, rot, loc


def stage_dedrift(f0, f1, loc, enabled):
    """Stage A — remove net horizontal pelvis travel via a linear ramp on the floor plane.

    Verified empirically on these exports: the armature object transform is identity and the
    rig is Z-up (pelvis head sits at Z≈0.5 m, X≈Y≈0 at rest), so horizontal = {X, Y} and
    vertical = Z. We ramp only the floor axes so the clip ends where it started on the ground;
    vertical (crouch) is choreography and is left alone. (Stage C's blend then forces the seam
    location itself to match exactly, so this is about keeping the crossfade region close.)
    """
    if not enabled:
        print("[loop_fix] A de-drift: SKIPPED (--no-dedrift)")
        return 0.0
    horiz = (0, 1)  # X, Y = floor plane; Z = up
    D = [loc[f1][i] - loc[f0][i] for i in range(3)]
    span = float(f1 - f0) or 1.0
    for f in range(f0, f1 + 1):
        frac = (f - f0) / span
        nl = list(loc[f])
        for i in horiz:
            nl[i] -= D[i] * frac
        loc[f] = tuple(nl)
    rem_cm = math.sqrt(sum(D[i] ** 2 for i in horiz)) * 100.0
    z_cm = abs(D[2]) * 100.0
    print(f"[loop_fix] A de-drift: removed {rem_cm:.1f} cm over {int(span)+1}f "
          f"on floor axes {horiz}; kept vertical Z ({z_cm:.1f} cm)")
    return rem_cm


def stage_cut(f0, f1, rot, min_keep, W, enabled):
    """Stage B — pick the loop region [a, b].

    Two independent criteria, because the seam blend (Stage C) already forces seam POSE gap
    and DRIFT to ~0; the only QC failure mode left is the velocity gap, and that depends
    SOLELY on the anchor frame `a` (loop start): the game wraps last→first, and the velocity
    just before vs just after the seam is `ang(orig[a-1],orig[a])` vs `ang(orig[a],orig[a+1])`
    — i.e. the original mocap's own local velocity discontinuity at `a`. So:

      a (anchor) — minimize max per-bone |v_in(a) − v_out(a)| over ALL bones (QC scores all
                   bones; fingers twitch fast, so all-bones is required, not core).
      b (cut)    — among end frames keeping >= min_keep, pick the best core-bone POSE match to
                   `a` (purely cosmetic: it governs how much the crossfade region deviates;
                   the seam pose itself is 0 regardless). Prefer a larger b to keep duration.

    `a` is constrained to a >= W (so the W blend-source frames [a-W, a] exist) and to leave
    room to keep >= min_keep (a <= f1 − K + 1).
    """
    span = f1 - f0 + 1
    K = max(2, int(math.ceil(min_keep * span)))          # minimum frames we must keep
    a_lo = f0 + W
    a_hi = f1 - K + 1                                     # so that b = a+K-1 <= f1
    if a_hi < a_lo:                                       # window too tight — relax to whole clip
        a_lo, a_hi = f0 + W, f1 - 1
    all_bones = list(rot[f0].keys())
    core = [b for b in CORE_BONES if b in rot[f0]]

    if not enabled:
        a, b = f0 + W, f1
        print(f"[loop_fix] B cut: SKIPPED (--no-trim) → anchor a={a} cut b={b}")
        return a, b

    # --- anchor a: smallest all-bone velocity discontinuity ---
    best_a, best_vd, vd_bone = a_lo, float("inf"), ""
    for a in range(a_lo, a_hi + 1):
        disc, wbone = 0.0, ""
        for bn in all_bones:
            v_in = ang(rot[a - 1][bn], rot[a][bn])
            v_out = ang(rot[a][bn], rot[a + 1][bn])
            d = abs(v_in - v_out)
            if d > disc:
                disc, wbone = d, bn
        if disc < best_vd - 1e-9:
            best_vd, best_a, vd_bone = disc, a, wbone

    # --- cut b: best core-bone pose match to anchor a, prefer larger b ---
    a = best_a
    b_lo = a + K - 1
    best_b, best_pose = f1, float("inf")
    for b in range(b_lo, f1 + 1):
        pg = max((ang(rot[b][c], rot[a][c]) for c in core), default=0.0)
        if pg < best_pose - 1e-9 or (abs(pg - best_pose) < 1e-9 and b > best_b):
            best_pose, best_b = pg, b

    print(f"[loop_fix] B cut: anchor a={a} (vel disc {best_vd:.2f}°/f, worst {vd_bone}); "
          f"cut b={best_b} (core pose gap {best_pose:.1f}°); keep {best_b - a + 1}/{span}f "
          f"({100.0 * (best_b - a + 1) / span:.1f}%)")
    return a, best_b


def build_output(f0, a, b, W, bones, rot, loc, do_blend):
    """Produce the rebased output frames [a..b] → [0..(b-a)], applying the seam blend.

    For k in [0,W], original frame b-W+k is blended toward original frame a-W+k:
      blend(original[b-W+k], original[a-W+k], smoothstep(k/W))
    At k=W the blended last frame == original[a] == the first output frame (loop start), so
    pose, drift, and the frames leading into the seam are continuous when the game wraps
    last→first. The frames [a-W, a) are the crossfade SOURCE (dropped from the output).
    """
    out = []   # list of (rot:{bone:q}, loc:(x,y,z))
    for i in range(0, b - a + 1):
        src = a + i                      # original frame index this output frame derives from
        rot_f = dict(rot[src])
        loc_f = loc[src]
        if do_blend and W > 0 and src >= b - W:
            k = src - (b - W)
            t = smoothstep(k / float(W))
            tgt_idx = a - W + k          # the frames BEFORE the loop start (crossfade target)
            tgt = rot[tgt_idx]
            for bn in bones:
                rot_f[bn] = slerp(rot[src][bn], tgt[bn], t)
            loc_f = lerp(loc[src], loc[tgt_idx], t)
        out.append((rot_f, loc_f))
    return out


def purge_animation(arm, keep_action):
    """Drop every action/NLA track except the one we are about to export.

    glTF import leaves the source action on the armature; if we leave it in place the NLA-track
    export ships BOTH clips and loop_qc (which grabs the first action) measures the original.
    Clear the armature's animation_data and remove all other actions from bpy.data.actions.
    """
    ad = arm.animation_data
    if ad is not None:
        ad.action = None
        for tr in list(ad.nla_tracks):
            ad.nla_tracks.remove(tr)
    for act in list(bpy.data.actions):
        if act != keep_action:
            bpy.data.actions.remove(act)


def write_action(arm, name, out_frames, root):
    """Bake the rebased frames into a fresh action (mirrors cc_retarget's keyframe_insert)."""
    # Make room for the clean name first: drop the imported action + any NLA tracks so the new
    # action gets the exact clip name (no .001 suffix) and nothing else ships.
    purge_animation(arm, keep_action=None)
    new_act = bpy.data.actions.new(name)
    if arm.animation_data is None:
        arm.animation_data_create()
    arm.animation_data.action = new_act
    if getattr(new_act, "slots", None) and arm.animation_data.action_slot is None:
        arm.animation_data.action_slot = new_act.slots[0]

    for i, (rot_f, loc_f) in enumerate(out_frames):
        for pb in arm.pose.bones:
            pb.rotation_quaternion = rot_f[pb.name]
            pb.keyframe_insert("rotation_quaternion", frame=i)
        arm.pose.bones[root].location = loc_f
        arm.pose.bones[root].keyframe_insert("location", frame=i)

    # Lock the action to exactly the frames we wrote.
    try:
        new_act.frame_range = (0, len(out_frames) - 1)
    except Exception:
        pass
    arm.animation_data.action = None  # re-added as an NLA track for export
    return new_act


def export(arm, acts, out):
    bpy.ops.object.select_all(action="DESELECT")
    arm.select_set(True)
    bpy.context.view_layer.objects.active = arm
    for act in acts:
        tr = arm.animation_data.nla_tracks.new()
        tr.name = act.name
        tr.strips.new(act.name, 1, act)
    bpy.ops.export_scene.gltf(
        filepath=out, use_selection=True,
        export_animations=True, export_animation_mode="NLA_TRACKS",
        export_force_sampling=True, export_optimize_animation_size=False)
    print(f"[loop_fix] EXPORTED {len(acts)} clip(s) → {out}")


def main():
    a = parse()
    try:
        arm, act = load(a.src)
    except Exception as e:
        print(f"[loop_fix] ERROR load: {e}")
        sys.exit(2)

    name = a.name or act.name
    f0, f1, root, bones, rot, loc = sample_frames(arm, act)
    print(f"[loop_fix] src={a.src}  action={act.name}  frames={f0}..{f1}  bones={len(bones)}")

    stage_dedrift(f0, f1, loc, enabled=not a.no_dedrift)
    anchor, F = stage_cut(f0, f1, rot, a.min_keep, a.blend_frames, enabled=not a.no_trim)

    W = a.blend_frames
    if anchor - W < f0:   # window too big to have W source frames before the anchor — clamp
        W = max(0, anchor - f0)
        print(f"[loop_fix] C blend: clamped W → {W} (not enough frames before anchor)")
    do_blend = (not a.no_blend) and W > 0
    n_out = F - anchor + 1
    print(f"[loop_fix] C blend: W={W} loop=[{anchor}..{F}] → out {n_out}f (kept "
          f"{100.0 * n_out / (f1 - f0 + 1):.1f}%)")

    out_frames = build_output(f0, anchor, F, W, bones, rot, loc, do_blend)
    new_act = write_action(arm, name, out_frames, root)
    export(arm, [new_act], a.out)
    sys.exit(0)


if __name__ == "__main__":
    try:
        main()
    except SystemExit:
        raise
    except Exception as e:
        import traceback
        traceback.print_exc()
        print(f"[loop_fix] ERROR: {e}")
        sys.exit(2)