# Probe the pre-decimated Tripo sculpt: is the sports top a SEPARABLE shell?
# Answers, per connected component (after welding UV-seam splits):
#   size, z-range, open-boundary edge count, and what fraction of its faces are painted
#   garment (colour-keyed on lena+glb_basecolor via each face's own UV texels).
#
#   blender --background --python 01_probe.py -- <copy.glb> <checkpoint.blend>
import bpy, bmesh, sys, time
import numpy as np

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


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


bpy.ops.wm.read_factory_settings(use_empty=True)
bpy.ops.import_scene.gltf(filepath=SRC)
ob = max([o for o in bpy.data.objects if o.type == 'MESH'],
         key=lambda o: len(o.data.vertices))
me = ob.data
log(f"imported: {len(me.vertices)}v {len(me.polygons)}f, object '{ob.name}'")

bm = bmesh.new()
bm.from_mesh(me)
bmesh.ops.remove_doubles(bm, verts=list(bm.verts), dist=1e-6)
bm.to_mesh(me)
bm.free()
me.update()
log(f"after weld 1e-6: {len(me.vertices)}v {len(me.polygons)}f")

# save checkpoint so later stages skip the import+weld cost
bpy.context.preferences.filepaths.save_version = 0  # no .blend1 autosave
bpy.ops.wm.save_as_mainfile(filepath=CKPT)
log(f"checkpoint saved: {CKPT}")

n_v = len(me.vertices)
n_f = len(me.polygons)

# ---- connected components over faces (edge-connected), union-find in numpy ----
# face -> vertices
loop_tot = np.empty(n_f, dtype=np.int32)
me.polygons.foreach_get("loop_total", loop_tot)
loop_start = np.empty(n_f, dtype=np.int32)
me.polygons.foreach_get("loop_start", loop_start)
loops_v = np.empty(len(me.loops), dtype=np.int32)
me.loops.foreach_get("vertex_index", loops_v)

parent = np.arange(n_v, dtype=np.int64)


def find(a):
    root = a
    while parent[root] != root:
        root = parent[root]
    while parent[a] != root:
        parent[a], a = root, parent[a]
    return root


# union via 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)
for a, b in ev:
    ra, rb = find(a), find(b)
    if ra != rb:
        parent[rb] = ra
log("union-find done")

root_of = np.array([find(i) for i in range(n_v)], dtype=np.int64)
uniq, inv, counts = np.unique(root_of, return_inverse=True, return_counts=True)
order = np.argsort(-counts)
log(f"components: {len(uniq)}")

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

# boundary edges per component: edge belongs to 1 face only
# count faces per edge
edge_face_count = np.zeros(n_e, dtype=np.int32)
for p in me.polygons:
    for ek in p.edge_keys:
        pass  # too slow; use loops instead
# faster: build edge keys from loops via me.polygons edge indices
# Blender exposes loop edges: me.loops[i].edge_index
loops_e = np.empty(len(me.loops), dtype=np.int32)
me.loops.foreach_get("edge_index", loops_e)
np.add.at(edge_face_count, loops_e, 1)
boundary_edge = edge_face_count == 1
log(f"boundary edges total: {boundary_edge.sum()}")

# ---- colour key per vertex (first-loop UV), calibrated on this texture ----
img = None
for i in bpy.data.images:
    if "basecolor" in i.name.lower():
        img = i
        break
if img is None:
    img = max(bpy.data.images, key=lambda i: i.size[0] * i.size[1])
w, h = img.size
log(f"texture '{img.name}' {w}x{h}")
px = np.empty(w * h * 4, dtype=np.float32)
img.pixels.foreach_get(px)
rgb = px.reshape(h, w, 4)[:, :, :3]

uvl = me.uv_layers.active.data
uv = np.empty(len(me.loops) * 2, dtype=np.float64)
uvl.foreach_get("uv", uv)
uv = uv.reshape(-1, 2)
# first loop per vertex
first_loop = np.full(n_v, -1, dtype=np.int64)
for li in range(len(loops_v) - 1, -1, -1):
    first_loop[loops_v[li]] = li
has_uv = first_loop >= 0
vx = np.clip(uv[first_loop, 0], 0, 1) * (w - 1)
vy = np.clip(uv[first_loop, 1], 0, 1) * (h - 1)   # bpy-imported UVs: already flipped
vcol = rgb[vy.astype(int), vx.astype(int)]
r, g, b = vcol[:, 0], vcol[:, 1], vcol[:, 2]
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 = r / np.maximum(b, 1e-5)

# calibrate: thigh skin (z 0.28-0.34) vs briefs centre (z 0.50-0.56 front)
thigh = (co[:, 2] > 0.28) & (co[:, 2] < 0.34)
briefs = (co[:, 2] > 0.50) & (co[:, 2] < 0.56) & (co[:, 1] < 0) & (np.abs(co[:, 0]) < 0.05)
bra = (co[:, 2] > 0.64) & (co[:, 2] < 0.72) & (co[:, 1] < 0) & (np.abs(co[:, 0]) < 0.05)
for nm, s in (("thigh skin", thigh), ("briefs", briefs), ("bra front", bra)):
    if s.sum():
        log(f"  [{nm}] n={s.sum()} sat p50={np.median(sat[s]):.3f} rb p50={np.median(rb[s]):.3f} "
            f"rgb ({np.median(r[s]):.3f},{np.median(g[s]):.3f},{np.median(b[s]):.3f})")

# pick thresholds midway between garment and skin medians
sat_thr = (np.median(sat[briefs]) + np.median(sat[thigh])) / 2 if briefs.sum() and thigh.sum() else 0.45
rb_thr = (np.median(rb[briefs]) + np.median(rb[thigh])) / 2 if briefs.sum() and thigh.sum() else 1.9
garment_v = (sat < sat_thr) & (rb < rb_thr)
log(f"thresholds: sat<{sat_thr:.3f} rb<{rb_thr:.2f} -> {garment_v.sum()} garment verts")

# ---- per-component report (top 12 by size) ----
# per-vertex boundary flag
vb = np.zeros(n_v, dtype=bool)
vb[ev[boundary_edge].ravel()] = True

print("\n=== COMPONENTS (top 12 by vertex count) ===")
for ci in order[:12]:
    root = uniq[ci]
    m = root_of == root
    zc = co[m, 2]
    gfrac = garment_v[m].mean()
    bcount = vb[m].sum()
    print(f"comp root={root}: verts={m.sum():7d}  z[{zc.min():.3f},{zc.max():.3f}]  "
          f"garment-painted={100*gfrac:5.1f}%  boundary-verts={bcount}")
print("PROBE_DONE")