# El Tuna (bumblebee) — custom hover rig + procedural clips, keyed every frame.
# blender -b -P bee_rig.py -- eltuna_striped.glb out.glb [out.blend]
# The model is posed three-quarters like the author's sheet, so landmarks are measured, not mirrored.
import bpy, sys, math, bmesh
from mathutils import Vector, Quaternion

a = sys.argv[sys.argv.index("--") + 1:]
src, dst, blend = a[0], a[1], (a[2] if len(a) > 2 else None)
bpy.ops.wm.read_factory_settings(use_empty=True)
bpy.ops.import_scene.gltf(filepath=src)
me = [o for o in bpy.context.scene.objects if o.type == "MESH"][0]; me.name = "ElTuna"
bpy.context.view_layer.objects.active = me; me.select_set(True)
bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
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()
G = min(v.co.z for v in me.data.vertices)

LEGS = [(-0.18, -0.31), (0.11, -0.37), (-0.38, -0.29), (0.16, 0.01), (-0.35, -0.01), (-0.60, 0.02)]
B = {"root": ((-0.15, 0.0, G), (-0.15, 0.0, G + 0.2), None),
     "body": ((-0.15, 0.0, -0.05), (-0.15, -0.30, 0.10), "root"),
     "abdomen": ((-0.20, 0.22, -0.10), (-0.32, 0.56, -0.42), "body"),
     "wing.L": ((0.27, -0.07, 0.10), (0.95, -0.08, 0.16), "body"),
     "wing.R": ((-0.50, -0.03, 0.20), (-0.95, 0.05, 0.15), "body"),
     "antenna.L": ((0.0, -0.40, 0.42), (0.12, -0.45, 0.76), "body"),
     "antenna.R": ((-0.30, -0.38, 0.42), (-0.80, -0.30, 0.70), "body")}
for i, (x, y) in enumerate(LEGS):
    B["leg.%d" % i] = ((x, y, -0.40), (x, y, G + 0.02), "body")

ad = bpy.data.armatures.new("ElTunaRig"); arm = bpy.data.objects.new("ElTunaRig", ad)
bpy.context.scene.collection.objects.link(arm); bpy.context.view_layer.objects.active = arm
bpy.ops.object.mode_set(mode="EDIT")
for n, (h, t, p) in B.items():
    eb = ad.edit_bones.new(n); eb.head = h; eb.tail = t
for n, (h, t, p) in B.items():
    if p:
        ad.edit_bones[n].parent = ad.edit_bones[p]
bpy.ops.object.mode_set(mode="OBJECT")


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


groups = {n: me.vertex_groups.new(name=n) for n in B if n != "root"}
stats = {}
for v in me.data.vertices:
    x, y, z = v.co
    if x > 0.25 and z > -0.06:
        r = smooth(x, 0.25, 0.33); w = {"wing.L": r, "body": 1 - r}; k = "wingL"
    elif x < -0.52 and z > -0.06:
        r = smooth(-x, 0.52, 0.60); w = {"wing.R": r, "body": 1 - r}; k = "wingR"
    elif z > 0.42 and y < -0.15:
        w = {("antenna.L" if x > -0.15 else "antenna.R"): 1.0}; k = "antenna"
    elif z < -0.42:
        i = min(range(len(LEGS)), key=lambda j: (x - LEGS[j][0]) ** 2 + (y - LEGS[j][1]) ** 2)
        d = math.hypot(x - LEGS[i][0], y - LEGS[i][1])
        if d < 0.12 or z < -0.55:
            r = smooth(-z, 0.42, 0.48); w = {"leg.%d" % i: r, "body": 1 - r}; k = "leg"
        else:
            w = {"abdomen": 1.0}; k = "belly"
    elif y > 0.20:
        r = smooth(y, 0.18, 0.30); w = {"abdomen": r, "body": 1 - r}; k = "abdomen"
    else:
        w = {"body": 1.0}; k = "body"   # head + thorax: rigid
    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

# ---------------- procedural clips ----------------
X, Y, Z = (1, 0, 0), (0, 1, 0), (0, 0, 1)
for pb in arm.pose.bones:
    pb.rotation_mode = "QUATERNION"


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


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


def key(act_name, frames, pose_fn):
    act = bpy.data.actions.new(act_name); arm.animation_data_create(); arm.animation_data.action = act
    for f in range(1, frames + 1):
        t = (f - 1) / frames
        pose = pose_fn(f, t)
        for pb in arm.pose.bones:
            q = Quaternion()
            for axis, deg in pose.get(pb.name, []):
                q = wrot(pb.name, axis, deg) @ q
            pb.rotation_quaternion = q
            pb.location = wvec(pb.name, pose.get("loc:" + pb.name, (0, 0, 0)))
            pb.keyframe_insert("rotation_quaternion", frame=f, group=pb.name)
            pb.keyframe_insert("location", frame=f, group=pb.name)
    act.use_fake_user = True; arm.animation_data.action = None
    return act


HOVER = 0.14
TAU = 2 * math.pi


def buzz(f, amp=32.0):
    # 6 beats per second at 24 fps: one beat every 4 frames
    s = math.sin(TAU * (f - 1) / 4.0)
    return {"wing.L": [(Y, -amp * s)], "wing.R": [(Y, amp * s)]}


