"""
Spike T-1089 / Synty Sidekick intake — step 2: transplant.

Route decision (from 01_recon_rest_pose.py, out/recon_rest_pose.json):
  Route 1 (rename-and-rebind, keep Sidekick weights verbatim) REJECTED on
  measured evidence: spine_02/spine_03 joint heads sit 0.160/0.168 m lower in
  the Sidekick rig than ours, shoulders 0.082 m lower, clavicles 0.150 m off.
  Rest pose would render fine (inverse binds absorb it) but every spine/shoulder
  rotation would pivot the garment around joints up to 17 cm away from where the
  body underneath pivots -> swim/clip under animation.
  Route 2 (proven pipeline): affine-align garment onto our average_m body,
  then Data Transfer POLYINTERP_NEAREST weights from OUR body mesh, normalize,
  bind to OUR armature.

Alignment transform (documented, measured from bone landmarks):
  scale  XY by ours_shoulder_x / sk_shoulder_x
  scale  Z  by (ours_upperarm_z - ours_pelvis_z) / (sk_upperarm_z - sk_pelvis_z)
  then translate so the pelvis landmarks coincide.

Run:
    tooling/blender --background --python \
        spikes/synty-intake/scripts/02_transplant_garment.py
"""

import bpy
import json
import os
from mathutils import Matrix, Vector

REPO = "/var/mnt/data/projects/settled-reach"
GARMENT_FBX = os.path.join(
    REPO, "spikes/synty-intake/raw/Assets/Synty/SidekickCharacters/Resources",
    "Meshes/Outfits/Starter/SK_SCFI_CIVL_09_10TORS_HU01.fbx")
OUR_ARMATURE_GLB = os.path.join(REPO, "client/assets/characters/skeleton/armature.glb")
BODY_DIR = os.path.join(REPO, "client/assets/characters/bodies/average_m")
# Segments covered by / adjacent to the torso garment — the weight-transfer source.
BODY_SEGS = ["seg_torso", "seg_hips", "seg_neck",
             "seg_arm_upper_l", "seg_arm_upper_r",
             "seg_leg_upper_l", "seg_leg_upper_r"]

OUT_GLB = os.path.join(REPO, "spikes/synty-intake/out",
                       "SK_SCFI_CIVL_09_10TORS_HU01_quaternius.glb")
OUT_LOG = os.path.join(REPO, "spikes/synty-intake/out", "transplant_log.json")

log = {}


def import_new(op, path, **kw):
    before = set(bpy.data.objects)
    op(filepath=path, **kw)
    return list(set(bpy.data.objects) - before)


def find_armature(objs):
    for o in objs:
        if o.type == 'ARMATURE':
            return o
    return None


def bone_head_world(arm, name):
    b = arm.data.bones.get(name)
    return (arm.matrix_world @ b.head_local) if b else None


def deselect_all():
    bpy.ops.object.select_all(action='DESELECT')


def set_active(obj):
    bpy.context.view_layer.objects.active = obj


# --- fresh scene ---
bpy.ops.object.select_all(action='SELECT')
bpy.ops.object.delete(use_global=False)

# --- 1. import garment ---
g_objs = import_new(bpy.ops.import_scene.fbx, GARMENT_FBX)
g_arm = find_armature(g_objs)
garment = next(o for o in g_objs if o.type == 'MESH')
log["garment"] = {"verts": len(garment.data.vertices),
                  "sidekick_vgroups": [vg.name for vg in garment.vertex_groups]}

# landmarks in Sidekick space (world)
sk_pelvis = bone_head_world(g_arm, "pelvis").copy()
sk_upperarm = bone_head_world(g_arm, "upperarm_l").copy()

# --- 2. import our armature ---
o_objs = import_new(bpy.ops.import_scene.gltf, OUR_ARMATURE_GLB)
our_arm = find_armature(o_objs)
our_pelvis = bone_head_world(our_arm, "pelvis").copy()
our_upperarm = bone_head_world(our_arm, "upperarm_l").copy()

# --- 3. detach garment from Sidekick armature, freeze world coords ---
deselect_all()
set_active(garment)
garment.select_set(True)
mw = garment.matrix_world.copy()
garment.parent = None
garment.matrix_world = mw
for mod in list(garment.modifiers):
    garment.modifiers.remove(mod)
# drop Sidekick body-blend shape keys (not part of our body-type system)
if garment.data.shape_keys:
    set_active(garment)
    bpy.ops.object.shape_key_remove(all=True)
# bake object transform into mesh data so object space == world space
bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)

# --- 4. affine alignment garment -> our proportions ---
sxy = our_upperarm.x / sk_upperarm.x
sz = (our_upperarm.z - our_pelvis.z) / (sk_upperarm.z - sk_pelvis.z)
S = Matrix.Diagonal(Vector((sxy, sxy, sz, 1.0)))
scaled_pelvis = Vector((sk_pelvis.x * sxy, sk_pelvis.y * sxy, sk_pelvis.z * sz))
t = our_pelvis - scaled_pelvis
M = Matrix.Translation(t) @ S
garment.data.transform(M)
garment.data.update()

