# Retarget Mesh2Motion (CC0) human clips onto a Meshy 24-bone skeleton, keeping the target's weights.
# blender -b -P retarget_m2m.py -- target.glb out.glb "Idle=Idle_A,Walk=Walk,..." src1.glb [src2.glb ...]
# Each source bone is first aligned (minimal rotation) onto the target's rest direction, so rest poses
# may differ (T-pose source, arms-down target). Then per frame:
#   target_world = (source_world * aligned_source_rest^-1) * target_rest      (rotations, world space)
# Only the hips carry translation, scaled by the hip-height ratio.
import bpy, sys, os
from mathutils import Matrix, Vector, Quaternion

a = sys.argv[sys.argv.index("--") + 1:]
tgt_path, out_path, spec, srcs = a[0], a[1], a[2], a[3:]
KEEP = [k for k in os.environ.get("KEEP_ACTIONS", "").split(",") if k]

bpy.ops.wm.read_factory_settings(use_empty=True)
bpy.ops.import_scene.gltf(filepath=tgt_path)
for o in list(bpy.context.scene.objects):
    if o.type == "MESH" and o.name.startswith("Icosphere"):
        bpy.data.objects.remove(o)
tgt = [o for o in bpy.context.scene.objects if o.type == "ARMATURE"][0]
tmesh = [o for o in bpy.context.scene.objects if o.type == "MESH"][0]
for act in list(bpy.data.actions):
    if not any(act.name.startswith(k) for k in KEEP):
        bpy.data.actions.remove(act)
    else:
        print("KEEP", act.name)
tgt.animation_data_create(); tgt.animation_data.action = None

src_of = {}
for sp in srcs:
    before_o, before_a = set(bpy.data.objects), set(bpy.data.actions)
    bpy.ops.import_scene.gltf(filepath=sp)
    new = [o for o in bpy.data.objects if o not in before_o]
    arm = [o for o in new if o.type == "ARMATURE"][0]
    for o in new:
        if o.type != "ARMATURE":
            bpy.data.objects.remove(o)
    for act in list(bpy.data.actions):
        if act not in before_a:
            src_of[act.name] = (arm, act)
            act.name = "src:" + act.name     # never collide with the clips we create
print("SOURCE CLIPS", len(src_of))

MAP = {"Hips": "pelvis", "Spine02": "spine_01", "Spine01": "spine_02", "Spine": "spine_03", "neck": "neck_01",
       "Head": "head"}
for s, t in (("Left", "l"), ("Right", "r")):
    MAP.update({s + "Shoulder": "clavicle_" + t, s + "Arm": "upperarm_" + t, s + "ForeArm": "lowerarm_" + t,
                s + "Hand": "hand_" + t, s + "UpLeg": "thigh_" + t, s + "Leg": "calf_" + t,
                s + "Foot": "foot_" + t, s + "ToeBase": "ball_" + t})
# direction used for alignment: (from, to) joint names on each skeleton
TDIR = {"Hips": ("Hips", "Spine02"), "Spine02": ("Spine02", "Spine01"), "Spine01": ("Spine01", "Spine"),
        "Spine": ("Spine", "neck"), "neck": ("neck", "Head"), "Head": ("Head", "head_end")}
SDIR = {"pelvis": ("pelvis", "spine_01"), "spine_01": ("spine_01", "spine_02"), "spine_02": ("spine_02", "spine_03"),
        "spine_03": ("spine_03", "neck_01"), "neck_01": ("neck_01", "head"), "head": ("head", "head_leaf")}
