# lena_leafbikini lane, stage 02: PROBE — what exactly are the leaves?
#
#   blender --background --factory-startup --python 02_probe_leaves.py -- <pristine.glb> [outdir]
#
# Before anything is cut, this answers the two questions that decide the method:
#
#   1. Are the leaves PAINTED or SCULPTED? Lena's underwear on the game body turned out to be
#      painted onto the body skin with no geometry of its own (characters/work/lena/README,
#      tools/make_lena_nude_body.py finding 1) — deleting it opened a hole because the garment
#      WAS the skin. If the leaves are the same, "remove at the seam" means a colour-keyed face
#      delete and nothing more. If they are real shells sitting proud of the body, the seam is a
#      geometric crease and the colour key is only a coarse pre-filter.
#   2. Where is the seam? Reported here as the distribution of per-vertex proudness (signed
#      distance from a heavily smoothed reference surface) inside vs outside the green key.
#
# Everything is numpy over foreach_get buffers: this mesh is ~1.03M verts and a per-vertex
# Python loop over it costs minutes.
import bpy, sys, os, time
import numpy as np

argv = sys.argv[sys.argv.index("--") + 1:]
GLB = os.path.abspath(argv[0])
OUT = os.path.abspath(argv[1]) if len(argv) > 1 else os.path.dirname(GLB)
os.makedirs(OUT, exist_ok=True)
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=GLB)
meshes = [o for o in bpy.data.objects if o.type == 'MESH']
log(f"objects: {[(o.name, o.type) for o in bpy.data.objects]}")
for o in meshes:
    log(f"  MESH '{o.name}': {len(o.data.vertices)}v {len(o.data.polygons)}f "
        f"uv={[l.name for l in o.data.uv_layers]} mats={[m.name for m in o.data.materials if m]}")
body = max(meshes, key=lambda o: len(o.data.vertices))
me = body.data
n = len(me.vertices)

# ── geometry ────────────────────────────────────────────────────────────────────────────────
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]
lo, hi = W.min(0), W.max(0)
log(f"bbox min={np.round(lo,4)} max={np.round(hi,4)} size={np.round(hi-lo,4)}")
log(f"height(z) {hi[2]-lo[2]:.4f}  -> 1 unit = {1.777/(hi[2]-lo[2]):.4f} of a 1.777 m body")

# ── textures ────────────────────────────────────────────────────────────────────────────────
imgs = {}
for mat in [m for m in me.materials if m]:
    for nd in mat.node_tree.nodes:
        if nd.type == 'TEX_IMAGE' and nd.image:
            tgt = [l.to_socket.name for o in nd.outputs for l in o.links]
            log(f"  tex '{nd.image.name}' {tuple(nd.image.size)} cs={nd.image.colorspace_settings.name} -> {tgt}")
            imgs[nd.image.name] = nd.image

# base colour = the image feeding Base Color
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)
    if not bsdf:
        continue
    lnk = bsdf.inputs["Base Color"].links
    if lnk:
        nd = lnk[0].from_node
        while nd.type != 'TEX_IMAGE' and nd.inputs:
            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("[probe] FATAL: no base-colour image found")
log(f"base colour image: '{base.name}' {tuple(base.size)}")

# ── per-vertex UV (first loop wins), then sample the albedo ─────────────────────────────────
nl = len(me.loops)
lv = np.empty(nl, dtype=np.int32); me.loops.foreach_get("vertex_index", lv)
uv = np.empty(nl * 2); me.uv_layers.active.data.foreach_get("uv", uv); uv = uv.reshape(-1, 2)
vuv = np.zeros((n, 2))
vuv[lv[::-1]] = uv[::-1]          # reversed scatter -> first loop of each vert 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]
del buf
# bpy-imported UVs are already v-flipped by the importer (see 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]                    # linear RGB per vertex
del px
log(f"sampled albedo for {n} verts from {w}x{h}")

# linear -> sRGB for a colour key that matches what the eye/Tripo saw
def to_srgb(x):
    return np.where(x <= 0.0031308, x * 12.92, 1.055 * np.maximum(x, 0) ** (1 / 2.4) - 0.055)


