# Stage 39: the remaining speckle is GEOMETRY (confirmed by clay render, 38_clay.py) — small
# shards, slivers and one-vertex spikes on the hips, thighs and belly. Diagnose first, fix only
# what the numbers show.
#
#   blender --background --python 39_despeckle.py -- <in.blend> <out.blend> [--apply] [out.glb]
#
# THE METRIC — and the one that did NOT work. Deviation from a locally smoothed copy of the
# surface was tried first and is useless at this density: with ~10 mm edges, three Laplacian
# passes erase real curvature, so |dev| > 1 mm flags 33% of the body — her hips and her toes score
# the same as a defect. Speckle is not "far from smooth", it is SHARP and ISOLATED. So:
#   sharpness = dihedral angle across an edge (a flap folds back on itself; a hip does not)
#   isolation = size of the connected cluster of sharp vertices (a real crease under the buttock
#               runs for hundreds of vertices; a shard is a handful)
# Requiring both is what separates a defect from anatomy, and neither test alone does.
#
# WHY NOT DELETE AND FILL. Already tried twice in this lane and refused both times: the 830
# boundary edges sit in open chains of 3-5, i.e. dangling flaps, not perforations. Nothing to fill.
# Flattening a flap onto the surface it hangs off is the operation that actually applies, and it
# changes no topology, so the atlas and its textures stay valid.
#
# Vertex moves are clamped, so this cannot quietly restyle her.
import bpy, bmesh, sys, os, time
import numpy as np

argv = sys.argv[sys.argv.index("--") + 1:]
BLEND, OUT = argv[0], argv[1]
APPLY = "--apply" in argv
GLB = next((a for a in argv[2:] if a.lower().endswith(".glb")), "")
t0 = time.time()
SHARD_MAX = 60           # faces; the body is one component of ~65k, so this is unambiguous debris
SMOOTH_N = 3
MAX_MOVE_MM = 3.0


