# Stage 11: heal the garment-line creases across the stomach/waist/hips.
# The lines are GEOMETRY (they show in clay): the briefs waistband ledge and leg-hem ridges
# (the briefs zone's 1x field kept hip anatomy AND the hem ridges), plus dotted pinch lines
# where the source mesh's open slits were welded. Strip = hemband mask ∪ slit-seam verts
# (mapped from the RAW pre-weld GLB's boundary edges) ∪ normal-kink verts in the torso band.
# Taubin (volume-preserving) on the strip: ridges round off, hips/butt keep their shape.
#   blender --background --python 11_line_heal.py -- <in.blend> <masks.npz> <raw.glb> <out.blend>
import bpy, sys, time, math
import numpy as np
from mathutils import Vector
from mathutils.kdtree import KDTree

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

ANG_THR = math.radians(12.0)
TAUBIN_PAIRS = 30
Z0, Z1 = 0.30, 0.87


def log(m):
    print(f"[heal {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)
M = np.load(MASKS)
hemband = M["hemband"]

# ---- slit-seam verts from the raw pre-weld GLB ----
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]
        key = (a, b) if a < b else (b, a)
        ecount[key] = ecount.get(key, 0) + 1
rbnd = np.zeros(rn, dtype=bool)
for (a, b), c in ecount.items():
    if c == 1:
        rbnd[a] = rbnd[b] = True
log(f"raw boundary verts: {rbnd.sum()}")
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)
# NB: positions have been sculpted since the weld — match generously but only in the torso
# band, and only trust matches within 6 mm (sculpted garment areas moved far more; their
# seams are already handled by the fills/melts there)
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 mapped: {seam.sum()}")

# ---- normal-kink verts (ledges/ridges) ----
nrm = np.empty(n_v * 3)
me.vertices.foreach_get("normal", nrm)
nrm = nrm.reshape(-1, 3)
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")
o_s = order[srt]
n_s = nbr[srt]
ptr = np.searchsorted(o_s, np.arange(n_v + 1))
cnt = np.maximum(ptr[1:] - ptr[:-1], 1)
N = nrm.copy()
for _ in range(5):
    acc = np.zeros_like(N)
    np.add.at(acc, o_s, N[n_s])
    N = acc / cnt[:, None]
    N /= np.maximum(np.linalg.norm(N, axis=1, keepdims=True), 1e-12)
ang = np.arccos(np.clip((nrm * N).sum(axis=1), -1, 1))
navel = (np.abs(co[:, 0]) < 0.022) & (co[:, 2] > 0.495) & (co[:, 2] < 0.555) & (co[:, 1] < 0)
band = (co[:, 2] > 0.42) & (co[:, 2] < Z1) & ~navel
kink = band & (ang > ANG_THR)
log(f"kink verts: {kink.sum()}")

strip = (hemband | seam | kink) & (co[:, 2] > Z0) & (co[:, 2] < Z1) & ~navel
for _ in range(2):
    hit = strip[ev[:, 0]] | strip[ev[:, 1]]
    s2 = strip.copy()
    s2[ev[:, 0]] |= hit
    s2[ev[:, 1]] |= hit
    strip = s2
strip &= ~navel
sidx = np.nonzero(strip)[0]
log(f"strip: {len(sidx)} verts")

Q = co.copy()
lam, mu = 0.5, -0.53
for _ in range(TAUBIN_PAIRS):
    for f in (lam, mu):
        acc = np.zeros_like(Q)
        np.add.at(acc, o_s, Q[n_s])
        mean = acc / cnt[:, None]
        Q[sidx] += f * (mean[sidx] - Q[sidx])
d = np.linalg.norm(Q - co, axis=1)
log(f"Taubin x{TAUBIN_PAIRS}: max move {d.max():.4f}")

me.vertices.foreach_set("co", Q.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("HEAL_DONE")