# lena_leafbikini lane, stage 04: build the LEAF MASK and prove it before anything is cut.
#
#   blender --background --factory-startup --python 04_leaf_mask.py -- <pristine.glb> <renderdir>
#            [--hue-lo 52] [--hue-hi 170] [--sat-min 0.10] [--val-max 0.90] [--z-max 0.85]
#            [--min-comp 200] [--close 6] [--grow 2] [--no-render]
#
# The mask lands at leaf_mask.npz beside this script, NOT in <renderdir>: renderdir is a
# review/ dir and those are gitignored scratch, while the mask is a KEEP-tier lane input
# (.agents/rules/working-files.md) that stage 05 consumes.
#
# WHAT THE LEAVES ACTUALLY ARE (stage 02 probe + stage 03 clay renders):
#   * REAL GEOMETRY, not paint. The clay render shows every leaf, curled tip and hip vine with
#     all materials stripped — so this is the OPPOSITE case to Lena's game-body underwear, which
#     was painted onto the skin (tools/make_lena_nude_body.py, finding 1). There the fix was to
#     re-fair a rim crease; here there is a solid shell to delete.
#   * ONE WELDED SHELL with the body. Tripo emitted a single 1,029,360 v / 1,992,503 f mesh, so
#     the leaves are not a separable object, material or UV island — they are a bulge in the body
#     surface. Nothing can be selected "by object"; the seam has to be found.
#   * The seam is a SHARP RIM. Where a leaf meets skin the surface folds back on itself, so the
#     boundary is simultaneously a colour edge (green -> beige) and a crease.
#
# WHY THE KEY IS HUE, NOT "GREENNESS". Stage 02 keyed on G - max(R,B) and topped out at 0.20:
# Tripo painted these leaves a dark, desaturated olive (sRGB ~0.27/0.35/0.20), so the channel
# gap is only ~0.05-0.09 and any threshold that catches the leaf also catches shadow noise on
# skin. Hue separates them completely instead of marginally — skin sits at ~20-30 deg (orange),
# leaf at ~80-110 deg (green) — and hue is invariant to exactly the thing that ruins the channel
# gap here, which is how dark the pixel is.
#
# The mask is then repaired ON THE MESH, not in texture space: small components dropped (JPEG
# speckle), holes closed (leaf highlights that blow out to near-white lose their hue), and grown
# by one ring so the cut lands just outside the rim rather than just inside it. Erring outward is
# deliberate: a hole one ring too big is invisible, a leftover leaf stub is not.
import bpy, sys, os, time, argparse
import numpy as np
from collections import deque

argv = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
ap = argparse.ArgumentParser()
ap.add_argument("glb")
ap.add_argument("outdir")
ap.add_argument("--hue-lo", type=float, default=52.0)
ap.add_argument("--hue-hi", type=float, default=170.0)
ap.add_argument("--sat-min", type=float, default=0.10)
ap.add_argument("--val-max", type=float, default=0.90,
                help="hue alone separates olive leaf from beige skin; this only rejects pixels "
                     "so blown out that their hue is noise")
ap.add_argument("--z-max", type=float, default=0.85,
                help="fraction of body height above which the key is ignored — her irises and "
                     "eyebrows are green too, and they are not leaves")
ap.add_argument("--min-comp", type=int, default=200)
ap.add_argument("--close", type=int, default=6, help="ring radius of the morphological close")
ap.add_argument("--grow", type=int, default=2, help="final dilation, in rings")
ap.add_argument("--no-render", action="store_true")
A = ap.parse_args(argv)
GLB, OUT = os.path.abspath(A.glb), os.path.abspath(A.outdir)
os.makedirs(OUT, exist_ok=True)
t0 = time.time()


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


# =============================================================================================
# load + welded adjacency
# =============================================================================================
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
n = len(me.vertices)
log(f"'{body.name}' {n}v {len(me.polygons)}f")

