#!/usr/bin/env python3
"""
Measure authored ground speed from root-motion animation clips.

Usage:
  blender --background --python tools/audit_animation_speeds.py -- \
      <library.glb> <clip> [<clip> ...]

The reported metres/second is the ground-plane displacement of the resolved
hips bone over the clip's own duration. Use the root-motion edition of an
animation library; an in-place edition correctly reports approximately zero.
"""
import os
import sys

import bpy
from mathutils import Vector

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rig_map import RigRoles


argv = sys.argv
argv = argv[argv.index("--") + 1:] if "--" in argv else []
if len(argv) < 2:
    print(__doc__)
    sys.exit(1)

LIBRARY = argv[0]
CLIPS = argv[1:]


def assign_action(armature, action):
    if not armature.animation_data:
        armature.animation_data_create()
    armature.animation_data.action = action
    try:
        if len(action.slots):
            armature.animation_data.action_slot = action.slots[0]
    except (AttributeError, TypeError):
        pass


def hips_world_position(armature, hips_name, frame):
    bpy.context.scene.frame_set(int(round(frame)))
    depsgraph = bpy.context.evaluated_depsgraph_get()
    evaluated = armature.evaluated_get(depsgraph)
    return (
        evaluated.matrix_world @ evaluated.pose.bones[hips_name].matrix
    ).translation


def main():
    bpy.ops.object.select_all(action="SELECT")
    bpy.ops.object.delete()
    bpy.ops.import_scene.gltf(filepath=LIBRARY)

    armatures = [obj for obj in bpy.data.objects if obj.type == "ARMATURE"]
    if len(armatures) != 1:
        print(f"ERROR: expected one armature, found {len(armatures)}")
        sys.exit(1)
    armature = armatures[0]
    roles = RigRoles(armature)
    if not roles.hips:
        print("ERROR: could not resolve hips bone")
        sys.exit(1)

    fps = (
        bpy.context.scene.render.fps
        / bpy.context.scene.render.fps_base
    )
    if CLIPS == ["--list"]:
        for action in sorted(bpy.data.actions, key=lambda item: item.name):
            print(action.name)
        return

    failed = False
    print(f"AUTHORED_SPEED_AUDIT fps={fps:.3f} hips={roles.hips}")
    for clip in CLIPS:
        action = bpy.data.actions.get(clip)
        if action is None:
            print(f"ERROR: missing clip {clip}")
            failed = True
            continue
        assign_action(armature, action)
        first, last = (float(value) for value in action.frame_range)
        start = hips_world_position(armature, roles.hips, first)
        end = hips_world_position(armature, roles.hips, last)
        displacement = end - start
        displacement.z = 0.0
        duration = (last - first) / fps
        speed = displacement.length / duration if duration > 1e-6 else 0.0
        print(
            f"AUTHORED_SPEED {clip} speed={speed:.5f}m/s "
            f"distance={displacement.length:.5f}m "
            f"duration={duration:.5f}s "
            f"delta=({displacement.x:.5f},{displacement.y:.5f})"
        )

    sys.exit(1 if failed else 0)


main()