# Fit the Mesh2Motion (CC0) bird rig + clips onto the BIG EGG pirate gull.
# blender -b -P m2m_fit.py -- bird-animations.glb reposed_gull.glb out.glb out.blend
# Joints are translated onto the gull, bone frames are kept: template clips replay the same
# world-space rotations around the gull's own pivots.
import bpy, sys, os, math, mathutils
from mathutils import Vector, Quaternion, Matrix

a = sys.argv[sys.argv.index("--") + 1:]
tpl, gull, dst, blend = a[0], a[1], a[2], (a[3] if len(a) > 3 else None)
bpy.ops.wm.read_factory_settings(use_empty=True)
bpy.ops.import_scene.gltf(filepath=tpl)
arm = [o for o in bpy.context.scene.objects if o.type == "ARMATURE"][0]
for o in [o for o in bpy.context.scene.objects if o.type == "MESH"]:
    bpy.data.objects.remove(o)
arm.name = "GullRig"; arm.data.name = "GullRig"
bpy.ops.import_scene.gltf(filepath=gull)
me = [o for o in bpy.context.scene.objects if o.type == "MESH"][0]
me.name = "PowderGull"
bpy.context.view_layer.objects.active = me; me.select_set(True)
bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
import bmesh
bm = bmesh.new(); bm.from_mesh(me.data)
bmesh.ops.remove_doubles(bm, verts=bm.verts, dist=1e-4)
bm.to_mesh(me.data); bm.free()
assert arm.matrix_world == Matrix.Identity(4), arm.matrix_world

# ---------------- targets ----------------
T = {
    "root": (0, 0, -0.82), "hips": (0, -0.02, -0.25), "spine_0": (0, -0.03, -0.05),
    "spine_1": (0, -0.04, 0.10), "spine_2": (0, -0.05, 0.20), "spine_3": (0, -0.05, 0.26),
    "head": (0, -0.06, 0.32),
    "mouth_upper": (0.10, -0.36, 0.53), "mouth_upper_tip": (0.12, -0.50, 0.51),
    "mouth_lower": (0.10, -0.36, 0.47), "mouth_lower_tip": (0.12, -0.48, 0.45),
    "tail_1": (0, 0.12, -0.30), "tail_2": (-0.04, 0.22, -0.42), "tail_3": (-0.09, 0.32, -0.54),
    "tail_tip": (-0.15, 0.42, -0.66),
    "UpperLeg_L": (0.11, -0.04, -0.30), "LowerLeg_L": (0.11, -0.05, -0.42), "AnkleLeg_L": (0.11, -0.06, -0.55),
    "Foot_L": (0.13, -0.12, -0.73), "Toes_L": (0.15, -0.22, -0.79), "Toes_tip_L": (0.17, -0.33, -0.80),
    "UpperLeg_R": (-0.17, -0.04, -0.30), "LowerLeg_R": (-0.17, -0.05, -0.42), "AnkleLeg_R": (-0.18, -0.06, -0.55),
    "Foot_R": (-0.21, -0.11, -0.73), "Toes_R": (-0.25, -0.22, -0.79), "Toes_tip_R": (-0.29, -0.33, -0.80),
}
# wings: affine map from the template wing frame to the gull's flattened wing
OUR_EDGE = {"L": (Vector((0.27, 0.04, 0.26)), Vector((0.95, 0.17, 0.25))),
            "R": (Vector((-0.27, 0.06, 0.26)), Vector((-0.95, 0.25, 0.25)))}
OUR_CHORD = 0.34
B = arm.data.bones
heads = {b.name: arm.matrix_world @ b.head_local for b in B}


def frame(o, e, side):
    s = (e - o).normalized()
    c = Vector((0, 1, 0)); c = (c - s * c.dot(s)).normalized()
    n = s.cross(c) if side == "L" else c.cross(s)
    return s, c, n