co = np.empty(n * 3); me.vertices.foreach_get("co", co); co = co.reshape(-1, 3)
M = np.array(body.matrix_world)
W = co @ M[:3, :3].T + M[:3, 3]
H = W[:, 2].max() - W[:, 2].min()
MM = 1000.0 * 1.777 / H                       # units -> mm at final 1.777 m body scale

# The importer splits every UV seam into separate Blender vertices, so mesh adjacency is
# shattered along seams (nude-body finding 5). Every ring operation below runs on the
# POSITION-WELDED graph or it would leak holes along the seams.
_, inv = np.unique(np.round(W, 6), axis=0, return_inverse=True)
inv = inv.astype(np.int64)
ng = int(inv.max()) + 1
ev = np.empty(len(me.edges) * 2, dtype=np.int32); me.edges.foreach_get("vertices", ev)
ea, eb = inv[ev[0::2]], inv[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]
cnt = np.bincount(src, minlength=ng)
ptr = np.concatenate([[0], np.cumsum(cnt)])
log(f"welded {n} -> {ng} points, {len(dst)//2} undirected edges")


def ring(m, k=1):
    """Dilate a boolean over welded adjacency by k rings."""
    out = m.copy()
    for _ in range(k):
        hit = np.add.reduceat(out[dst].astype(np.int32), ptr[:-1]) > 0
        hit[cnt == 0] = False
        out = out | hit
    return out


def shrink(m, k=1):
    return ~ring(~m, k)


def components(m):
    """Connected components of a boolean over welded adjacency (exact BFS; the mask is small)."""
    lab = np.full(ng, -1, dtype=np.int64)
    comps = []
    for s in np.nonzero(m)[0]:
        if lab[s] >= 0:
            continue
        cid = len(comps)
        q = deque([s]); lab[s] = cid; size = 0
        while q:
            c = q.popleft(); size += 1
            for j in range(ptr[c], ptr[c + 1]):
                nb = dst[j]
                if m[nb] and lab[nb] < 0:
                    lab[nb] = cid; q.append(nb)
        comps.append(size)
    return lab, np.array(comps)


# =============================================================================================
# albedo -> HSV, per welded point
# =============================================================================================
base = None
for mat in [m for m in me.materials if m]:
    bsdf = next((x for x in mat.node_tree.nodes if x.type == 'BSDF_PRINCIPLED'), None)
    lnk = bsdf and bsdf.inputs["Base Color"].links
    if lnk:
        nd = lnk[0].from_node
        while nd.type != 'TEX_IMAGE':
            up = [i for i in nd.inputs if i.links]
            if not up:
                break
            nd = up[0].links[0].from_node
        if nd.type == 'TEX_IMAGE':
            base = nd.image
if base is None:
    raise SystemExit("[mask] FATAL: no base-colour image")

nl = len(me.loops)
lv = np.empty(nl, dtype=np.int32); me.loops.foreach_get("vertex_index", lv)
uvb = np.empty(nl * 2); me.uv_layers.active.data.foreach_get("uv", uvb); uvb = uvb.reshape(-1, 2)
vuv = np.zeros((n, 2))
vuv[lv[::-1]] = uvb[::-1]                     # first loop of each vertex wins
w, h = base.size
buf = np.empty(w * h * 4, dtype=np.float32); base.pixels.foreach_get(buf)
px = buf.reshape(h, w, 4)[:, :, :3].copy(); del buf
# bpy-imported UVs are already v-flipped by the importer (memory: gltf-uv-flip-vs-blender-images)
xi = np.clip((vuv[:, 0] * (w - 1)).astype(np.int32), 0, w - 1)
yi = np.clip((vuv[:, 1] * (h - 1)).astype(np.int32), 0, h - 1)
C = px[yi, xi].astype(np.float64); del px
S = np.clip(np.where(C <= 0.0031308, C * 12.92, 1.055 * np.maximum(C, 0) ** (1 / 2.4) - 0.055), 0, 1)

