# 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, 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.22), "spine_3": (0, -0.05, 0.30),
    "head": (0, -0.06, 0.40),
    "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}
stats = {}
for v in me.data.vertices:
    p = v.co; x, y, z = p
    side = "L" if x > 0 else "R"
    if z > 0.32 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()}
        w["spine_1"] = w.get("spine_1", 0) + (1 - r); k = "wing"
    elif y > 0.24 and z < -0.30:
        w = envelope(p, tail); k = "tail"
    elif z < -0.30:
        top = smooth(z, -0.36, -0.28)
        legside = "L" if x > -0.03 else "R"
        w = {n: (1 - top) * val for n, val in envelope(p, legs[legside]).items()}
        w["hips"] = w.get("hips", 0) + top; k = "leg"
    else:
        c0 = smooth(z, -0.18, -0.08); c1 = smooth(z, 0.0, 0.08); n2 = smooth(z, 0.20, 0.26); n3 = smooth(z, 0.27, 0.32)
        w = {"hips": 1 - c0, "spine_0": c0 * (1 - c1), "spine_1": c1 * (1 - n2), "spine_2": n2 * (1 - n3), "head": n3}
        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

# ---------------- diagnostics: how far each template clip moves the body frame ----------------
def use_action(act):
    arm.animation_data.action = act
    if getattr(arm.animation_data, "action_slot", None) is None and act.slots:
        arm.animation_data.action_slot = act.slots[0]


for act in bpy.data.actions:
    use_action(act)
    bpy.context.scene.frame_set(int(act.frame_range[0]))
    out = []
    for n in ("root", "hips", "spine_1", "head"):
        q = arm.pose.bones[n].matrix.to_quaternion() @ arm.data.bones[n].matrix_local.to_quaternion().inverted()
        out.append("%s %.0f°" % (n, math.degrees(q.angle)))
    print("CLIPDIAG %-10s" % act.name, ", ".join(out))
arm.animation_data.action = None
if blend:
    bpy.ops.wm.save_as_mainfile(filepath=blend)
print("FIT DONE")