def dangle(t, amp=6.0, back=0.0):
    return {"leg.%d" % i: [(X, back + amp * math.sin(TAU * t + i * 0.9))] for i in range(len(LEGS))}


def hover_pose(f, t, pitch=0.0, back=0.0, bob=0.03):
    p = {"loc:root": (0, 0, HOVER + bob * math.sin(TAU * t)),
         "body": [(X, pitch + 2.5 * math.sin(TAU * t + 1.0)), (Y, 3 * math.sin(TAU * t))],
         "abdomen": [(X, -4 * math.sin(TAU * t + 0.5))],
         "antenna.L": [(X, 9 * math.sin(TAU * t))], "antenna.R": [(X, 9 * math.sin(TAU * t + 1.3))]}
    p.update(buzz(f)); p.update(dangle(t, back=back))
    return p


key("Idle", 24, lambda f, t: hover_pose(f, t))
key("Walk", 24, lambda f, t: hover_pose(f, t, pitch=14, back=12, bob=0.02))


def attack(f, t):
    p = hover_pose(f, t)
    # wind-up (back), sting thrust (forward + abdomen swung under and forward), recover
    w = smooth(f, 1, 6) * (1 - smooth(f, 6, 10))
    s = smooth(f, 6, 10) * (1 - smooth(f, 13, 20))
    p["body"] = [(X, -18 * w + 32 * s)]
    p["abdomen"] = [(X, 10 * w - 55 * s)]
    p["loc:root"] = (0, 0.06 * w - 0.22 * s, HOVER + 0.04 * w - 0.06 * s)
    return p


key("Attack", 20, attack)


def hurt(f, t):
    p = hover_pose(f, t)
    j = smooth(f, 1, 3) * (1 - smooth(f, 4, 12))
    p["body"] = [(X, -22 * j), (Y, 25 * j)]
    p["loc:root"] = (0, 0.10 * j, HOVER + 0.05 * j)
    p.update(buzz(f, 32 + 20 * j))
    return p


key("Hurt", 12, hurt)


def death(f, t):
    p = hover_pose(f, t)
    fall = smooth(f, 4, 22); roll = smooth(f, 6, 24)
    stop = smooth(f, 4, 14)
    p.update(buzz(f, 32 * (1 - stop)))
    p["wing.L"] = p["wing.L"] + [(Y, 35 * stop)]          # wings droop open
    p["wing.R"] = p["wing.R"] + [(Y, -35 * stop)]
    p["root"] = [(Y, 170 * roll)]                          # roll belly-up
    p["loc:root"] = (0, 0.05 * roll, (HOVER * (1 - fall)) + DEATH_LIFT * roll)
    tw = math.sin(TAU * f / 6.0) * 12 * smooth(f, 22, 26) * (1 - smooth(f, 30, 36))
    for i in range(len(LEGS)):
        p["leg.%d" % i] = [(X, -25 * roll + tw * (1 if i % 2 else -1))]
    return p


DEATH_LIFT = 0.0
for _ in range(3):  # adjust the lying height so the bee rests on the ground, belly up
    if "Death" in bpy.data.actions:
        bpy.data.actions.remove(bpy.data.actions["Death"])
    act = key("Death", 36, death)
    arm.animation_data.action = act
    if act.slots:
        arm.animation_data.action_slot = act.slots[0]
    bpy.context.scene.frame_set(36); dg = bpy.context.evaluated_depsgraph_get(); e = me.evaluated_get(dg); m = e.to_mesh()
    low = min((e.matrix_world @ vv.co).z for vv in m.vertices); e.to_mesh_clear(); arm.animation_data.action = None
    print("DEATH lying low %.3f (lift %.3f)" % (low - G, DEATH_LIFT))
    if abs(low - G) < 0.005:
        break
    DEATH_LIFT += G - low

# per-frame snap on Death: never below the ground while falling, resting on it once down
act = bpy.data.actions["Death"]; arm.animation_data.action = act
if act.slots:
    arm.animation_data.action_slot = act.slots[0]
root_pb = arm.pose.bones["root"]; root_inv = arm.data.bones["root"].matrix_local.to_3x3().inverted()
for f in range(1, 37):
    bpy.context.scene.frame_set(f); dg = bpy.context.evaluated_depsgraph_get(); e = me.evaluated_get(dg); m = e.to_mesh()
    low = min((e.matrix_world @ vv.co).z for vv in m.vertices); e.to_mesh_clear()
    dz = (G - low) if f >= 24 else max(0.0, G - low)
    if abs(dz) > 1e-4:
        root_pb.location = root_pb.location + root_inv @ Vector((0, 0, dz))
        root_pb.keyframe_insert("location", frame=f, group="root")
arm.animation_data.action = None

for act in bpy.data.actions:
    arm.animation_data.action = act
    if act.slots:
        arm.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 = me.evaluated_get(dg); m = e.to_mesh()
        lows.append(round(min((e.matrix_world @ vv.co).z for vv in m.vertices) - G, 3)); e.to_mesh_clear()
    print("GROUND", act.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, [x.name for x in bpy.data.actions])