for side in ("L", "R"):
    wnames = [b.name for b in B if b.name.startswith("wing") and b.name.endswith("_" + side)]
    o_t, e_t = heads["wing_2_" + side], heads["wing_tip_" + side]
    s_t, c_t, n_t = frame(o_t, e_t, side)
    span_t = (e_t - o_t).length
    chord_t = max((heads[n] - o_t).dot(c_t) - max(0, (heads[n] - o_t).dot(s_t)) * 0 for n in wnames)
    o_o, e_o = OUR_EDGE[side]
    s_o, c_o, n_o = frame(o_o, e_o, side)
    span_o = (e_o - o_o).length
    k = span_o / span_t
    # chord measured from the leading edge line at the same span position
    lead_c = {}
    for n in wnames:
        if "feather" not in n:
            lead_c[(heads[n] - o_t).dot(s_t)] = (heads[n] - o_t).dot(c_t)
    sp = sorted(lead_c)

    def lead_at(u):
        if u <= sp[0]: return lead_c[sp[0]]
        if u >= sp[-1]: return lead_c[sp[-1]]
        for i in range(len(sp) - 1):
            if sp[i] <= u <= sp[i + 1]:
                t = (u - sp[i]) / (sp[i + 1] - sp[i]); return lead_c[sp[i]] * (1 - t) + lead_c[sp[i + 1]] * t
    depth = max((heads[n] - o_t).dot(c_t) - lead_at((heads[n] - o_t).dot(s_t)) for n in wnames)
    for n in wnames:
        d = heads[n] - o_t
        u, cc, nn = d.dot(s_t), d.dot(c_t), d.dot(n_t)
        behind = cc - lead_at(u)
        T[n] = tuple(o_o + s_o * (u * k) + c_o * (behind / depth * OUR_CHORD) + n_o * (nn * k))
    print("WING", side, "span_t %.2f depth_t %.2f -> span %.2f chord %.2f" % (span_t, depth, span_o, OUR_CHORD))

missing = [b.name for b in B if b.name not in T]
print("UNMAPPED", missing)

# ---------------- move joints (frames kept) ----------------
bpy.ops.object.select_all(action="DESELECT")
bpy.context.view_layer.objects.active = arm; arm.select_set(True)
bpy.ops.object.mode_set(mode="EDIT")
eb = arm.data.edit_bones
for b in eb:
    b.use_connect = False
for b in eb:
    if b.name in T:
        d = Vector(T[b.name]) - b.head
        b.head += d; b.tail += d
bpy.ops.object.mode_set(mode="OBJECT")

# ---------------- segments for envelope weights ----------------
A = arm.matrix_world; seg = {}
for b in arm.data.bones:
    h = A @ b.head_local
    ch = list(b.children)
    if ch:
        t = max((A @ c.head_local for c in ch), key=lambda q: (q - h).length)
    elif b.parent:
        ph = A @ b.parent.head_local; d = h - ph
        t = h + d.normalized() * max(0.03, 0.6 * d.length) if d.length > 1e-6 else h
    else:
        t = h
    seg[b.name] = (h, t)


def dist(p, s0, s1):
    ab = s1 - s0; l = ab.length_squared
    u = 0 if l < 1e-12 else max(0, min(1, (p - s0).dot(ab) / l))
    return (s0 + ab * u - p).length


def envelope(p, names, power=6, keep=2):
    ds = sorted((max(dist(p, *seg[n]), 1e-4), n) for n in names)[:keep]
    w = {n: 1 / d ** power for d, n in ds}; t = sum(w.values())
    return {n: x / t for n, x in w.items()}


def smooth(t, a0, a1):
    u = max(0.0, min(1.0, (t - a0) / (a1 - a0)))
    return u * u * (3 - 2 * u)


