#!/usr/bin/env python3
"""Build the compact CMU backward-locomotion animation library.

Usage:
  blender --background --python tools/build_cmu_locomotion_library.py -- \
      <76_11.bvh> <output.glb>

Subject 76 trial 11 is CMU's optical capture described as "quick large steps
backwards." The script finds a naturally matching gait-cycle boundary, keeps
the captured bone motion intact, removes gameplay-owned horizontal root travel,
and exports only that cycle.
"""

import os
import sys

import bpy


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

SOURCE, OUTPUT = argv
SEARCH_START = range(170, 281)
SEARCH_END = range(300, 411)
MIN_CYCLE_FRAMES = 95
MAX_CYCLE_FRAMES = 175
SIGNATURE_BONES = (
    "Hips", "LowerBack", "Spine", "Spine1", "Neck", "Head",
    "LeftArm", "LeftForeArm", "RightArm", "RightForeArm",
    "LeftUpLeg", "LeftLeg", "LeftFoot",
    "RightUpLeg", "RightLeg", "RightFoot",
)


def action_fcurves(action):
    """Blender 5.x channelbag f-curves for an imported BVH action."""
    curves = []
    for layer in action.layers:
        for strip in layer.strips:
            for slot in action.slots:
                channelbag = strip.channelbag(slot)
                if channelbag is not None:
                    curves.extend(channelbag.fcurves)
    return curves


def pose_signature(armature, frame):
    bpy.context.scene.frame_set(frame)
    signature = []
    for name in SIGNATURE_BONES:
        bone = armature.pose.bones.get(name)
        if bone is None:
            continue
        signature.append(bone.matrix.to_quaternion().normalized())
    hips = armature.pose.bones["Hips"].matrix.translation
    return signature, hips.copy()


def find_cycle(armature):
    samples = {
        frame: pose_signature(armature, frame)
        for frame in range(min(SEARCH_START), max(SEARCH_END) + 1)
    }
    best = None
    for start in SEARCH_START:
        start_rotations, start_hips = samples[start]
        for end in SEARCH_END:
            duration = end - start
            if duration < MIN_CYCLE_FRAMES or duration > MAX_CYCLE_FRAMES:
                continue
            end_rotations, end_hips = samples[end]
            rotation_error = sum(
                rotation.rotation_difference(other).angle ** 2
                for rotation, other in zip(start_rotations, end_rotations)
            ) / max(len(start_rotations), 1)
            # Match the captured vertical phase but intentionally ignore
            # horizontal displacement: it proves the performer travelled.
            vertical_error = (end_hips.z - start_hips.z) ** 2 * 0.02
            score = rotation_error + vertical_error
            if best is None or score < best[0]:
                best = (score, start, end)
    if best is None:
        raise RuntimeError("No plausible gait cycle found in the capture")
    return best


def crop_action(action, start, end):
    shift = start - 1
    curves = action_fcurves(action)
    for curve in curves:
        points = curve.keyframe_points
        for index in range(len(points) - 1, -1, -1):
            if points[index].co.x < start or points[index].co.x > end:
                points.remove(points[index])
        for point in points:
            point.co.x -= shift
            point.handle_left.x -= shift
            point.handle_right.x -= shift

    # The gameplay body owns horizontal travel. Preserve every captured
    # rotation and the vertical compression/lift, but keep the cycle in place.
    hips_path = 'pose.bones["Hips"].location'
    for curve in curves:
        if curve.data_path != hips_path or curve.array_index not in (0, 1):
            continue
        baseline = curve.evaluate(1.0)
        for point in curve.keyframe_points:
            point.co.y -= baseline
            point.handle_left.y -= baseline
            point.handle_right.y -= baseline


def main():
    bpy.ops.wm.read_factory_settings(use_empty=True)
    bpy.context.scene.render.fps = 120
    bpy.context.scene.render.fps_base = 1.0
    bpy.ops.import_anim.bvh(
        filepath=SOURCE,
        axis_forward="-Z",
        axis_up="Y",
    )
    armature = next(
        (obj for obj in bpy.data.objects if obj.type == "ARMATURE"), None
    )
    if armature is None or armature.animation_data is None:
        raise RuntimeError(f"{SOURCE}: no animated armature")
    action = armature.animation_data.action
    if action is None:
        raise RuntimeError(f"{SOURCE}: no action")

    score, start, end = find_cycle(armature)
    crop_action(action, start, end)
    action.name = "RunBackward"
    action.use_fake_user = True
    bpy.context.scene.frame_start = 1
    bpy.context.scene.frame_end = end - start + 1

    track = armature.animation_data.nla_tracks.new()
    track.name = "RunBackward"
    strip = track.strips.new("RunBackward", 0, action)
    strip.name = "RunBackward"
    track.mute = True
    armature.animation_data.action = None

    os.makedirs(os.path.dirname(os.path.abspath(OUTPUT)), exist_ok=True)
    bpy.ops.export_scene.gltf(
        filepath=OUTPUT,
        export_format="GLB",
        export_yup=True,
        export_skins=True,
        export_animations=True,
        export_animation_mode="NLA_TRACKS",
        export_bake_animation=False,
        export_optimize_animation_size=True,
        export_optimize_animation_keep_anim_armature=False,
    )
    print(
        f"CMU 76_11 frames {start}..{end} -> RunBackward "
        f"({(end - start) / 120.0:.3f}s, phase score {score:.6f})"
    )
    print(f"Wrote {OUTPUT}")


main()