# lena_leafbikini lane, stage 01: graft the AccuRig skeleton onto the pristine leaf-bikini GLB.
#
#   blender --background --factory-startup --python 01_graft.py -- <pristine.glb> <rigged.fbx> <bait.fbx> <out.blend> <out.glb>
#
# Lane rule (`.agents/plans/rig-graft-lane-2026-08-04.md`): the FBX is a disposable rig carrier, the
# GLB is the sole source of truth for mesh and materials, and it only ever GAINS bones and weights.
# Nothing here touches a vertex position, a UV or a texture — asserted at the end, not hoped for.
#
# NEAREST-SURFACE, not index-exact — and that is the difference from work/mako/02_graft.py. Mako's
# AccuRig run took the shipping mesh whole and handed back the same 122,062 verts in order, so
# weights copied index-to-index. Here AccuRig was fed a 51,682v bait decimated 20:1 from a
# 1,029,360v pristine mesh, so there is no index correspondence and proximity is the only mechanism.
# This is the case the lane was actually designed for: the bait's hands were budgeted at 40k tris
# (not the July lane's 10k) precisely to bound how crisply per-finger weights land back on the
# full-res hands. See the FINGER BLEED gate below for the check that this bought what it claims.
#
# THE BAIT IS THE BRIDGE, and it is why no scale is guessed here. AccuRig returns the carrier moved
# and rescaled (the same family as the July lane's 4 cm offset), so the returned skeleton is in some
# AccuRig space, not the GLB's. But the rigged FBX is index-exact against the BAIT (51,682v both,
# same order), and the bait is in the GLB's space BY CONSTRUCTION — rigbait_decimate.py applies no
# scale and no recenter, which the probe confirms exactly: bait and GLB share height 0.9792 and
# half-span 0.4583 to four decimals. So fitting rigged->bait index-exact recovers AccuRig's transform
# in closed form, and the residual of that fit is a gate: it must be near zero, or the carrier that
# came back is not the carrier that went in.
import bpy, sys, os, time
import numpy as np
from mathutils import Vector, Matrix

argv = sys.argv[sys.argv.index("--") + 1:]
GLB, FBX, BAIT, OUTB, OUTG = argv[0], argv[1], argv[2], os.path.abspath(argv[3]), os.path.abspath(argv[4])
t0 = time.time()


def log(m):
    print(f"[graft {time.time()-t0:6.1f}s] {m}", flush=True)


def world_co(o):
    n = len(o.data.vertices)
    a = np.empty(n * 3)
    o.data.vertices.foreach_get("co", a)
    a = a.reshape(-1, 3)
    M = np.array(o.matrix_world)
    return a @ M[:3, :3].T + M[:3, 3]


# ── target: the pristine GLB ──────────────────────────────────────────────────
bpy.ops.wm.read_factory_settings(use_empty=True)
bpy.ops.import_scene.gltf(filepath=GLB)
for _o in bpy.data.objects:
    if _o.type == "MESH":
        bpy.context.view_layer.objects.active = _o
        _o.select_set(True)
bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
body = max([o for o in bpy.data.objects if o.type == 'MESH'], key=lambda o: len(o.data.vertices))
me = body.data
n = len(me.vertices)
V_ref = world_co(body)
UNITM = 1.750 / (V_ref[:, 2].max() - V_ref[:, 2].min())   # rough units->mm scale for reporting
log(f"target: '{body.name}' {n}v {len(me.polygons)}f uv={[l.name for l in me.uv_layers]} "
    f"mats={[m.name for m in me.materials if m]}")

# ── carriers: the AccuRig return, and the bait it was made from ───────────────
before = {o.name for o in bpy.data.objects}
bpy.ops.import_scene.fbx(filepath=FBX)
new = [o for o in bpy.data.objects if o.name not in before]
arm = next(o for o in new if o.type == 'ARMATURE')
src = max([o for o in new if o.type == 'MESH'], key=lambda o: len(o.data.vertices))
log(f"carrier: armature '{arm.name}' {len(arm.data.bones)} bones, mesh {len(src.data.vertices)}v")

before = {o.name for o in bpy.data.objects}
bpy.ops.import_scene.fbx(filepath=BAIT)
bait = max([o for o in bpy.data.objects if o.name not in before and o.type == 'MESH'],
           key=lambda o: len(o.data.vertices))
log(f"bait: '{bait.name}' {len(bait.data.vertices)}v")

