# Stage 52 (v02): transplant the PRISTINE ORIGINAL head onto the decimated nude body.
#
#   blender --background --python 52_head_transplant.py -- <body.blend> <out.blend>
#
# Why: bisecting the shard-eye defect (head renders per stage) proved it is in the hires MASTER,
# not in any v02 step — the early heal chain ran whole-mesh welds and reset custom split normals,
# and the eye/lash shells only read as an eye through Tripo's authored normals. Every descendant
# inherits it, including the currently shipped game body. The one place the eyes are right is the
# untouched original, so the head comes from there — geometry, normals, UVs and all.
#
# Why a cut works HERE and not at the underwear: the bra had no skin beneath it (ray census,
# _Bare.glb's black cavity). The throat is ordinary skin on both sides of a plane, so both meshes
# take a clean planar bisect and the two rings bridge.
#
# The two lineages keep their own materials: body faces sample the master's maps (with the nude
# fill), head faces sample the ORIGINAL maps — the transfer stage resolves images per material.
# Bridge faces are new geometry with no source UVs; each takes a single nearby body-ring UV
# (all three corners the same texel), i.e. a flat splat of throat skin, invisible on a uniform
# throat and never smearing across two different atlases.
import bpy, bmesh, sys, os, time
import numpy as np
from mathutils import Vector

argv = sys.argv[sys.argv.index("--") + 1:]
BODY_BLEND, OUT = argv[0], argv[1]
ORIG_GLB = r"C:\Users\Jeremy\tinqs\animation\characters\originals\female\female_lena_tripo.glb"
# TX_ZCUT env overrides the plane. The default 0.775 is what the SHIPPED hybrid used and must
# keep reproducing it — but that plane clips the T-POSE ARMS, giving five boundary rings, and at
# full density bridge_loops paired the wrong ones (a geometry shelf flared out of the shoulders).
# For full-res builds use 0.80: pure throat, exactly one ring per side.
Z_CUT = float(os.environ.get("TX_ZCUT", "0.775"))
t0 = time.time()


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


def bisect_keep(ob, keep_above):
    bpy.ops.object.select_all(action='DESELECT')
    ob.select_set(True)
    bpy.context.view_layer.objects.active = ob
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.select_all(action='SELECT')
    # plane normal is +z, so OUTER = above the cut. clear_outer removes the ABOVE side —
    # the first run had these inverted and stitched the shard head onto the clothed body.
    bpy.ops.mesh.bisect(plane_co=(0, 0, Z_CUT), plane_no=(0, 0, 1),
                        clear_inner=keep_above, clear_outer=not keep_above, use_fill=False)
    bpy.ops.object.mode_set(mode='OBJECT')


bpy.ops.wm.open_mainfile(filepath=BODY_BLEND)
body = max([o for o in bpy.data.objects if o.type == 'MESH'], key=lambda o: len(o.data.vertices))
log(f"body in: {len(body.data.vertices)} v")
bisect_keep(body, keep_above=False)
log(f"body below z={Z_CUT}: {len(body.data.vertices)} v")
# give the body a custom-normals layer BEFORE joining, or join drops the head's authored ones
body.data.shade_smooth()
try:
    ln = np.empty(len(body.data.loops) * 3, dtype=np.float32)
    body.data.corner_normals.foreach_get("vector", ln)
    body.data.normals_split_custom_set(ln.reshape(-1, 3))
except Exception as e:
    log(f"  custom-normal seed on body: {e}")

before = {o.name for o in bpy.data.objects}
bpy.ops.import_scene.gltf(filepath=ORIG_GLB)
new = [o for o in bpy.data.objects if o.name not in before]
# drop importer widgets (glTF_not_exported) and non-meshes
for o in list(new):
    if o.type != 'MESH' or any(c.name.startswith("glTF_not_exported") for c in o.users_collection):
        if o.type == 'MESH' or o.type == 'ARMATURE' or o.type == 'EMPTY':
            bpy.data.objects.remove(o, do_unlink=True)
            new.remove(o)
head = max([o for o in new if o.type == 'MESH'], key=lambda o: len(o.data.vertices))
bpy.ops.object.select_all(action='DESELECT')
head.select_set(True)
bpy.context.view_layer.objects.active = head
bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
log(f"original in: {len(head.data.vertices)} v")
bisect_keep(head, keep_above=True)
log(f"original above z={Z_CUT}: {len(head.data.vertices)} v")