names = [b.name for b in arm.data.bones]
wing = {s: [n for n in names if n.startswith("wing") and n.endswith("_" + s) and n != "wing_1_" + s] for s in "LR"}
legs = {s: [n for n in names if n.endswith("_" + s) and ("Leg" in n or "Foot" in n or "Toes" in n) and "tip" not in n] for s in "LR"}
tail = ["tail_1", "tail_2", "tail_3"]
groups = {n: me.vertex_groups.new(name=n) for n in names}
# body chain: each z-slice belongs to the bone whose segment it lies on (head -> next joint),
# blended only around the joints; never weight vertices below a bone's pivot to that bone.
JOINTS = [(-0.05, "hips", "spine_0", 0.06), (0.10, "spine_0", "spine_1", 0.05),
          (0.20, "spine_1", "spine_2", 0.03), (0.26, "spine_2", "spine_3", 0.03)]
HEAD_Z0, HEAD_Z1 = 0.29, 0.345   # neck -> rigid head transition (head pivot at the junction, z 0.32)


def chain_w(z):
    if z >= HEAD_Z0:
        t = smooth(z, HEAD_Z0, HEAD_Z1)
        return {"spine_3": 1 - t, "head": t}
    owner = "hips"
    for zj, lo, hi, b in JOINTS:
        if z < zj - b:
            return {owner: 1.0}
        if z <= zj + b:
            t = smooth(z, zj - b, zj + b)
            return {lo: 1 - t, hi: t}
        owner = hi
    return {owner: 1.0}


LEG_R = 0.075
stats = {}
for v in me.data.vertices:
    p = v.co; x, y, z = p
    side = "L" if x > 0 else "R"
    if z > HEAD_Z1 and not (abs(x) > 0.30 and z < 0.36):
        w = {"head": 1.0}; k = "head"
    elif abs(x) > 0.25 and z > -0.05:
        r = smooth(abs(x), 0.25, 0.31)
        w = {n: r * val for n, val in envelope(p, wing[side]).items()}
        for n, val in chain_w(z).items():
            w[n] = w.get(n, 0) + (1 - r) * val
        k = "wing"
    elif y > 0.24 and z < -0.30:
        w = envelope(p, tail); k = "tail"
    elif z < -0.28:
        legside = "L" if x > -0.03 else "R"
        dleg = min(dist(p, *seg[n]) for n in legs[legside])
        follow = 1.0 if z < -0.62 else 1 - smooth(dleg, 0.05, LEG_R)   # feet always follow
        top = smooth(z, -0.36, -0.28)
        lw = (1 - top) * follow
        w = {n: lw * val for n, val in envelope(p, legs[legside]).items()}
        w["hips"] = w.get("hips", 0) + (1 - lw); k = "leg" if follow > 0.5 else "skirt"
    else:
        w = chain_w(z); k = "body"
    stats[k] = stats.get(k, 0) + 1
    for n, val in w.items():
        if val > 1e-3:
            groups[n].add([v.index], val, "REPLACE")
print("ZONES", stats)
me.parent = arm
mod = me.modifiers.new("Armature", "ARMATURE"); mod.object = arm


# ======================= clip authoring =======================
SCN = bpy.context.scene
arm.animation_data_create()
PB = arm.pose.bones
for pb in PB:
    pb.rotation_mode = "QUATERNION"
for act in list(bpy.data.actions):
    act.name = "_tpl_" + act.name
SRC = {n: bpy.data.actions["_tpl_" + n] for n in ("Flap", "Glide", "Idle", "Walk")}
X_AX = (1, 0, 0); Z_AX = (0, 0, 1)
WING = [n for n in names if n.startswith("wing")]
LEGS = [n for n in names if "Leg" in n or "Foot" in n or "Toes" in n]
NECK = {"spine_2": 0.5, "spine_3": 0.5, "head": 0.6}


def use_action(act):
    arm.animation_data.action = act
    if act.slots:
        arm.animation_data.action_slot = act.slots[0]