def log(m):
    print(f"[spk {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, n_f = len(me.vertices), len(me.polygons)
co = np.empty(n_v * 3); me.vertices.foreach_get("co", co); co = co.reshape(-1, 3)
UNITM = 1.777 / (co[:, 2].max() - co[:, 2].min())
MM = UNITM * 1000.0
log(f"in: {n_v}v {n_f}f   1 unit = {MM:.1f} mm")

ev = np.empty(len(me.edges) * 2, dtype=np.int32); me.edges.foreach_get("vertices", ev)
ev = ev.reshape(-1, 2)
o_ = np.concatenate([ev[:, 0], ev[:, 1]])
n_ = np.concatenate([ev[:, 1], ev[:, 0]])
srt = np.argsort(o_, kind="stable")
o_s, n_s = o_[srt], n_[srt]
ptr = np.searchsorted(o_s, np.arange(n_v + 1))
cnt = np.maximum(np.diff(ptr), 1)


def nbr_mean(X):
    acc = np.add.reduceat(X[n_s], ptr[:-1], axis=0)
    acc[np.diff(ptr) == 0] = X[np.diff(ptr) == 0]
    return acc / cnt[:, None]


def deviation(P):
    nrm = np.empty(n_v * 3); me.vertices.foreach_get("normal", nrm); nrm = nrm.reshape(-1, 3)
    sm = P.copy()
    for _ in range(SMOOTH_N):
        sm = nbr_mean(sm)
    return ((P - sm) * nrm).sum(axis=1), sm


SHARP_DEG = float(os.environ.get('SHARP_DEG', '70.0'))
CLUSTER_MAX = int(os.environ.get('CLUSTER_MAX', '14'))


def sharp_verts(deg_thresh):
    """vertices touching an edge whose two faces fold by more than deg_thresh"""
    bm_ = bmesh.new(); bm_.from_mesh(me)
    ang, vs = [], np.zeros(len(bm_.verts), dtype=bool)
    for e in bm_.edges:
        if len(e.link_faces) != 2:
            continue
        a = e.calc_face_angle_signed(None)
        if a is None:
            continue
        d = abs(np.degrees(a))
        ang.append(d)
        if d > deg_thresh:
            vs[e.verts[0].index] = True
            vs[e.verts[1].index] = True
    bm_.free()
    return np.array(ang), vs


ang, _ = sharp_verts(1e9)
print("\n=== DIHEDRAL ANGLE ACROSS EDGES (degrees; a fold, not a curve) ===")
for p in (50, 90, 99, 99.5, 99.9):
    print(f"  p{p:<5} {np.percentile(ang, p):7.2f}")
print(f"  max {ang.max():.2f}")
for t in (30, 50, 70, 90):
    print(f"  edges folding > {t}deg: {int((ang > t).sum()):6d} of {len(ang)}")

_, sharp = sharp_verts(SHARP_DEG)
print(f"\nvertices on a >{SHARP_DEG:.0f}deg fold: {int(sharp.sum())}")

# cluster the sharp vertices: anatomy runs long, defects are isolated
def speckle_mask(verbose=False):
    """sharp AND isolated. Recomputable, because deleting debris renumbers vertices — computing
    this once up front and reusing it after a topology change indexed the wrong array."""
    nv = len(me.vertices)
    ev_ = np.empty(len(me.edges) * 2, dtype=np.int32); me.edges.foreach_get("vertices", ev_)
    ev_ = ev_.reshape(-1, 2)
    _, sharp_ = sharp_verts(SHARP_DEG)
    adj_ = {}
    for a_, b_ in ev_:
        if sharp_[a_] and sharp_[b_]:
            adj_.setdefault(a_, []).append(b_)
            adj_.setdefault(b_, []).append(a_)
    lab_ = -np.ones(nv, dtype=np.int64)
    sizes_ = []
    for s_ in np.nonzero(sharp_)[0]:
        if lab_[s_] >= 0:
            continue
        stack, cells = [s_], []
        lab_[s_] = len(sizes_)
        while stack:
            v_ = stack.pop()
            cells.append(v_)
            for w_ in adj_.get(v_, ()):
                if lab_[w_] < 0:
                    lab_[w_] = len(sizes_)
                    stack.append(w_)
        sizes_.append(len(cells))
    sizes_ = np.array(sizes_) if sizes_ else np.array([0])
    small_ = np.array([lab_[i] >= 0 and sizes_[lab_[i]] <= CLUSTER_MAX for i in range(nv)])
    if verbose:
        print("")
        print(f"vertices on a >{SHARP_DEG:.0f}deg fold: {int(sharp_.sum())}")
        print(f"sharp clusters: {len(sizes_)}  (sizes: p50 {np.percentile(sizes_,50):.0f}, "
              f"p90 {np.percentile(sizes_,90):.0f}, max {sizes_.max()})")
        print(f"SPECKLE = sharp AND in a cluster of <= {CLUSTER_MAX} verts: {int(small_.sum())} "
              f"verts in {int((sizes_ <= CLUSTER_MAX).sum())} clusters")
        print(f"  (kept as anatomy: {int(sharp_.sum() - small_.sum())} verts in "
              f"{int((sizes_ > CLUSTER_MAX).sum())} long folds)")
    return small_


small = speckle_mask(verbose=True)
spike = small
if spike.any():
    z = (co[:, 2] - co[:, 2].min()) / (co[:, 2].max() - co[:, 2].min())
    hist, edges = np.histogram(z[spike], bins=10, range=(0, 1))
    print("\nheight distribution of the speckle (0=feet, 1=head):")
    for c, lo, hi in zip(hist, edges[:-1], edges[1:]):
        if c:
            print(f"   {lo:.1f}-{hi:.1f}: {c:5d}")

# stray components
bm = bmesh.new(); bm.from_mesh(me)
bm.verts.ensure_lookup_table()
parent = np.arange(n_v, dtype=np.int64)


def find(x):
    r = x
    while parent[r] != r:
        r = parent[r]
    while parent[x] != r:
        parent[x], x = r, parent[x]
    return r


for a, b in ev:
    ra, rb = find(a), find(b)
    if ra != rb:
        parent[rb] = ra
roots = np.array([find(i) for i in range(n_v)])
_, comp, sizes = np.unique(roots, return_inverse=True, return_counts=True)
big = np.argmax(sizes)
debris = np.nonzero(sizes < SHARD_MAX)[0]
n_debris_v = int(sizes[debris].sum()) if len(debris) else 0
print(f"\ncomponents: {len(sizes)}   largest {sizes[big]} v   "
      f"debris (<{SHARD_MAX} faces worth): {len(debris)} components / {n_debris_v} verts")
slivers = [f for f in bm.faces if f.calc_area() * MM * MM < 0.02]
print(f"slivers (<0.02 mm2): {len(slivers)}")
bm.free()

if not APPLY:
    print("DIAGNOSE ONLY — rerun with --apply to fix")
    print("DESPECKLE_DONE")
    sys.exit(0)

# ---------------------------------------------------------------- fix
# 1. delete debris components outright
if n_debris_v:
    keep = comp == big
    bm = bmesh.new(); bm.from_mesh(me)
    bm.verts.ensure_lookup_table()
    kill = [bm.verts[i] for i in range(n_v) if sizes[comp[i]] < SHARD_MAX]
    bmesh.ops.delete(bm, geom=kill, context='VERTS')
    bm.to_mesh(me)
    bm.free()
    log(f"deleted {len(kill)} debris verts")
    n_v = len(me.vertices)
    co = np.empty(n_v * 3); me.vertices.foreach_get("co", co); co = co.reshape(-1, 3)
    ev = np.empty(len(me.edges) * 2, dtype=np.int32); me.edges.foreach_get("vertices", ev)
    ev = ev.reshape(-1, 2)
    o_ = np.concatenate([ev[:, 0], ev[:, 1]]); n_ = np.concatenate([ev[:, 1], ev[:, 0]])
    srt = np.argsort(o_, kind="stable"); o_s, n_s = o_[srt], n_[srt]
    ptr = np.searchsorted(o_s, np.arange(n_v + 1))
    cnt = np.maximum(np.diff(ptr), 1)
    spike = speckle_mask()
    log(f"speckle mask rebuilt after debris deletion: {int(spike.sum())} verts")

# 2. flatten spikes onto the local surface, clamped, re-measuring each round
co0 = co.copy()
for it in range(6):
    _, sm = deviation(co)
    if it > 0:
        _, sh2 = sharp_verts(SHARP_DEG)
        spike = spike & sh2          # stop touching anything that is no longer folded
    if not spike.any():
        log(f"round {it}: no folded speckle left")
        break
    tgt = co.copy()
    tgt[spike] = sm[spike]
    step = tgt - co
    n_step = np.linalg.norm(step, axis=1)
    lim = (MAX_MOVE_MM / MM)
    too = n_step > lim
    step[too] *= (lim / n_step[too])[:, None]
    co = co + step
    me.vertices.foreach_set("co", co.reshape(-1))
    me.update()
    log(f"round {it}: {int(spike.sum())} spikes, moved p99 "
        f"{np.percentile(n_step[spike], 99)*MM:.3f} mm")

# 3. collapse whatever refused to flatten. A FOLDED flap cannot be smoothed away: its Laplacian
# target is computed from neighbours that include the flap itself, so the target sits inside the
# fold and the iteration stalls (it did — 891 spikes fell to 335 and then stopped). Merging each
# stubborn cluster to a single point removes the fold outright instead of trying to relax it.
_, sh3 = sharp_verts(SHARP_DEG)
if sh3.any():
    adj2 = {}
    for a, b in ev:
        if sh3[a] and sh3[b]:
            adj2.setdefault(a, []).append(b)
            adj2.setdefault(b, []).append(a)
    lab2 = -np.ones(n_v, dtype=np.int64)
    groups = []
    for s in np.nonzero(sh3)[0]:
        if lab2[s] >= 0:
            continue
        stack, cells = [s], []
        lab2[s] = len(groups)
        while stack:
            v = stack.pop()
            cells.append(v)
            for w in adj2.get(v, ()):
                if lab2[w] < 0:
                    lab2[w] = len(groups)
                    stack.append(w)
        groups.append(cells)
    tight = [g for g in groups if len(g) <= CLUSTER_MAX]
    log(f"collapsing {len(tight)} stubborn clusters ({sum(len(g) for g in tight)} verts)")
    # DISSOLVE, do not merge. Collapsing a cluster to a single point leaves a fan vertex in the
    # middle of where the flap was; on a ring that is not planar that vertex is itself a new
    # spike, and measured at 55 deg the merge made folds >70 deg WORSE (361 -> 444). Dissolving
    # removes the offending vertices outright and lets the surrounding ring close over the hole,
    # which is what "remove a dangling flap" actually means.
    bm = bmesh.new(); bm.from_mesh(me)
    bm.verts.ensure_lookup_table()
    victims = [bm.verts[i] for g in tight for i in g]
    victims = [v for v in victims if v.is_valid]
    try:
        bmesh.ops.dissolve_verts(bm, verts=victims)
    except Exception as e:
        log(f"  dissolve_verts failed: {e}")
    bmesh.ops.triangulate(bm, faces=[f for f in bm.faces if len(f.verts) > 3])
    bmesh.ops.dissolve_degenerate(bm, dist=1e-6, edges=bm.edges)
    bm.to_mesh(me)
    bm.free()
    me.update()
    n_v = len(me.vertices)
    co = np.empty(n_v * 3); me.vertices.foreach_get("co", co); co = co.reshape(-1, 3)
    ev = np.empty(len(me.edges) * 2, dtype=np.int32); me.edges.foreach_get("vertices", ev)
    ev = ev.reshape(-1, 2)
    o_ = np.concatenate([ev[:, 0], ev[:, 1]]); n_ = np.concatenate([ev[:, 1], ev[:, 0]])
    srt = np.argsort(o_, kind="stable"); o_s, n_s = o_[srt], n_[srt]
    ptr = np.searchsorted(o_s, np.arange(n_v + 1))
    cnt = np.maximum(np.diff(ptr), 1)
    co0 = co0[:n_v] if len(co0) > n_v else np.pad(co0, ((0, n_v - len(co0)), (0, 0)))
    ang4, sh4 = sharp_verts(SHARP_DEG)
    log(f"after collapse: {n_v}v {len(me.polygons)}f, "
        f"edges folding >{SHARP_DEG:.0f}deg: {int((ang4 > SHARP_DEG).sum())} "
        f"(was {int((ang > SHARP_DEG).sum())})")
    for t in (30, 50, 70, 90):
        print(f"  edges folding > {t}deg: {int((ang4 > t).sum()):6d}  (before {int((ang > t).sum())})")

# Only meaningful if the collapse stage did not renumber the vertices. It used to print
# regardless and reported nonsense (p50 243 mm) by differencing two arrays whose indices no
# longer refer to the same points — a number that looks alarming and means nothing.
if len(co) == len(co0):
    moved = np.linalg.norm(co - co0, axis=1) * MM
    print(f"\nTOTAL MOVEMENT: {int((moved>0.01).sum())} verts moved, "
          f"p50 {np.percentile(moved[moved>0.01], 50) if (moved>0.01).any() else 0:.3f} mm, "
          f"max {moved.max():.3f} mm")
else:
    print(f"\nTOTAL MOVEMENT: not comparable — the collapse stage renumbered vertices "
          f"({len(co0)} -> {len(co)}). Judge from the fold counts above.")
dev2, _ = deviation(co)
d2 = dev2 * MM
print("AFTER:")
for t in (0.5, 1.0, 1.5, 2.0, 3.0):
    print(f"  |dev| > {t:.1f} mm: {int((np.abs(d2)>t).sum()):6d} verts")
print(f"  height unchanged: {co[:,2].max()-co[:,2].min():.5f} units")

bpy.ops.wm.save_as_mainfile(filepath=OUT)
log(f"WROTE {OUT}")
if GLB:
    for o in bpy.data.objects:
        o.select_set(o is ob)
    bpy.context.view_layer.objects.active = ob
    bpy.ops.export_scene.gltf(filepath=os.path.abspath(GLB), export_format='GLB',
                              use_selection=True, export_image_format='AUTO',
                              export_jpeg_quality=95, export_yup=True, export_apply=False)
    log(f"EXPORTED {GLB}")
print("DESPECKLE_DONE")