# Stage 2: isolate the garment on the welded hires sculpt, split it into functional zones, and
# save the masks + adjacency for the sculpt stage.
#
#   blender --background --python 02_isolate.py -- <00_welded.blend> <masks.npz>
#
# Zones (per vertex):
#   cups    — bra front where the breasts live (gets wall + amplified mounds)
#   shield  — sternum strip the bra bridges (gets carved cleavage)
#   toprest — rest of the top: band, straps, back panel (gets faired flat to skin)
#   briefs  — briefs incl. waistband/leg rims (gets faired featureless)
# plus 'hemband' — narrow ring around every garment/skin boundary (local crease fairing).
#
# The colour key was calibrated on THIS texture by 01_probe.py: garment sat~0.39/RB~1.65,
# skin sat~0.64/RB~2.80; thresholds are the midpoints. The z-guards keep the low-saturation
# face features (eyes, teeth, lips) out of the key's reach.
import bpy, sys, time
import numpy as np

argv = sys.argv[sys.argv.index("--") + 1:]
BLEND, OUT = argv[0], argv[1]
t0 = time.time()

SAT_THR, RB_THR = 0.518, 2.22
Z_GARMENT_LO, Z_GARMENT_HI = 0.40, 0.85      # below the head; briefs to over-shoulder straps
Z_TOP_SPLIT = 0.595                          # briefs/top divide (waistband top 0.578 + margin)
GROW_RINGS = 6            # ~2 cm real at this density; must swallow the 1 cm hem lips

# cups zone (front bra): z and |x| bounds from the measured landmarks
CUP_Z_LO, CUP_Z_HI = 0.615, 0.745
CUP_X_MAX = 0.085
SHIELD_X = 0.014                             # sternum strip half-width


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


bpy.ops.wm.open_mainfile(filepath=BLEND)
ob = max([o for o in bpy.data.objects if o.type == 'MESH'],
         key=lambda o: len(o.data.vertices))
me = ob.data
n_v = len(me.vertices)
log(f"loaded {n_v}v")

co = np.empty(n_v * 3, dtype=np.float64)
me.vertices.foreach_get("co", co)
co = co.reshape(-1, 3)

# --- per-vertex colour ---
img = next(i for i in bpy.data.images if "basecolor" in i.name.lower())
w, h = img.size
px = np.empty(w * h * 4, dtype=np.float32)
img.pixels.foreach_get(px)
rgb = px.reshape(h, w, 4)[:, :, :3]

loops_v = np.empty(len(me.loops), dtype=np.int32)
me.loops.foreach_get("vertex_index", loops_v)
uv = np.empty(len(me.loops) * 2, dtype=np.float64)
me.uv_layers.active.data.foreach_get("uv", uv)
uv = uv.reshape(-1, 2)
first_loop = np.full(n_v, len(loops_v), dtype=np.int64)
np.minimum.at(first_loop, loops_v, np.arange(len(loops_v), dtype=np.int64))
first_loop = np.minimum(first_loop, len(loops_v) - 1)
vx = (np.clip(uv[first_loop, 0], 0, 1) * (w - 1)).astype(int)
vy = (np.clip(uv[first_loop, 1], 0, 1) * (h - 1)).astype(int)
vcol = rgb[vy, vx]
mx = vcol.max(axis=1)
mn = vcol.min(axis=1)
sat = np.where(mx > 1e-5, (mx - mn) / np.maximum(mx, 1e-5), 0)
rb = vcol[:, 0] / np.maximum(vcol[:, 2], 1e-5)