log["alignment"] = {
    "scale_xy": round(sxy, 5), "scale_z": round(sz, 5),
    "translate": [round(v, 5) for v in t],
    "sk_pelvis": [round(v, 5) for v in sk_pelvis],
    "our_pelvis": [round(v, 5) for v in our_pelvis],
    "sk_upperarm_l": [round(v, 5) for v in sk_upperarm],
    "our_upperarm_l": [round(v, 5) for v in our_upperarm],
}

# residual landmark error after affine, for all 12 bound bones
residuals = {}
for bn in ["pelvis", "thigh_l", "thigh_r", "spine_01", "spine_02", "spine_03",
           "neck_01", "head", "clavicle_l", "upperarm_l", "clavicle_r",
           "upperarm_r"]:
    sk = bone_head_world(g_arm, bn)
    ours = bone_head_world(our_arm, bn if bn != "head" else "Head")
    sk_t = M @ sk
    d = sk_t - ours
    residuals[bn] = {"residual_m": round(d.length, 5),
                     "delta": [round(v, 5) for v in d]}
log["post_affine_bone_residuals"] = residuals

# --- 5. delete Sidekick armature ---
deselect_all()
g_arm.select_set(True)
set_active(g_arm)
bpy.ops.object.delete(use_global=False)

# --- 6. import + join our body segments as the weight-transfer source ---
seg_meshes = []
for seg in BODY_SEGS:
    objs = import_new(bpy.ops.import_scene.gltf,
                      os.path.join(BODY_DIR, seg + ".glb"))
    picked = None
    for o in objs:
        if o.type == 'MESH' and len(o.vertex_groups) > 0:
            picked = o
        elif o.type == 'MESH':
            # debris (bounds icosphere etc.) — delete
            deselect_all()
            o.select_set(True)
            set_active(o)
            bpy.ops.object.delete(use_global=False)
    # delete the segment's own armature copy
    arm = find_armature(objs)
    if arm and picked is not None:
        # unparent mesh first, keep transform
        pmw = picked.matrix_world.copy()
        picked.parent = None
        picked.matrix_world = pmw
    if arm:
        deselect_all()
        arm.select_set(True)
        set_active(arm)
        bpy.ops.object.delete(use_global=False)
    if picked is None:
        print(f"WARNING: no weighted mesh in {seg}")
        continue
    for mod in list(picked.modifiers):
        picked.modifiers.remove(mod)
    seg_meshes.append(picked)

log["weight_source_segments"] = [
    {"name": m.name, "verts": len(m.data.vertices)} for m in seg_meshes]

deselect_all()
for m in seg_meshes:
    m.select_set(True)
set_active(seg_meshes[0])
if len(seg_meshes) > 1:
    bpy.ops.object.join()
body = bpy.context.view_layer.objects.active
body.name = "weight_source_body"
bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)

# --- 7. weight transfer: our body -> garment ---
# wipe Sidekick vertex groups first so only OUR weights remain
for vg in list(garment.vertex_groups):
    garment.vertex_groups.remove(vg)

deselect_all()
set_active(garment)
garment.select_set(True)

dt = garment.modifiers.new(name="WeightTransfer", type='DATA_TRANSFER')
dt.object = body
dt.use_vert_data = True
dt.data_types_verts = {'VGROUP_WEIGHTS'}
dt.vert_mapping = 'POLYINTERP_NEAREST'
dt.layers_vgroup_select_src = 'ALL'
dt.layers_vgroup_select_dst = 'NAME'
bpy.ops.object.datalayout_transfer(modifier=dt.name)
bpy.ops.object.modifier_apply(modifier=dt.name)

# normalize
bpy.ops.object.mode_set(mode='WEIGHT_PAINT')
bpy.ops.object.vertex_group_normalize_all(lock_active=False)
bpy.ops.object.mode_set(mode='OBJECT')

log["transferred_vgroups"] = [vg.name for vg in garment.vertex_groups]

# sanity: count verts with (near-)zero total weight
zero = 0
for v in garment.data.vertices:
    tot = sum(g.weight for g in v.groups)
    if tot < 1e-4:
        zero += 1
log["zero_weight_verts"] = zero

# --- 8. bind garment to OUR armature ---
garment.parent = our_arm
garment.matrix_parent_inverse = our_arm.matrix_world.inverted()
am = garment.modifiers.new(name="Armature", type='ARMATURE')
am.object = our_arm

# --- 9. delete the body source, export armature + garment ---
deselect_all()
body.select_set(True)
set_active(body)
bpy.ops.object.delete(use_global=False)

os.makedirs(os.path.dirname(OUT_GLB), exist_ok=True)
deselect_all()
our_arm.select_set(True)
garment.select_set(True)
set_active(our_arm)
bpy.ops.export_scene.gltf(
    filepath=OUT_GLB,
    export_format='GLB',
    use_selection=True,
    export_apply=False,
    export_animations=False,
    export_skins=True,
    export_morph=False,
    export_extras=False,
)
log["export"] = {"path": OUT_GLB, "size_kb": os.path.getsize(OUT_GLB) // 1024}

with open(OUT_LOG, "w") as f:
    json.dump(log, f, indent=2)

print("\n=== TRANSPLANT SUMMARY ===")
print(json.dumps(log, indent=2))