# Stage 12: membrane-heal the garment dig-in lines (waistband ledge, leg-hem creases,
# belly/underbust dashes). These are DENTS the briefs/bra pressed into the body — the briefs
# zone's 1x anatomy field kept them, and Taubin/melting preserves exactly this mid-frequency
# shape. Fix = the proven pattern: narrow bands along the crease curves (mid-freq roughness
# detector + mapped slit seams), rim-anchored bi-harmonic membrane across them.
#   blender --background --python 12_dent_membrane.py -- <in.blend> <raw.glb> <out.blend>
import bpy, sys, time
import numpy as np
from mathutils import Vector
from mathutils.kdtree import KDTree

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

ROUGH_THR = 0.0006
Z0, Z1 = 0.42, 0.655        # waist/hip/belly lines up to under the mounds; crotch already healed
GROW_FREE = 3
t0 = time.time()


def log(m):
    print(f"[dent {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)
co = np.empty(n_v * 3)
me.vertices.foreach_get("co", co)
co = co.reshape(-1, 3)

n_e = len(me.edges)
ev0 = np.empty(n_e * 2, dtype=np.int32)
me.edges.foreach_get("vertices", ev0)
ev0 = ev0.reshape(-1, 2)
o_r = np.concatenate([ev0[:, 0], ev0[:, 1]])
n_r = np.concatenate([ev0[:, 1], ev0[:, 0]])
s_r = np.argsort(o_r, kind="stable")
o_rs = o_r[s_r]
n_rs = n_r[s_r]
ptr_r = np.searchsorted(o_rs, np.arange(n_v + 1))
cnt_r = np.maximum(np.diff(ptr_r), 1)

sm = co.copy()
for _ in range(8):
    su = np.add.reduceat(sm[n_rs], ptr_r[:-1], axis=0)
    emp = np.diff(ptr_r) == 0
    su[emp] = sm[emp]
    sm = su / cnt_r[:, None]
rough = np.linalg.norm(co - sm, axis=1)

# slit seams from the raw pre-weld GLB (the dotted dash lines)
before = set(bpy.data.objects)
bpy.ops.import_scene.gltf(filepath=RAW)
new = [o for o in bpy.data.objects if o not in before]
raw = max([o for o in new if o.type == 'MESH'], key=lambda o: len(o.data.vertices))
rme = raw.data
rn = len(rme.vertices)
rco = np.empty(rn * 3)
rme.vertices.foreach_get("co", rco)
rco = rco.reshape(-1, 3)
l_tot = np.empty(len(rme.polygons), dtype=np.int32)
rme.polygons.foreach_get("loop_total", l_tot)
l_start = np.empty(len(rme.polygons), dtype=np.int32)
rme.polygons.foreach_get("loop_start", l_start)
l_v = np.empty(len(rme.loops), dtype=np.int32)
rme.loops.foreach_get("vertex_index", l_v)
ecount = {}
for fs, ft in zip(l_start, l_tot):
    idxs = l_v[fs:fs + ft]
    for k in range(ft):
        a, b = idxs[k], idxs[(k + 1) % ft]
        kk = (a, b) if a < b else (b, a)
        ecount[kk] = ecount.get(kk, 0) + 1
rbnd = np.zeros(rn, dtype=bool)
for (a, b), c in ecount.items():
    if c == 1:
        rbnd[a] = rbnd[b] = True
for o in new:
    bpy.data.objects.remove(o, do_unlink=True)

kd = KDTree(n_v)
for i in range(n_v):
    kd.insert(Vector(co[i]), i)
kd.balance()
seam = np.zeros(n_v, dtype=bool)
for p in rco[rbnd]:
    if not (Z0 < p[2] < Z1):
        continue
    hit = kd.find(Vector(p))
    if hit[0] is not None and hit[2] < 0.006:
        seam[hit[1]] = True
log(f"seam verts: {seam.sum()}")

navel = (np.abs(co[:, 0]) < 0.022) & (co[:, 2] > 0.495) & (co[:, 2] < 0.555) & (co[:, 1] < 0)
band = (co[:, 2] > Z0) & (co[:, 2] < Z1) & ~navel
hot = band & ((rough > ROUGH_THR) | seam)
log(f"hot line verts: {hot.sum()}")


def grow_edges(mask, rings, ev):
    m = mask.copy()
    for _ in range(rings):
        hit = m[ev[:, 0]] | m[ev[:, 1]]
        m2 = m.copy()
        m2[ev[:, 0]] |= hit
        m2[ev[:, 1]] |= hit
        m = m2
    return m


free_m = grow_edges(hot, GROW_FREE, ev0) & ~navel
collar = grow_edges(free_m, 2, ev0) & ~free_m
S = np.nonzero(free_m | collar)[0]
in_S = np.zeros(n_v, dtype=bool)
in_S[S] = True
glb = np.full(n_v, -1, dtype=np.int64)
glb[S] = np.arange(len(S))
se = ev0[in_S[ev0].all(axis=1)]
a_ = glb[se[:, 0]]
b_ = glb[se[:, 1]]
deg = np.zeros(len(S))
np.add.at(deg, a_, 1.0)
np.add.at(deg, b_, 1.0)
free = free_m[S]
log(f"free {free.sum()}, collar {(~free).sum()}")


def Ls(X):
    out = deg[:, None] * X
    np.add.at(out, a_, -X[b_])
    np.add.at(out, b_, -X[a_])
    return out


def A_op(U):
    X = np.zeros((len(S), 3))
    X[free] = U
    return Ls(Ls(X))[free]


Xc = np.zeros((len(S), 3))
Xc[~free] = co[S[~free]]
rhs = -Ls(Ls(Xc))[free]
U = co[S[free]].copy()
r = rhs - A_op(U)
p = r.copy()
rs = (r * r).sum()
rs0 = rs
for it in range(120000):
    Ap = A_op(p)
    al = rs / max((p * Ap).sum(), 1e-30)
    U += al * p
    r -= al * Ap
    rs2 = (r * r).sum()
    if rs2 < 1e-18 or rs2 < rs0 * 1e-14:
        break
    p = r + (rs2 / rs) * p
    rs = rs2
co_new = co.copy()
co_new[S[free]] = U
d = np.linalg.norm(co_new - co, axis=1)
log(f"membrane: {free.sum()} verts (CG {it}, rel {rs2/max(rs0,1e-30):.2e}), max move {d.max():.4f}")

me.vertices.foreach_set("co", co_new.reshape(-1))
me.update()
if me.has_custom_normals:
    vn = np.empty(n_v * 3, dtype=np.float32)
    me.vertices.foreach_get("normal", vn)
    me.normals_split_custom_set_from_vertices(vn.reshape(-1, 3))
bpy.context.preferences.filepaths.save_version = 0  # no .blend1 autosave
bpy.ops.wm.save_as_mainfile(filepath=OUT)
log(f"WROTE {OUT}")
print("DENT_DONE")