# BOTH sides must share one UV layer name or join makes two half-empty layers and the head
# samples texel (0,0) — measured: the whole head rendered as clay. The body blend's layer is
# 'UVMap' (34_v04 lineage); rename both to UVMap_tripo so join merges them into one.
for l in head.data.uv_layers:
    l.name = "UVMap_tripo"
for l in body.data.uv_layers:
    l.name = "UVMap_tripo"

# remember a body-side ring UV for the bridge splat
bm = bmesh.new(); bm.from_mesh(body.data)
uvl = bm.loops.layers.uv.get("UVMap_tripo") or bm.loops.layers.uv.active
ring_uv = None
for v in bm.verts:
    if abs(v.co.z - Z_CUT) < 1e-4 and v.link_loops:
        # take the uv of a loop one step AWAY from the ring (interior throat texel)
        for l_ in v.link_loops:
            ring_uv = tuple(l_[uvl].uv)
            break
        if ring_uv:
            break
bm.free()
log(f"bridge splat uv: {ring_uv}")

# join: body is the target so its material stays slot 0; head material appends as slot 1
n_faces_before_join = len(body.data.polygons)
bpy.ops.object.select_all(action='DESELECT')
body.select_set(True)
head.select_set(True)
bpy.context.view_layer.objects.active = body
bpy.ops.object.join()
me = body.data
log(f"joined: {len(me.vertices)} v, {len(me.polygons)} f, materials {[m.name for m in me.materials]}")

# bridge the two rings
n_faces_before_bridge = len(me.polygons)
bm = bmesh.new(); bm.from_mesh(me)
bm.edges.ensure_lookup_table()
ring_edges = [e for e in bm.edges
              if len(e.link_faces) == 1
              and abs(e.verts[0].co.z - Z_CUT) < 1e-4 and abs(e.verts[1].co.z - Z_CUT) < 1e-4]
# how many separate rings? more than 2 means the plane clipped limbs and the bridge will pair wrong
_par = {}
def _find(a):
    while _par.get(a, a) != a:
        _par[a] = _par.get(_par[a], _par[a]); a = _par[a]
    return a
for e in ring_edges:
    a, b_ = _find(e.verts[0].index), _find(e.verts[1].index)
    if a != b_: _par[a] = b_
_roots = {_find(e.verts[0].index) for e in ring_edges}
log(f"ring boundary edges: {len(ring_edges)} in {len(_roots)} loops"
    + ("  <-- WARNING: >2 loops, plane clips limbs" if len(_roots) > 2 else ""))
try:
    res = bmesh.ops.bridge_loops(bm, edges=ring_edges)
    new_faces = res.get("faces", [])
except Exception as e:
    log(f"bridge_loops failed: {e}")
    new_faces = []
log(f"bridge created {len(new_faces)} faces")
uvl = bm.loops.layers.uv.get("UVMap_tripo") or bm.loops.layers.uv.active
for f in new_faces:
    f.material_index = 0            # body material: throat skin
    f.smooth = True
    for l_ in f.loops:
        l_[uvl].uv = ring_uv        # flat splat of one throat texel
bm.to_mesh(me)
bm.free()
me.update()

# One material for the whole mesh. The transfer stage resamples every map into ONE new atlas,
# so a second material would leave head faces pointing at untransferred images with replaced
# UVs. Sampling the head from the MASTER's maps is exact: every texture pass in the lane masked
# the head out, so its texels are the original's bit-for-bit.
mi = np.zeros(len(me.polygons), dtype=np.int32)
me.polygons.foreach_set("material_index", mi)
while len(me.materials) > 1:
    me.materials.pop(index=1)
log(f"unified material: {[m.name for m in me.materials]}")

# verify
co = np.empty(len(me.vertices) * 3); me.vertices.foreach_get("co", co); co = co.reshape(-1, 3)
bm = bmesh.new(); bm.from_mesh(me)
open_ring = [e for e in bm.edges if len(e.link_faces) == 1
             and abs(e.verts[0].co.z - Z_CUT) < 1e-3]
bm.free()
print(f"\nVERIFY: {len(me.vertices)} v / {len(me.polygons)} f")
print(f"  height {co[:,2].max()-co[:,2].min():.5f} (feet {co[:,2].min():+.5f})")
print(f"  open boundary edges left at the cut: {len(open_ring)}   (0 = ring fully bridged)")
print(f"  materials: {[m.name for m in me.materials]}")

bpy.ops.wm.save_as_mainfile(filepath=OUT)
log(f"WROTE {OUT}")
print("TRANSPLANT_DONE")