# AccuRig ships a "T-Pose" take and Blender's importer binds it as an ACTION. An action re-applies
# its pose on every evaluation and on every file load, so clearing pose values while it exists is
# futile — the in-session assert passes and the saved file comes back posed. Drop it first.
# (Verbatim from work/mako/02_graft.py; the failure it describes is not hypothetical.)
dropped = []
for o in (arm, src):
    if o.animation_data:
        dropped.append(o.name)
        o.animation_data_clear()
    dat = o.data
    if dat and getattr(dat, "animation_data", None):
        dat.animation_data_clear()
    sk = getattr(dat, "shape_keys", None)
    if sk and sk.animation_data:
        sk.animation_data_clear()
for act in list(bpy.data.actions):
    if act.users == 0:
        bpy.data.actions.remove(act)
log(f"animation data cleared on {dropped}; actions left: {[a.name for a in bpy.data.actions]}")

assert len(src.data.vertices) == len(bait.data.vertices), (
    f"rigged carrier {len(src.data.vertices)}v vs bait {len(bait.data.vertices)}v — AccuRig did not "
    f"return the mesh it was given, so the bait cannot bridge the two spaces")

# ── recover AccuRig's transform: fit rigged -> bait, index-exact ──────────────
R = world_co(src)
B = world_co(bait)
ca, cb = B.mean(axis=0), R.mean(axis=0)
A_, B_ = B - ca, R - cb
s = float((A_ * B_).sum() / (B_ * B_).sum())
t = ca - s * cb
res = np.linalg.norm(B - (s * R + t), axis=1)
log(f"fit: scale {s:.6f} (1/{1/s:.6f})   translation {np.round(t, 6)} in target units")
log(f"  carrier offset removed: {np.round(-t / s * 1000, 3)} mm in carrier space")
res_mm = res * UNITM * 1000
frac_1mm = float((res_mm > 1.0).mean())
log(f"  index-exact residual vs bait: mean {res_mm.mean():.4f} mm  p50 {np.percentile(res_mm,50):.4f} "
    f"mm  p99 {np.percentile(res_mm,99):.4f} mm  max {res_mm.max():.4f} mm  "
    f">1mm {100*frac_1mm:.2f}%")
# What this gate is actually for: catching a carrier that came back RESAMPLED — a different surface
# wearing the same vertex count, which would poison every one of the 1M nearest-surface lookups
# downstream. It is NOT for catching AccuRig's joint-region cleanup, which is localized and harmless.
# Measured on this carrier (2026-08-13): p50 0.0009 mm — bit-exact for the bulk — with 1.01% of verts
# over 1 mm, all of them one cluster at height fraction 0.69-0.75 around x=-0.23, i.e. a single
# armpit, where AccuRig routinely tidies the arm/torso crease. The hands, which is where a bad
# transfer actually hurts, come back at mean 0.020 mm / max 1.45 mm.
# So gate on the SHAPE of the distribution, not on a lone max: a global resample moves the mean and
# the tail together, a local touch-up moves neither. An earlier max-only threshold of 1.0 mm failed
# this carrier on 523 verts out of 51,682 and would have blocked a graft that is fine.
assert res_mm.mean() < 0.5, (
    f"mean carrier deviation {res_mm.mean():.3f} mm — AccuRig resampled the whole surface, and "
    f"nearest-surface transfer would inherit that error everywhere")
assert frac_1mm < 0.03, (
    f"{100*frac_1mm:.2f}% of carrier verts moved over 1 mm — too widespread to be joint cleanup")
assert res_mm.max() < 10.0, (
    f"carrier deviates by up to {res_mm.max():.3f} mm — that is a limb moving, not a crease tidied")

M_fit = Matrix.Translation(Vector(t)) @ Matrix.Diagonal(Vector((s, s, s))).to_4x4()

# Move BOTH the skeleton and the carrier mesh into target space. Mako's graft moved only the
# armature because index-to-index weight copying never reads carrier positions; nearest-surface
# does, so the carrier surface has to physically land on the target surface.
bpy.ops.object.select_all(action='DESELECT')
src.select_set(True)
bpy.context.view_layer.objects.active = src
bpy.ops.object.parent_clear(type='CLEAR_KEEP_TRANSFORM')
for o in (arm, src):
    o.matrix_world = M_fit @ o.matrix_world
    bpy.ops.object.select_all(action='DESELECT')
    o.select_set(True)
    bpy.context.view_layer.objects.active = o
    bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
log(f"armature + carrier placed; armature scale now {tuple(round(v,5) for v in arm.scale)}")

