# Stage 53 (v02): rig the transplant body by transferring weights from the rigged v01.
#
#   blender --background --python 53_rig_transfer.py -- <reatlased.blend> <rig_v01.blend> <out.blend> <out.glb>
#
# The v02 topology is new (head transplant + re-decimation), so the index-exact graft that rigged
# v01 is off the table. Nearest-surface transfer is safe HERE because source and target are the
# SAME body in the SAME frame (v01 was decimated from the same master this body's torso came
# from): every target vertex sits on or microns from the source surface. The July finger-mangling
# happened transferring decimated->FULL-RES fingers; v02's hands are at v01's own density, and
# the full-res head takes trivially rigid weights (Head + face bones).
#
# Verification is the same standard as the graft: weights must sum to 1 with zero unweighted
# verts, and the DEFORMED rest pose must sit exactly on the undeformed mesh.
import bpy, sys, os, time
import numpy as np

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


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


bpy.ops.wm.open_mainfile(filepath=TARGET)
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 = np.empty(n * 3); me.vertices.foreach_get("co", V_ref); V_ref = V_ref.reshape(-1, 3)
UNITM = 1.777 / (V_ref[:, 2].max() - V_ref[:, 2].min())
log(f"target: {n} v")

with bpy.data.libraries.load(RIGSRC, link=False) as (src, dst):
    dst.objects = src.objects
donor, arm = None, None
for o in dst.objects:
    if o is None:
        continue
    bpy.context.scene.collection.objects.link(o)
    if o.type == 'MESH':
        donor = o
    if o.type == 'ARMATURE':
        arm = o
log(f"donor: {len(donor.data.vertices)} v, {len(donor.vertex_groups)} groups; "
    f"armature: {len(arm.data.bones)} bones")

# transfer weights nearest-surface
for vg in list(body.vertex_groups):
    body.vertex_groups.remove(vg)
bpy.context.view_layer.objects.active = body
dt = body.modifiers.new("wts", 'DATA_TRANSFER')
dt.object = donor
dt.use_vert_data = True
dt.data_types_verts = {'VGROUP_WEIGHTS'}
dt.vert_mapping = 'POLYINTERP_NEAREST'
dt.layers_vgroup_select_src = 'ALL'
bpy.ops.object.datalayout_transfer(modifier=dt.name)
bpy.ops.object.modifier_apply(modifier=dt.name)
log(f"weights transferred into {len(body.vertex_groups)} groups")
bpy.ops.object.vertex_group_limit_total(group_select_mode='ALL', limit=4)
bpy.ops.object.vertex_group_normalize_all(group_select_mode='ALL', lock_active=False)

# 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()
bpy.data.objects.remove(donor, do_unlink=True)

# verify
tot = np.zeros(n)
for v in me.vertices:
    tot[v.index] = sum(g.weight for g in v.groups)
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()
drift = np.linalg.norm(Dv - V_ref, axis=1).max() * UNITM * 1000
print("\n=== VERIFY ===")
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())}")
print(f"  DEFORMED at rest: z {Dv[:,2].min():.5f}..{Dv[:,2].max():.5f}  drift {drift:.4f} mm")
assert tot.min() > 0.99 and drift < 0.01, "rig transfer failed verification"

bpy.ops.wm.save_as_mainfile(filepath=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("RIG53_DONE")