S = np.clip(to_srgb(C), 0, 1)
R, G, B = S[:, 0], S[:, 1], S[:, 2]
mx, mn = S.max(1), S.min(1)
sat = np.where(mx > 1e-5, (mx - mn) / np.maximum(mx, 1e-5), 0.0)
# green dominance: G is the max channel and beats both others
gdom = (G - np.maximum(R, B))
log(f"albedo sRGB: mean R{R.mean():.3f} G{G.mean():.3f} B{B.mean():.3f}  sat mean {sat.mean():.3f}")
for thr in (0.0, 0.02, 0.05, 0.08, 0.12, 0.20):
    m = gdom > thr
    log(f"  G-dominance > {thr:.2f}: {m.sum():7d} verts ({100*m.sum()/n:5.2f}%)"
        + (f"  z {W[m,2].min():.3f}..{W[m,2].max():.3f}" if m.any() else ""))

# ── proudness: signed offset from a smoothed reference surface ─────────────────────────────
# Adjacency over POSITION-WELDED points (the importer splits every UV seam; see nude-body
# finding 5). Built as a flat CSR from the edge list so smoothing is pure numpy.
key = np.round(W, 6)
_, inv = np.unique(key, axis=0, return_inverse=True)
ng = inv.max() + 1
log(f"welded: {n} verts -> {ng} unique positions ({n-ng} seam duplicates)")
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]]
keep = ea != eb
ea, eb = ea[keep], eb[keep]
src = np.concatenate([ea, eb]); dst = np.concatenate([eb, ea])
order = np.argsort(src, kind='stable')
src, dst = src[order], dst[order]
cnt = np.bincount(src, minlength=ng)
ptr = np.concatenate([[0], np.cumsum(cnt)])
cnt_safe = np.maximum(cnt, 1)

P = np.zeros((ng, 3)); np.add.at(P, inv, W); P /= np.bincount(inv, minlength=ng)[:, None]


def nbr_mean(X):
    s = np.add.reduceat(X[dst], ptr[:-1], axis=0)
    s[cnt == 0] = X[cnt == 0]
    return s / cnt_safe[:, None]


Q = P.copy()
for _ in range(60):                       # heavy Taubin: sheds the leaves, keeps the body
    Q += 0.55 * (nbr_mean(Q) - Q)
    Q += -0.58 * (nbr_mean(Q) - Q)
# reference normal from the smoothed surface, via the vertex-normal buffer of the ORIGINAL
vn = np.empty(n * 3); me.vertices.foreach_get("normal", vn); vn = vn.reshape(-1, 3)
N = np.zeros((ng, 3)); np.add.at(N, inv, vn)
N /= np.maximum(np.linalg.norm(N, axis=1, keepdims=True), 1e-12)
proud_g = np.einsum('ij,ij->i', P - Q, N)
proud = proud_g[inv]
mm = 1000.0 * 1.777 / (hi[2] - lo[2])     # units -> mm on a 1.777 m body
log(f"proudness (mm, body-scaled): mean {proud.mean()*mm:+.2f}  p50 {np.median(proud)*mm:+.2f}  "
    f"p99 {np.percentile(proud,99)*mm:+.2f}  max {proud.max()*mm:+.2f}")
for thr in (0.05, 0.12):
    m = gdom > thr
    if m.sum() < 100:
        continue
    log(f"  green(G-dom>{thr}): proud p50 {np.median(proud[m])*mm:+.2f} mm  "
        f"p90 {np.percentile(proud[m],90)*mm:+.2f} mm")
    log(f"  skin (G-dom<=0 ): proud p50 {np.median(proud[~(gdom>0)])*mm:+.2f} mm  "
        f"p90 {np.percentile(proud[~(gdom>0)],90)*mm:+.2f} mm")
    break

np.savez_compressed(os.path.join(OUT, "probe_leaves.npz"),
                    gdom=gdom.astype(np.float32), sat=sat.astype(np.float32),
                    proud=proud.astype(np.float32), W=W.astype(np.float32),
                    inv=inv.astype(np.int32))
log(f"WROTE {os.path.join(OUT, 'probe_leaves.npz')}")