for s, t in (("Left", "l"), ("Right", "r")):
    TDIR.update({s + "Shoulder": (s + "Shoulder", s + "Arm"), s + "Arm": (s + "Arm", s + "ForeArm"),
                 s + "ForeArm": (s + "ForeArm", s + "Hand"), s + "Hand": (s + "ForeArm", s + "Hand"),
                 s + "UpLeg": (s + "UpLeg", s + "Leg"), s + "Leg": (s + "Leg", s + "Foot"),
                 s + "Foot": (s + "Foot", s + "ToeBase"), s + "ToeBase": (s + "Foot", s + "ToeBase")})
    SDIR.update({"clavicle_" + t: ("clavicle_" + t, "upperarm_" + t), "upperarm_" + t: ("upperarm_" + t, "lowerarm_" + t),
                 "lowerarm_" + t: ("lowerarm_" + t, "hand_" + t), "hand_" + t: ("lowerarm_" + t, "hand_" + t),
                 "thigh_" + t: ("thigh_" + t, "calf_" + t), "calf_" + t: ("calf_" + t, "foot_" + t),
                 "foot_" + t: ("foot_" + t, "ball_" + t), "ball_" + t: ("foot_" + t, "ball_" + t)})


def rot_of(m):
    return m.to_3x3().normalized().to_quaternion()


TA = tgt.matrix_world; TA_rot = rot_of(TA); TA_inv = TA.inverted()
t_rest_w = {b.name: TA @ b.matrix_local for b in tgt.data.bones}
order = []
def walk(b):
    order.append(b); [walk(c) for c in b.children]
for b in tgt.data.bones:
    if b.parent is None:
        walk(b)
GROUND = min((tmesh.matrix_world @ v.co).z for v in tmesh.data.vertices)
# feet = vertices whose dominant bone is a lower-leg/foot bone: standing clips are snapped on them only
# (long arms or claws may hang lower than the feet and must not lift the body)
FOOT_BONES = {"LeftLeg", "RightLeg", "LeftFoot", "RightFoot", "LeftToeBase", "RightToeBase"}
gname = {g.index: g.name for g in tmesh.vertex_groups}
FEET = [v.index for v in tmesh.data.vertices
        if v.groups and gname[max(v.groups, key=lambda g: g.weight).group] in FOOT_BONES]
FEET_GROUND = min((tmesh.matrix_world @ tmesh.data.vertices[i].co).z for i in FEET)
WHOLE_BODY = [k for k in os.environ.get("WHOLE_BODY", "Death,Roll,Stomp").split(",") if k]
print("FEET vertices", len(FEET), "ground", round(GROUND, 3), "feet ground", round(FEET_GROUND, 3))
tgt_hips_h = t_rest_w["Hips"].translation.z - GROUND

for pb in tgt.pose.bones:
    pb.rotation_mode = "QUATERNION"


