# lena_leafbikini lane, stage 05: CUT the leaves out at their seam. Leave the holes open.
#
#   blender --background --factory-startup --python 05_cut_leaves.py -- \
#           <pristine.glb> <leaf_mask.npz> <out.glb> [--blend <out.blend>]
#
# Deliberately NOT healed. Jeremy asked for the model with holes so the seam itself can be
# judged before any repair is designed, and that is the whole deliverable of this stage: the
# leaf shells are gone, the skin is untouched, and the rim of each hole is exactly the line
# where leaf stopped and Lena started. Filling them is a separate decision (and a harder one —
# see the crotch/gusset history in characters/work/lena/06*.py: a membrane over a wide footprint
# flattens the anatomy it spans, which is why the nude lane only ever faired a narrow rim band).
#
# THE CUT RULE: a face dies only if ALL of its vertices are masked. Combined with the mask's
# 2-ring outward grow (stage 04), the surviving rim therefore sits ~1 ring OUTSIDE the green,
# i.e. just onto the skin. That asymmetry is on purpose in both directions:
#   * "any vertex masked" would erode a ring further into her skin and leave a ragged, spiky
#     boundary — single triangles hanging off the rim wherever the mask edge zig-zags;
#   * a cut one ring short leaves a rim of leaf root standing proud of the skin, which is the
#     one outcome that would make this file useless for judging the seam.
#
# Nothing is smoothed, welded, re-normalled or re-textured here. The material, UVs and all three
# 4096^2 maps ride through untouched, so what changes between input and output is exactly "some
# faces are missing" — asserted below on the surviving vertex positions, not assumed.
import bpy, sys, os, time, argparse
import numpy as np

argv = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
ap = argparse.ArgumentParser()
ap.add_argument("glb")
ap.add_argument("mask")
ap.add_argument("out")
ap.add_argument("--blend", default="")
ap.add_argument("--despike", type=int, default=6,
                help="passes of dangling-face erosion at the new rims (0 = off)")
ap.add_argument("--min-island", type=int, default=400,
                help="drop surviving shells smaller than this many welded points")
A = ap.parse_args(argv)
GLB, MASK, OUT = os.path.abspath(A.glb), os.path.abspath(A.mask), os.path.abspath(A.out)
os.makedirs(os.path.dirname(OUT), exist_ok=True)
t0 = time.time()


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


def weld_ids(me):
    """Position-welded vertex ids. EVERY topology question below has to be asked on these: the
    glTF importer splits each UV seam into separate Blender vertices, so on the raw mesh a face
    sitting on a texture-chart border looks exactly like a face on a hole rim. Eroding "dangling"
    faces off the split mesh would chew Lena open along her UV seams."""
    n = len(me.vertices)
    co = np.empty(n * 3); me.vertices.foreach_get("co", co)
    _, inv = np.unique(np.round(co.reshape(-1, 3), 6), axis=0, return_inverse=True)
    return inv.astype(np.int64), int(inv.max()) + 1


def face_edge_table(me, wid):
    """Per-loop: the welded edge it starts, and how many faces share that edge."""
    nf, nl = len(me.polygons), len(me.loops)
    lv = np.empty(nl, dtype=np.int32); me.loops.foreach_get("vertex_index", lv)
    ls = np.empty(nf, dtype=np.int32); me.polygons.foreach_get("loop_start", ls)
    lt = np.empty(nf, dtype=np.int32); me.polygons.foreach_get("loop_total", lt)
    fol = np.repeat(np.arange(nf), lt)
    nxt = ls[fol] + ((np.arange(nl) - ls[fol] + 1) % lt[fol])
    a, b = wid[lv], wid[lv[nxt]]
    key = np.stack([np.minimum(a, b), np.maximum(a, b)], axis=1)
    _, ei, ec = np.unique(key, axis=0, return_inverse=True, return_counts=True)
    return ls.astype(np.int64), lt, ec[ei]


def kill_faces(me, kill):
    """Delete the flagged faces plus anything left orphaned by them."""
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.select_all(action='DESELECT')
    bpy.ops.object.mode_set(mode='OBJECT')
    me.polygons.foreach_set("select", kill)
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.select_mode(type='FACE')
    bpy.ops.mesh.delete(type='FACE')
    bpy.ops.object.mode_set(mode='OBJECT')