R, G, B = S[:, 0], S[:, 1], S[:, 2]
mx = S.max(1); mn = S.min(1); d = mx - mn
hue = np.zeros(n)
nz = d > 1e-6
im = np.argmax(S, axis=1)
sel = nz & (im == 0); hue[sel] = 60 * (((G[sel] - B[sel]) / d[sel]) % 6)
sel = nz & (im == 1); hue[sel] = 60 * ((B[sel] - R[sel]) / d[sel] + 2)
sel = nz & (im == 2); hue[sel] = 60 * ((R[sel] - G[sel]) / d[sel] + 4)
sat = np.where(mx > 1e-6, d / np.maximum(mx, 1e-6), 0.0)
log("hue histogram (all verts, 20 deg bins):")
hh, _ = np.histogram(hue, bins=18, range=(0, 360))
log("  " + "  ".join(f"{i*20:3d}:{c*100.0/n:5.2f}%" for i, c in enumerate(hh) if c))

zf = (W[:, 2] - W[:, 2].min()) / H
raw = (hue >= A.hue_lo) & (hue <= A.hue_hi) & (sat >= A.sat_min) & (mx <= A.val_max)
log(f"hue key [{A.hue_lo},{A.hue_hi}] sat>={A.sat_min} val<={A.val_max}: {raw.sum()} verts "
    f"({100*raw.sum()/n:.2f}%)  z {zf[raw].min():.3f}..{zf[raw].max():.3f} of height")
head = raw & (zf > A.z_max)
raw &= ~head
log(f"head gate z<={A.z_max}: dropped {head.sum()} verts (irises/eyebrows)")

# per-welded-point: a point is leaf if ANY of its seam copies keyed (a seam copy can sample the
# far side of a texture chart boundary)
m = np.zeros(ng, dtype=bool)
np.logical_or.at(m, inv, raw)

# =============================================================================================
# repair the mask on the mesh
# =============================================================================================
lab, sizes = components(m)
if len(sizes):
    log(f"components: {len(sizes)}, largest {sorted(sizes)[-8:]}")
    m &= np.isin(lab, np.nonzero(sizes >= A.min_comp)[0])
    log(f"  after min-comp {A.min_comp}: {m.sum()} points in {(sizes>=A.min_comp).sum()} comps")

if A.close:
    m = shrink(ring(m, A.close), A.close)     # close: fill specular blowouts inside a leaf
    log(f"after close({A.close}): {m.sum()} points")

# any hole left inside the mask is a mask hole, not a real skin island: fill enclosed holes by
# dropping small components of the COMPLEMENT that do not touch the rest of the body
hlab, hsizes = components(~m)
if len(hsizes):
    big = int(np.argmax(hsizes))              # the body itself
    fill = (~m) & (hlab >= 0) & (hlab != big)
    small = np.isin(hlab, np.nonzero(hsizes < 5000)[0]) & fill
    if small.any():
        m |= small
        log(f"filled {small.sum()} enclosed hole points ({(hsizes<5000).sum()-0} small comps)")

if A.grow:
    m = ring(m, A.grow)
    log(f"after grow({A.grow}): {m.sum()} points")

lab2, sizes2 = components(m)
log(f"FINAL mask: {m.sum()} welded points ({100*m.sum()/ng:.2f}%), {len(sizes2)} islands, "
    f"sizes {sorted(sizes2)[::-1][:10]}")
vm = m[inv]
log(f"        -> {vm.sum()} mesh verts; z {zf[vm].min():.3f}..{zf[vm].max():.3f} of height, "
    f"|x| max {np.abs(W[vm,0]).max()/H:.3f}")
for cid in np.argsort(sizes2)[::-1][:8]:
    sel = lab2 == cid
    zs = zf[sel[inv]]
    log(f"    island {int(sizes2[cid]):7d} pts  z {zs.min():.3f}..{zs.max():.3f}")