# Clear the carrier's POSE, and do it AFTER transform_apply. The mesh binds to the REST skeleton, so
# any pose riding along is pure displacement — and this FBX arrives with one. Clearing it by setting
# location/rotation/scale BEFORE the apply does not survive: transform_apply reintroduces non-identity
# rotations and the saved body deforms into the air with every edge length unchanged (rigid motion,
# which reads as "weights are fine, model is gone"). Use the pose operator, last, and assert it took.
bpy.context.view_layer.objects.active = arm
bpy.ops.object.mode_set(mode='POSE')
bpy.ops.pose.select_all(action='SELECT')
bpy.ops.pose.transforms_clear()
bpy.ops.object.mode_set(mode='OBJECT')
bpy.context.view_layer.update()
worst_pb = max(np.abs(np.array(pb.matrix_basis) - np.eye(4)).max() for pb in arm.pose.bones)
log(f"pose cleared; largest deviation from identity across {len(arm.pose.bones)} bones: {worst_pb:.2e}")
assert worst_pb < 1e-6, "pose did not clear — the bind would be displaced"

# how far the carrier surface sits from the target surface it must hand weights to
S_now = world_co(src)
gap = np.linalg.norm(S_now - B, axis=1)
log(f"carrier now sits on target space: max deviation from bait {gap.max()*UNITM*1000:.5f} mm")

# ── nearest-surface weight transfer ──────────────────────────────────────────
# The carrier still carries an Armature modifier from the FBX; data_transfer reads EVALUATED
# geometry, so leaving it in would sample a deformed surface. The pose is identity by now, but
# relying on that is a silent dependency — remove it.
for m_ in list(src.modifiers):
    if m_.type == 'ARMATURE':
        src.modifiers.remove(m_)
for vg in list(body.vertex_groups):
    body.vertex_groups.remove(vg)

bpy.ops.object.select_all(action='DESELECT')
body.select_set(True)
src.select_set(True)
bpy.context.view_layer.objects.active = src        # ACTIVE is the source, selected is the target
log(f"transferring {len(src.vertex_groups)} vertex groups onto {n} verts (POLYINTERP_NEAREST)...")
bpy.ops.object.data_transfer(
    use_reverse_transfer=False,
    data_type='VGROUP_WEIGHTS',
    vert_mapping='POLYINTERP_NEAREST',   # project onto nearest face, barycentric-interpolate
    layers_select_src='ALL',
    layers_select_dst='NAME',
    mix_mode='REPLACE',
)
log(f"transfer done; target now has {len(body.vertex_groups)} groups")

# ── bind ─────────────────────────────────────────────────────────────────────
for m_ in list(body.modifiers):
    if m_.type == 'ARMATURE':
        body.modifiers.remove(m_)
mod = body.modifiers.new("Armature", 'ARMATURE')
mod.object = arm
body.parent = arm
body.matrix_parent_inverse = arm.matrix_world.inverted()

# ── the carriers are disposable: nothing visual survives from either FBX ─────
carrier_mats = [m_ for m_ in src.data.materials if m_] + [m_ for m_ in bait.data.materials if m_]
bpy.data.objects.remove(src, do_unlink=True)
bpy.data.objects.remove(bait, do_unlink=True)
for m_ in carrier_mats:
    if m_.users == 0:
        bpy.data.materials.remove(m_)
for im in list(bpy.data.images):
    if im.users == 0 and im.name not in {"Render Result", "Viewer Node"}:
        bpy.data.images.remove(im)

# ── verify ───────────────────────────────────────────────────────────────────
V2 = world_co(body)
moved = np.linalg.norm(V2 - V_ref, axis=1).max()

gidx = {g.index: g.name for g in body.vertex_groups}
tot = np.zeros(n)
nz = np.zeros(n, dtype=np.int32)
for v in me.vertices:
    ssum = 0.0
    c = 0
    for g in v.groups:
        w = g.weight
        if w > 1e-6:
            ssum += w
            c += 1
    tot[v.index] = ssum
    nz[v.index] = c