bpy.ops.wm.read_factory_settings(use_empty=True)
bpy.ops.import_scene.gltf(filepath=GLB)
body = max([o for o in bpy.data.objects if o.type == 'MESH'], key=lambda o: len(o.data.vertices))
me = body.data
n0, f0 = len(me.vertices), len(me.polygons)
log(f"in : '{body.name}' {n0}v {f0}f  mats={[m.name for m in me.materials if m]}")

co = np.empty(n0 * 3); me.vertices.foreach_get("co", co)
P0 = co.reshape(-1, 3).copy()

z = np.load(MASK)
inv, m = z["inv"].astype(np.int64), z["mask"]
if len(inv) != n0:
    raise SystemExit(f"[cut] FATAL: mask was built for {len(inv)} verts, this GLB has {n0}. "
                     f"Re-run 04_leaf_mask.py against {os.path.basename(GLB)}.")
vm = m[inv]
log(f"mask: {vm.sum()} of {n0} verts ({100*vm.sum()/n0:.2f}%)")

# ── faces whose every vertex is masked ──────────────────────────────────────────────────────
nl = len(me.loops)
lv = np.empty(nl, dtype=np.int32); me.loops.foreach_get("vertex_index", lv)
ls = np.empty(f0, dtype=np.int32); me.polygons.foreach_get("loop_start", ls)
lt = np.empty(f0, dtype=np.int32); me.polygons.foreach_get("loop_total", lt)
hits = np.add.reduceat(vm[lv].astype(np.int32), ls.astype(np.int64))
kill = hits == lt
log(f"faces: {kill.sum()} fully masked, {(hits > 0).sum() - kill.sum()} straddle the rim (kept)")

# ── delete them ─────────────────────────────────────────────────────────────────────────────
bpy.context.view_layer.objects.active = body
body.select_set(True)
kill_faces(me, kill)
log(f"cut: {len(me.vertices)}v {len(me.polygons)}f")

# ── clean the new rims ──────────────────────────────────────────────────────────────────────
# Two passes, both of which can only ever remove things that are already not part of a smooth
# surface. Neither is cosmetic: the colour key loses the deep shadow pockets UNDER overlapping
# leaves (a leaf underside in shadow reads near-black, so its hue is noise), and what survives
# there is a leaf fragment either hanging off the rim by a single edge or floating free.
#
# A. DANGLING-FACE EROSION. A face with two or three of its edges on a real (welded) boundary is
#    a spike: it is joined to the surface along at most one edge. On a closed surface no such
#    face exists, so this cannot bite into her skin — it can only walk back the ragged tongues
#    left where the mask edge zig-zagged. Iterated, because removing a spike can expose the next.
for i in range(A.despike):
    wid, _ = weld_ids(me)
    ls, lt, ecount = face_edge_table(me, wid)
    nb = np.add.reduceat((ecount == 1).astype(np.int32), ls)
    spike = nb >= 2
    if not spike.any():
        log(f"despike pass {i+1}: none left")
        break
    kill_faces(me, spike)
    log(f"despike pass {i+1}: {int(spike.sum())} dangling faces -> {len(me.polygons)}f")