# above z=0.80 only the shoulder straps qualify — keep the chin/lips (low-saturation paint)
# out by requiring lateral offset there
zone_guard = (co[:, 2] > Z_GARMENT_LO) & (co[:, 2] < Z_GARMENT_HI)              & ((co[:, 2] < 0.80) | (np.abs(co[:, 0]) > 0.025))
key = (sat < SAT_THR) & (rb < RB_THR) & zone_guard
# The shoulder straps' paint is much closer to skin (median sat 0.57-0.61 vs the bra body's
# 0.39) — the main key catches only ~a third of each strap and the rest survived as raised
# geometry. In the strap corridor a relaxed threshold seeds them; the ring grow fills the rest.
strap_zone = (co[:, 2] > 0.72) & (co[:, 2] < 0.85)              & (np.abs(co[:, 0]) > 0.03) & (np.abs(co[:, 0]) < 0.105)
key |= strap_zone & (sat < 0.55)
log(f"colour key (+strap corridor): {key.sum()} verts")

# --- adjacency (numpy CSR-ish over edges) ---
n_e = len(me.edges)
ev = np.empty(n_e * 2, dtype=np.int32)
me.edges.foreach_get("vertices", ev)
ev = ev.reshape(-1, 2)
order = np.concatenate([ev[:, 0], ev[:, 1]])
nbr = np.concatenate([ev[:, 1], ev[:, 0]])
srt = np.argsort(order, kind="stable")
order_s = order[srt]
nbr_s = nbr[srt]
ptr = np.searchsorted(order_s, np.arange(n_v + 1))
log("adjacency built")


def grow(mask, rings):
    out = mask.copy()
    for _ in range(rings):
        sel = np.zeros(n_v, dtype=bool)
        # mark all neighbours of current selection
        active = np.nonzero(out)[0]
        # gather neighbour slices
        for a in active:
            sel[nbr_s[ptr[a]:ptr[a + 1]]] = True
        out |= sel
    return out


# component filter: drop specks (<200 verts) via BFS on the keyed set
from collections import deque
comp_id = np.full(n_v, -1, dtype=np.int64)
cid = 0
keep = np.zeros(n_v, dtype=bool)
for s in np.nonzero(key)[0]:
    if comp_id[s] >= 0:
        continue
    q = deque([s])
    comp_id[s] = cid
    members = [s]
    while q:
        c = q.popleft()
        for nb in nbr_s[ptr[c]:ptr[c + 1]]:
            if key[nb] and comp_id[nb] < 0:
                comp_id[nb] = cid
                q.append(nb)
                members.append(nb)
    if len(members) >= 200:
        keep[members] = True
    cid += 1
log(f"component filter: {keep.sum()} verts in {cid} raw components")

garment = grow(keep, GROW_RINGS)
log(f"grown +{GROW_RINGS}: {garment.sum()}")

top = garment & (co[:, 2] >= Z_TOP_SPLIT)
briefs = garment & (co[:, 2] < Z_TOP_SPLIT)

front = co[:, 1] < 0.0
cups = top & front & (co[:, 2] > CUP_Z_LO) & (co[:, 2] < CUP_Z_HI) \
       & (np.abs(co[:, 0]) < CUP_X_MAX) & (np.abs(co[:, 0]) > SHIELD_X)
shield = top & front & (co[:, 2] > CUP_Z_LO) & (co[:, 2] < CUP_Z_HI) \
         & (np.abs(co[:, 0]) <= SHIELD_X)
toprest = top & ~cups & ~shield

# hem band: garment boundary vs non-garment, +/-2 rings
edge_g = garment[ev]
bnd_edges = ev[edge_g[:, 0] != edge_g[:, 1]]
bmask = np.zeros(n_v, dtype=bool)
bmask[bnd_edges.ravel()] = True
hemband = grow(bmask, 8)   # hem lips are ~3 rings wide; blend needs room beyond them
log(f"zones: cups={cups.sum()} shield={shield.sum()} toprest={toprest.sum()} "
    f"briefs={briefs.sum()} hemband={hemband.sum()}")

np.savez_compressed(OUT, garment=garment, cups=cups, shield=shield, toprest=toprest,
                    briefs=briefs, hemband=hemband, key_raw=keep)
log(f"WROTE {OUT}")
print("ISOLATE_DONE")