MASK_OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "leaf_mask.npz")
np.savez_compressed(MASK_OUT,
                    mask=m, inv=inv.astype(np.int32), hue=hue.astype(np.float32),
                    sat=sat.astype(np.float32), val=mx.astype(np.float32))
log(f"WROTE {MASK_OUT}")

# =============================================================================================
# prove it: bake the mask to vertex colour and render
# =============================================================================================
if not A.no_render:
    import math
    from mathutils import Vector
    lay = me.color_attributes.new(name="LeafMask", type='FLOAT_COLOR', domain='POINT')
    colv = np.zeros((n, 4)); colv[:, 3] = 1.0
    colv[:, 0] = np.where(vm, 1.0, 0.45)
    colv[:, 1] = np.where(vm, 0.05, 0.44)
    colv[:, 2] = np.where(vm, 0.05, 0.42)
    lay.data.foreach_set("color", colv.ravel())
    me.materials.clear()
    mat = bpy.data.materials.new("MaskDbg"); mat.use_nodes = True
    nt = mat.node_tree
    vc = nt.nodes.new("ShaderNodeVertexColor"); vc.layer_name = "LeafMask"
    bsdf = next(x for x in nt.nodes if x.type == 'BSDF_PRINCIPLED')
    nt.links.new(vc.outputs["Color"], bsdf.inputs["Base Color"])
    bsdf.inputs["Roughness"].default_value = 0.5
    me.materials.append(mat)

    wl = bpy.data.worlds.new("W"); wl.color = (0.20, 0.20, 0.22)
    bpy.context.scene.world = wl
    scn = bpy.context.scene
    scn.render.engine = 'BLENDER_EEVEE' if bpy.app.version >= (4, 2) else 'BLENDER_EEVEE_NEXT'
    scn.render.resolution_x = scn.render.resolution_y = 1000
    cam = bpy.data.objects.new("Cam", bpy.data.cameras.new("Cam")); cam.data.lens = 85
    bpy.context.collection.objects.link(cam); scn.camera = cam
    key = bpy.data.objects.new("Key", bpy.data.lights.new("Key", 'SUN')); key.data.energy = 3.0
    bpy.context.collection.objects.link(key)
    fill = bpy.data.objects.new("Fill", bpy.data.lights.new("Fill", 'SUN')); fill.data.energy = 1.2
    bpy.context.collection.objects.link(fill)
    z0 = W[:, 2].min()
    for name, f_lo, f_hi, yawdeg in [("chest_front", 0.60, 0.82, 0), ("chest_34", 0.60, 0.82, 40),
                                     ("hip_front", 0.38, 0.60, 0), ("hip_34", 0.38, 0.60, 40),
                                     ("hip_back", 0.38, 0.60, 180), ("full_front", 0.0, 1.0, 0)]:
        z_lo, z_hi = z0 + f_lo * H, z0 + f_hi * H
        band = W[(W[:, 2] >= z_lo) & (W[:, 2] <= z_hi)]
        band = band if len(band) else W
        ctr = Vector((0.0, float(band[:, 1].mean()), (z_lo + z_hi) / 2))
        span = max(float(band[:, 0].max() - band[:, 0].min()), z_hi - z_lo)
        if f_hi - f_lo < 0.5:
            span = min(span, (z_hi - z_lo) * 1.25)
        yaw = math.radians(yawdeg); dist = span * 2.9
        cam.location = ctr + Vector((math.sin(yaw) * dist, -math.cos(yaw) * dist, 0.0))
        cam.rotation_euler = (ctr - cam.location).to_track_quat('-Z', 'Y').to_euler()
        key.rotation_euler = (math.radians(60), 0, math.radians(35 + yawdeg))
        fill.rotation_euler = (math.radians(75), 0, math.radians(yawdeg - 110))
        scn.render.filepath = os.path.join(OUT, f"mask_{name}.png")
        bpy.ops.render.render(write_still=True)
        print("RENDER_OK", scn.render.filepath, flush=True)