# B. DETACHED SHELLS. Whatever is no longer connected to the body is a leaf that came away whole.
if A.min_island > 0:
    wid, ng = weld_ids(me)
    nl = len(me.loops)
    lv = np.empty(nl, dtype=np.int32); me.loops.foreach_get("vertex_index", lv)
    ev = np.empty(len(me.edges) * 2, dtype=np.int32); me.edges.foreach_get("vertices", ev)
    ea, eb = wid[ev[0::2]], wid[ev[1::2]]
    k = ea != eb
    src = np.concatenate([ea[k], eb[k]]); dst = np.concatenate([eb[k], ea[k]])
    o = np.argsort(src, kind='stable'); src, dst = src[o], dst[o]
    ptr = np.concatenate([[0], np.cumsum(np.bincount(src, minlength=ng))])
    from collections import deque
    lab = np.full(ng, -1, dtype=np.int64); sizes = []
    for s in range(ng):
        if lab[s] >= 0:
            continue
        cid = len(sizes); q = deque([s]); lab[s] = cid; sz = 0
        while q:
            c = q.popleft(); sz += 1
            for j in range(ptr[c], ptr[c + 1]):
                if lab[dst[j]] < 0:
                    lab[dst[j]] = cid; q.append(dst[j])
        sizes.append(sz)
    sizes = np.array(sizes)
    log(f"shells: {len(sizes)}, sizes {sorted(sizes)[::-1][:8]}")
    drop = np.isin(lab, np.nonzero(sizes < A.min_island)[0])
    if drop.any():
        vd = drop[wid]
        ls = np.empty(len(me.polygons), dtype=np.int32); me.polygons.foreach_get("loop_start", ls)
        lt = np.empty(len(me.polygons), dtype=np.int32); me.polygons.foreach_get("loop_total", lt)
        hit = np.add.reduceat(vd[lv].astype(np.int32), ls.astype(np.int64)) == lt
        kill_faces(me, hit)
        log(f"dropped {int((sizes < A.min_island).sum())} detached shells "
            f"({int(drop.sum())} points, {int(hit.sum())} faces) -> {len(me.polygons)}f")

n1, f1 = len(me.vertices), len(me.polygons)
log(f"out: {n1}v ({n0-n1} removed, {100*(n0-n1)/n0:.1f}%) "
    f"{f1}f ({f0-f1} removed, {100*(f0-f1)/f0:.1f}%)")

# ── gates ───────────────────────────────────────────────────────────────────────────────────
# 1. NOTHING MOVED. Every surviving vertex position must appear in the input — this stage is a
#    deletion and only a deletion, so any displacement at all is a bug, not a tolerance.
co = np.empty(n1 * 3); me.vertices.foreach_get("co", co)
P1 = co.reshape(-1, 3)
Q = np.round(np.concatenate([P0, P1]) * 1e6).astype(np.int64)
_, iq = np.unique(Q, axis=0, return_inverse=True)
i0, i1 = iq[:n0], iq[n0:]
stray = int((~np.isin(i1, i0)).sum())
log(f"gate positions: {stray} survivors not present in the input ({'PASS' if not stray else 'FAIL'})")
if stray:
    raise SystemExit("[cut] FATAL: the cut moved geometry; it must only remove it")

# 2. the survivors must be one connected shell plus nothing else, and the leaf zone must now be
#    open: count boundary edges before/after.
import bmesh
bm = bmesh.new(); bm.from_mesh(me)
bnd = sum(1 for e in bm.edges if len(e.link_faces) == 1)
loose = sum(1 for v in bm.verts if not v.link_faces)
log(f"gate topology: {bnd} boundary edges, {loose} loose verts")
if loose:
    bmesh.ops.delete(bm, geom=[v for v in bm.verts if not v.link_faces], context='VERTS')
    bm.to_mesh(me)
    log(f"  removed {loose} loose verts -> {len(me.vertices)}v")
bm.free()

# 3. the holes must be where the leaves were and nowhere else
me.update()
H = P0[:, 2].max() - P0[:, 2].min()
zf = (P0[vm, 2] - P0[:, 2].min()) / H
log(f"gate extent: cut geometry spanned z {zf.min():.3f}..{zf.max():.3f} of body height "
    f"(bust + briefs bands only)")

# ── export ──────────────────────────────────────────────────────────────────────────────────
if A.blend:
    bpy.ops.wm.save_as_mainfile(filepath=os.path.abspath(A.blend))
    log(f"WROTE {A.blend}")
for o in bpy.data.objects:
    o.select_set(True)
bpy.ops.export_scene.gltf(filepath=OUT, export_format='GLB', use_selection=True,
                          export_yup=True, export_skins=False, export_animations=False,
                          export_apply=False, export_image_format='AUTO',
                          export_tangents=False, export_normals=True)
log(f"WROTE {OUT} ({os.path.getsize(OUT)/1e6:.2f} MB)")