print("\n=== VERIFY ===")
print(f"  mesh vertices moved: {moved*UNITM*1000:.6f} mm   (must be 0)")
print(f"  height {V2[:,2].max()-V2[:,2].min():.5f} units, feet min-Z {V2[:,2].min():+.6f}")
print(f"  weight sums: min {tot.min():.4f} mean {tot.mean():.4f} max {tot.max():.4f}")
print(f"  unweighted vertices: {int((tot < 1e-6).sum())} of {n}")
print(f"  influences/vert: mean {nz.mean():.2f} max {nz.max()}")
print(f"  bones {len(arm.data.bones)}   vertex groups {len(body.vertex_groups)}")
print(f"  uv layers {[l.name for l in me.uv_layers]}   materials {[m_.name for m_ in me.materials if m_]}")
print(f"  images {[(i.name, i.size[0]) for i in bpy.data.images if i.size[0]]}")
assert moved < 1e-9, "the mesh moved — the graft must be additive only"
assert int((tot < 1e-6).sum()) == 0, "unweighted vertices — nearest-surface transfer left holes"

# FINGER BLEED gate. This is the failure this whole lane is shaped around: a vertex on one finger
# finding a NEIGHBOURING finger's surface nearer than its own, which mangles hands the moment a
# clip curls them. A clean transfer keeps each finger's flesh dominated by its own chain, so for
# every finger bone, measure how much of the weight it hands out lands on verts whose dominant
# influence belongs to a DIFFERENT digit. Reported per digit rather than asserted blind: the number
# is the QC evidence, and the pose sweep in 02_qc.py is what confirms it visually.
DIGITS = ("Thumb", "Index", "Mid", "Ring", "Pinky")
dom = np.zeros(n, dtype=np.int32) - 1
best = np.zeros(n)
for v in me.vertices:
    for g in v.groups:
        if g.weight > best[v.index]:
            best[v.index] = g.weight
            dom[v.index] = g.group


def digit_of(name):
    for i, d in enumerate(DIGITS):
        if f"_{d}" in name:
            return (0 if "_L_" in name else 1) * 5 + i
    return -1


dom_digit = np.full(n, -1, dtype=np.int32)
for gi, gname in gidx.items():
    dd = digit_of(gname)
    if dd >= 0:
        dom_digit[dom == gi] = dd
print("\n=== FINGER BLEED (weight landing on a different digit) ===")
for side, S in (("L", 0), ("R", 1)):
    for i, d in enumerate(DIGITS):
        own = S * 5 + i
        names = [gi for gi, gname in gidx.items()
                 if f"_{S and 'R' or 'L'}_" in gname and f"_{d}" in gname]
        if not names:
            continue
        tot_w = 0.0
        bleed_w = 0.0
        for v in me.vertices:
            for g in v.groups:
                if g.group in names and g.weight > 1e-6:
                    tot_w += g.weight
                    if dom_digit[v.index] != own and dom_digit[v.index] >= 0:
                        bleed_w += g.weight
        pct = 100.0 * bleed_w / tot_w if tot_w else 0.0
        print(f"  {side}_{d:<6} weight {tot_w:9.1f}   onto other digits {pct:5.2f}%")

# The rest mesh being right proves nothing: the armature modifier is what the engine will run.
# Evaluate it and require the deformed body to sit exactly where the undeformed one does.
dg = bpy.context.evaluated_depsgraph_get()
evo = body.evaluated_get(dg)
tmp = evo.to_mesh()
Dv = np.empty(len(tmp.vertices) * 3)
tmp.vertices.foreach_get("co", Dv)
Dv = Dv.reshape(-1, 3)
evo.to_mesh_clear()
Mw = np.array(body.matrix_world)
Dv = Dv @ Mw[:3, :3].T + Mw[:3, 3]
drift = np.linalg.norm(Dv - V_ref, axis=1).max() * UNITM * 1000
print(f"\n  DEFORMED at rest pose: z {Dv[:,2].min():.5f}..{Dv[:,2].max():.5f}  "
      f"max drift vs undeformed {drift:.6f} mm")
assert drift < 0.01, (f"the armature displaces the body by {drift:.3f} mm at rest — a leftover pose, "
                      f"not a weighting problem")

bpy.ops.wm.save_as_mainfile(filepath=OUTB)
log(f"SAVED {OUTB}")
bpy.ops.object.select_all(action='DESELECT')
arm.select_set(True)
body.select_set(True)
bpy.context.view_layer.objects.active = arm
bpy.ops.export_scene.gltf(filepath=OUTG, export_format='GLB', use_selection=True,
                          export_image_format='AUTO', export_jpeg_quality=95,
                          export_yup=True, export_apply=False, export_skins=True)
log(f"EXPORTED {OUTG}  ({os.path.getsize(OUTG)/1e6:.2f} MB)")
print("GRAFT_DONE")