def sample(src, f):
    use_action(SRC[src]); SCN.frame_set(int(f), subframe=f - int(f))
    return {pb.name: (pb.location.copy(), pb.rotation_quaternion.copy(), pb.scale.copy()) for pb in PB}


def world_rot(bone, axis, deg):
    m = arm.data.bones[bone].matrix_local.to_3x3().inverted()
    return Quaternion((m @ Vector(axis)).normalized(), math.radians(deg))


def world_vec(bone, vec):
    return arm.data.bones[bone].matrix_local.to_3x3().inverted() @ Vector(vec)


def copy(pose):
    return {n: (l.copy(), q.copy(), s.copy()) for n, (l, q, s) in pose.items()}


def take(pose, other, bones):
    out = copy(pose)
    for n in bones:
        out[n] = tuple(c.copy() for c in other[n])
    return out


def mix(pa, pb_, t):
    return {n: (pa[n][0].lerp(pb_[n][0], t), pa[n][1].slerp(pb_[n][1], t), pa[n][2].lerp(pb_[n][2], t)) for n in pa}


def mod(pose, rot=None, loc=None, neck=True):
    out = copy(pose)
    if neck:
        for n, f in NECK.items():
            l, q, s = out[n]; out[n] = (l, Quaternion().slerp(q, f), s)
    for n, rots in (rot or {}).items():
        l, q, s = out[n]
        for axis, deg in rots:
            q = world_rot(n, axis, deg) @ q
        out[n] = (l, q, s)
    for n, vec in (loc or {}).items():
        l, q, s = out[n]; out[n] = (l + world_vec(n, vec), q, s)
    return out


def write(name, keys):
    act = bpy.data.actions.new(name); use_action(act)
    for f, pose in keys:
        for pb in PB:
            l, q, s = pose[pb.name]
            pb.location = l; pb.rotation_quaternion = q; pb.scale = s
            for path in ("location", "rotation_quaternion", "scale"):
                pb.keyframe_insert(path, frame=f, group=pb.name)
    act.use_fake_user = True
    print("CLIP", name, "keys", len(keys), "range", tuple(act.frame_range))
    return act


# drift check on the template locomotion
for src in ("Walk", "Idle", "Flap"):
    ys = [sample(src, f)["hips"][0] for f in range(int(SRC[src].frame_range[0]), int(SRC[src].frame_range[1]) + 1)]
    print("DRIFT", src, "hips loc range", tuple(round(max(v[i] for v in ys) - min(v[i] for v in ys), 3) for i in range(3)))

FLY_PITCH, FLY_LIFT = 25, 0.25
ATTACK_LOW, DEATH_Z1, DEATH_Z2 = float(os.environ.get("ATTACK_LOW", "0.04")), float(os.environ.get("DEATH_Z1", "0.15")), float(os.environ.get("DEATH_Z2", "0.22"))
OUT = {}
OUT["Idle"] = [(f, mod(sample("Idle", f))) for f in range(1, 31)]
OUT["Walk"] = [(f, mod(sample("Walk", f))) for f in range(1, 13)]
fly = lambda src, f: mod(sample(src, f), rot={"hips": [(X_AX, FLY_PITCH)], "head": [(X_AX, -0.6 * FLY_PITCH)]},
                         loc={"root": (0, 0, FLY_LIFT)})
OUT["Fly"] = [(f, fly("Flap", f)) for f in range(1, 16)]
OUT["Glide"] = [(f, fly("Glide", f)) for f in range(1, 16)]