def build(out_name, src_name):
    sarm, sact = src_of[src_name]
    SA = sarm.matrix_world
    s_rest_w = {b.name: SA @ b.matrix_local for b in sarm.data.bones}
    ref = {}
    for tb, sb in MAP.items():
        t0, t1 = TDIR[tb]; s0, s1 = SDIR[sb]
        tdir = (t_rest_w[t1].translation - t_rest_w[t0].translation).normalized()
        sdir = (s_rest_w[s1].translation - s_rest_w[s0].translation).normalized()
        ref[tb] = sdir.rotation_difference(tdir) @ rot_of(s_rest_w[sb])
    src_hips_h = s_rest_w["pelvis"].translation.z
    k = tgt_hips_h / src_hips_h
    sarm.animation_data_create(); sarm.animation_data.action = sact
    if sact.slots:
        sarm.animation_data.action_slot = sact.slots[0]
    act = bpy.data.actions.new(out_name)
    assert act.name == out_name, ("name clash", act.name)
    tgt.animation_data.action = act
    f0, f1 = map(int, sact.frame_range)
    scn = bpy.context.scene
    for f in range(f0, f1 + 1):
        scn.frame_set(f)
        pose_arm = {}
        for b in order:
            n = b.name
            rest_arm = b.matrix_local
            if b.parent:
                rest_rel = b.parent.matrix_local.inverted() @ rest_arm
                chain = pose_arm[b.parent.name] @ rest_rel
            else:
                rest_rel = rest_arm
                chain = rest_arm.copy()
            if n in MAP:
                spb = sarm.pose.bones[MAP[n]]
                w_src = rot_of(SA @ spb.matrix)
                w_tgt = (w_src @ ref[n].inverted()) @ rot_of(t_rest_w[n])
                r_arm = TA_rot.inverted() @ w_tgt
                if n == "Hips":
                    d = (SA @ spb.matrix).translation - s_rest_w["pelvis"].translation
                    head = TA_inv @ (t_rest_w["Hips"].translation + d * k)
                else:
                    head = chain.translation
                m = Matrix.Translation(head) @ r_arm.to_matrix().to_4x4()
            else:
                m = chain
            pose_arm[n] = m
            parent_pose = pose_arm[b.parent.name] if b.parent else Matrix.Identity(4)
            basis = rest_rel.inverted() @ parent_pose.inverted() @ m
            loc, rq, sc = basis.decompose()
            pb = tgt.pose.bones[n]
            pb.location, pb.rotation_quaternion, pb.scale = loc, rq, Vector((1, 1, 1))
            for path in ("location", "rotation_quaternion", "scale"):
                pb.keyframe_insert(path, frame=f - f0 + 1, group=n)
    act.use_fake_user = True
    # ground snap: per frame, move the hips so the lowest vertex sits on the rest ground
    if act.slots:
        tgt.animation_data.action_slot = act.slots[0]
    hips = tgt.pose.bones["Hips"]; rest_rot = tgt.data.bones["Hips"].matrix_local.to_3x3().normalized()
    TA3inv = TA.to_3x3().inverted()
    shifts = []
    for f in range(1, f1 - f0 + 2):
        scn.frame_set(f); dg = bpy.context.evaluated_depsgraph_get()
        e = tmesh.evaluated_get(dg); me_ = e.to_mesh()
        if any(out_name.startswith(k) for k in WHOLE_BODY):
            low = min((e.matrix_world @ v.co).z for v in me_.vertices); ref_ground = GROUND
        else:
            low = min((e.matrix_world @ me_.vertices[i].co).z for i in FEET); ref_ground = FEET_GROUND
        e.to_mesh_clear()
        dz = ref_ground - low; shifts.append(dz)
        hips.location = hips.location + rest_rot.inverted() @ (TA3inv @ Vector((0, 0, dz)))
        hips.keyframe_insert("location", frame=f, group="Hips")
    tgt.animation_data.action = None
    print("CLIP %-8s <- %-22s frames %d  snap %.3f..%.3f" % (out_name, src_name, f1 - f0 + 1, min(shifts), max(shifts)))
    return act


for item in spec.split(","):
    out_name, src_name = item.split("=")
    build(out_name, src_name)

# ground check per clip
for act in [x for x in bpy.data.actions if x.name in [i.split("=")[0] for i in spec.split(",")]]:
    tgt.animation_data.action = act
    if act.slots:
        tgt.animation_data.action_slot = act.slots[0]
    f0, f1 = map(int, act.frame_range); lows = []
    for f in range(f0, f1 + 1, max(1, (f1 - f0) // 6)):
        bpy.context.scene.frame_set(f); dg = bpy.context.evaluated_depsgraph_get()
        e = tmesh.evaluated_get(dg); m = e.to_mesh()
        lows.append(round(min((e.matrix_world @ v.co).z for v in m.vertices), 3)); e.to_mesh_clear()
    print("GROUND", act.name, lows)
tgt.animation_data.action = None

# drop sources before export
for sarm, _ in {v[0].name: v for v in src_of.values()}.values():
    bpy.data.objects.remove(sarm)
keep_names = set(i.split("=")[0] for i in spec.split(",")) | set(x.name for x in bpy.data.actions if any(x.name.startswith(k) for k in KEEP))
for act in list(bpy.data.actions):
    if act.name not in keep_names:
        bpy.data.actions.remove(act)
bpy.ops.object.select_all(action="DESELECT"); tmesh.select_set(True); tgt.select_set(True)
bpy.context.view_layer.objects.active = tgt
bpy.ops.export_scene.gltf(filepath=out_path, export_format="GLB", use_selection=True,
                          export_animations=True, export_animation_mode="ACTIONS")
print("EXPORTED", out_path, sorted(keep_names))