idle0 = sample("Idle", 1)
up = take(idle0, sample("Flap", 1), WING)
mid = take(idle0, sample("Flap", 5), WING)
glide = take(take(idle0, sample("Glide", 8), WING), sample("Glide", 8), LEGS)
OUT["Attack"] = [
    (1, mod(idle0)),
    (5, mod(up, rot={"hips": [(X_AX, -12)], "head": [(X_AX, -10)]}, loc={"root": (0, 0.06, 0.0)})),
    (9, mod(take(mid, glide, LEGS), rot={"hips": [(X_AX, 35)], "head": [(X_AX, -10)]}, loc={"root": (0, -0.25, 0.20)})),
    (12, mod(glide, rot={"hips": [(X_AX, 55)], "head": [(X_AX, 20)]}, loc={"root": (0, -0.45, 0.14)})),
    (15, mod(glide, rot={"hips": [(X_AX, 42)], "head": [(X_AX, 35)]}, loc={"root": (0, -0.42, ATTACK_LOW)})),
    (20, mod(mix(idle0, glide, 0.3), rot={"hips": [(X_AX, 10)]}, loc={"root": (0, -0.15, 0.09)})),
    (24, mod(idle0)),
]
flare = mix(idle0, take(idle0, sample("Flap", 3), WING), 0.55)
OUT["Hurt"] = [
    (1, mod(idle0)),
    (4, mod(flare, rot={"hips": [(X_AX, -18)], "head": [(X_AX, -25), (Z_AX, 15)]}, loc={"root": (0, 0.08, 0.03)})),
    (8, mod(idle0, rot={"hips": [(X_AX, -6)], "head": [(X_AX, -8)]}, loc={"root": (0, 0.03, 0.0)})),
    (14, mod(idle0)),
]
splay = take(take(idle0, sample("Glide", 8), WING), sample("Glide", 8), LEGS)
OUT["Death"] = [
    (1, mod(idle0)),
    (5, mod(flare, rot={"hips": [(X_AX, -18)], "head": [(X_AX, -25), (Z_AX, 15)]}, loc={"root": (0, 0.08, 0.03)})),
    (12, mod(up, rot={"hips": [(X_AX, -25)], "head": [(X_AX, -30)]}, loc={"root": (0, 0.12, 0.06)})),
    (22, mod(splay, rot={"root": [(X_AX, -60)], "head": [(X_AX, -20), (Z_AX, 20)]}, loc={"root": (0, 0.25, DEATH_Z1)})),
    (30, mod(splay, rot={"root": [(X_AX, -88)], "head": [(X_AX, -10), (Z_AX, 35)],
                         "UpperLeg_L": [(X_AX, -35)], "UpperLeg_R": [(X_AX, -25)]}, loc={"root": (0, 0.32, DEATH_Z2)})),
    (36, mod(splay, rot={"root": [(X_AX, -90)], "head": [(X_AX, -8), (Z_AX, 38)],
                         "UpperLeg_L": [(X_AX, -38)], "UpperLeg_R": [(X_AX, -28)]}, loc={"root": (0, 0.33, DEATH_Z2)})),
]
for name, keys in OUT.items():
    write(name, keys)
for act in [a for a in bpy.data.actions if a.name.startswith("_tpl_")]:
    bpy.data.actions.remove(act)
arm.animation_data.action = None

# ground check
for name in OUT:
    act = bpy.data.actions[name]; use_action(act)
    f0, f1 = map(int, act.frame_range); lows = []
    for f in range(f0, f1 + 1, max(1, (f1 - f0) // 6)):
        SCN.frame_set(f); dg = bpy.context.evaluated_depsgraph_get(); e = me.evaluated_get(dg); m = e.to_mesh()
        lows.append(round(min((e.matrix_world @ vv.co).z for vv in m.vertices), 3)); e.to_mesh_clear()
    print("GROUND", name, lows)
arm.animation_data.action = None
if blend:
    bpy.ops.wm.save_as_mainfile(filepath=blend)
bpy.ops.object.select_all(action="DESELECT"); me.select_set(True); arm.select_set(True)
bpy.context.view_layer.objects.active = arm
bpy.ops.export_scene.gltf(filepath=dst, export_format="GLB", use_selection=True,
                          export_animations=True, export_animation_mode="ACTIONS")
print("EXPORTED", dst, list(OUT))