# Stage 26: the right operator for the lines (ring-median despeckle) + a real cleavage fillet.
#
#   blender --background --python 26_finish.py -- <in.blend> <out.blend> [iters] [alpha]
#
# PART A — LINES, with a median filter instead of a membrane.
# 19_wide_heal measured the seam cross-section: the relief is ONE VERTEX WIDE (|offset| 0.23 mm
# median at the centre and already back to the 0.074 mm background by ring 2). That measurement
# picks the operator, and it is not a membrane:
#   * a collar-fixed membrane needs clean ground to stand on. With a 1-vertex defect the collar
#     lands ON the defect's shoulders, so the interpolant faithfully reproduces what it is pinned
#     to — which is why stages 11/14/17 moved 12k+ verts for almost no visible change, and why
#     widening the band (GROW 4, 8) made it worse rather than better.
#   * a MEDIAN over the 1-ring is exact for this defect class. Take the scalar offset of each
#     vertex from the locally smooth surface and replace it with the median of its neighbourhood:
#     an isolated 1-vertex ridge or groove is a rank outlier and is deleted completely, while a
#     feature many vertices wide has offset ~= its own median and is returned UNCHANGED. So the
#     underbust fold, the clavicles, the navel and the gluteal fold survive by construction
#     rather than by a hand-tuned threshold — which is what every earlier stage got wrong.
# Only the normal component is filtered, so no vertex slides tangentially and the UV atlas stays
# valid. Displacement is clamped, as in stage 17, so nothing can reshape her.
#
# PART B — CLEAVAGE, with a blended membrane in the medial corridor.
# 20_cleavage's fill-only curvature filter took the sternum radius from 9.5 mm to only 15.8 mm and
# left the notch depth at 44 mm, because a uniform-Laplacian curvature test measures sharpness at
# the ~2 mm vertex scale and the defect is a 44 mm-deep macro V. Fixing macro shape needs an
# operator with macro reach: a bi-harmonic membrane across the corridor, which interpolates the
# cups' own slopes into a smooth valley. It is applied at a fraction alpha so the cleavage is
# rounded rather than erased — alpha=1 would bridge the cups into a web.
#
# NOTE ON A BUG THIS FIXES: 20_cleavage.py (and 18/25) render clay LAST and then save, so the
# saved .blend keeps the clay material in the slots and loses the texture wiring — that is why
# 21_tone.py aborted with "could not find the wired basecolor" on 23_cleavage.blend. Here the
# original materials are restored and the file is saved BEFORE any render.
import bpy, sys, time
import numpy as np

argv = sys.argv[sys.argv.index("--") + 1:]
BLEND, OUT = argv[0], argv[1]
ITERS = int(argv[2]) if len(argv) > 2 else 3
ALPHA = float(argv[3]) if len(argv) > 3 else 0.65
t0 = time.time()

UNIT_MM = 1815.0
SMOOTH_K = 6            # scale the offset is measured against
CLAMP_MM = 2.0          # lines are <=1.5 mm of relief
Z_LO, Z_HI = 0.04, 0.90
X_MAX = 0.36
# cleavage corridor
CZ0, CZ1 = 0.690, 0.782
CX = 0.028
COLLAR = 3
CLV_CLAMP_MM = 26.0     # the notch is 44 mm deep; allow a real fillet but not a bridge


def log(m):
    print(f"[fin {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)
orig_mats = [ms.material for ms in ob.material_slots]
log(f"in: {n_v}v {len(me.polygons)}f mats={[m.name if m else None for m in orig_mats]}")

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)

order = np.concatenate([ev[:, 0], ev[:, 1]])
nbr = np.concatenate([ev[:, 1], ev[:, 0]])
srt = np.argsort(order, kind="stable")
o_s, n_s = order[srt], nbr[srt]
ptr = np.searchsorted(o_s, np.arange(n_v + 1))
cnt = np.maximum(np.diff(ptr), 1)
empty = np.diff(ptr) == 0
MAXV = min(int(np.diff(ptr).max()), 24)
log(f"valence: mean {np.diff(ptr).mean():.2f}, max {np.diff(ptr).max()} (using {MAXV})")

# padded neighbour index table for the ring median
cols = np.arange(MAXV)[None, :]
base = ptr[:-1][:, None]
lim = ptr[1:][:, None]
idx = np.minimum(base + cols, max(len(n_s) - 1, 0))
valid = (base + cols) < lim
nb_tab = n_s[idx]


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


def smooth_n(X, k):
    Y = X.copy()
    for _ in range(k):
        Y = nbr_mean(Y)
    return Y


def ring_median(d):
    """Median of each vertex's 1-ring plus itself. Invalid slots -> nan, nanmedian ignores them."""
    T = np.where(valid, d[nb_tab], np.nan)
    T = np.concatenate([T, d[:, None]], axis=1)
    return np.nanmedian(T, axis=1)


def vnormals():
    nrm = np.empty(n_v * 3)
    me.vertices.foreach_get("normal", nrm)
    return nrm.reshape(-1, 3)


zone = (co[:, 2] > Z_LO) & (co[:, 2] < Z_HI) & (np.abs(co[:, 0]) < X_MAX)
log(f"zone: {int(zone.sum())} verts")


def kink_stats(P):
    nrm = vnormals()
    N = nrm.copy()
    for _ in range(5):
        N = nbr_mean(N)
        N /= np.maximum(np.linalg.norm(N, axis=1, keepdims=True), 1e-12)
    ang = np.degrees(np.arccos(np.clip((nrm * N).sum(axis=1), -1, 1)))
    t = (P[:, 2] > 0.28) & (P[:, 2] < 0.90)
    return [int((t & (ang > k)).sum()) for k in (6, 12, 20)]


log(f"kink [>6,>12,>20] before: {kink_stats(co)}")

# =============================================================================
# PART A — ring-median despeckle
# =============================================================================
P = co.copy()
for it in range(1, ITERS + 1):
    me.vertices.foreach_set("co", P.reshape(-1))
    me.update()
    N = vnormals()
    S = smooth_n(P, SMOOTH_K)
    d = ((P - S) * N).sum(axis=1)
    dm = ring_median(d)
    delta = np.where(zone, dm - d, 0.0)
    dmm = np.abs(delta) * UNIT_MM
    over = dmm > CLAMP_MM
    if over.any():
        delta[over] = np.sign(delta[over]) * (CLAMP_MM / UNIT_MM)
    Q = P + N * delta[:, None]
    mv = np.abs(delta) * UNIT_MM
    log(f"median pass {it}: {int((mv>0.01).sum())} verts adjusted, "
        f"max {mv.max():.3f} mm, median(adjusted) "
        f"{np.median(mv[mv>0.01]) if (mv>0.01).any() else 0:.3f} mm, clamped {int(over.sum())}")
    P = Q

me.vertices.foreach_set("co", P.reshape(-1))
me.update()
log(f"kink [>6,>12,>20] after median: {kink_stats(P)}")
dA = np.linalg.norm(P - co, axis=1) * UNIT_MM
log(f"part A displacement: max {dA.max():.3f} mm, p99 {np.percentile(dA,99):.3f} mm")

# =============================================================================
# PART B — cleavage fillet
# =============================================================================
def grow(mask, rings):
    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


def bilaplacian(X, free_m, collar_rings=COLLAR, maxit=8000):
    collar = grow(free_m, collar_rings) & ~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 = ev[in_S[ev].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]

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

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

    Xc = np.zeros((len(S), 3))
    Xc[~free] = X[S[~free]]
    rhs = -Ls(Ls(Xc))[free]
    U = X[S[free]].copy()
    r = rhs - A_op(U)
    p = r.copy()
    rs = (r * r).sum()
    rs0 = max(rs, 1e-30)
    it = 0
    for it in range(maxit):
        Ap = A_op(p)
        den = (p * Ap).sum()
        if abs(den) < 1e-30:
            break
        al = rs / den
        U += al * p
        r -= al * Ap
        rs2 = (r * r).sum()
        if rs2 < 1e-20 or rs2 < rs0 * 1e-13:
            rs = rs2
            break
        p = r + (rs2 / rs) * p
        rs = rs2
    Y = X.copy()
    Y[S[free]] = U
    return Y, int(free.sum()), it, rs / rs0


def probe(X, tag):
    front = X[:, 1] < 0
    print(f"\n=== CLEAVAGE PROFILE [{tag}] ===")
    print("   z      sternum y    fillet radius     notch vs apex")
    for z0 in np.arange(0.650, 0.7801, 0.015):
        row = []
        for x0 in np.arange(-0.05, 0.0501, 0.005):
            m = front & (np.abs(X[:, 0] - x0) < 0.0035) & (np.abs(X[:, 2] - z0) < 0.004)
            row.append(X[m, 1].min() if m.sum() else np.nan)
        row = np.array(row)
        if np.isnan(row).all():
            continue
        mid = len(row) // 2
        seg = row[max(0, mid - 3):mid + 4]
        rad = float('inf')
        if len(seg) >= 3 and not np.isnan(seg).any():
            d2 = (seg[:-2] - 2 * seg[1:-1] + seg[2:]) / (0.005 ** 2)
            k = float(np.nanmax(d2))
            rad = 1.0 / k if k > 1e-6 else float('inf')
        ma = front & (np.abs(np.abs(X[:, 0]) - 0.034) < 0.005) & (np.abs(X[:, 2] - z0) < 0.004)
        notch = ((row[mid] - X[ma, 1].min()) * UNIT_MM) if ma.sum() else np.nan
        print(f"  {z0:.3f}   {row[mid]:+.4f}     {rad*UNIT_MM:8.1f} mm     {notch:+7.1f} mm"
              + ("   <-- SHARP" if rad * UNIT_MM < 30 else ""))


probe(P, "before fillet")
corridor = (P[:, 1] < 0) & (np.abs(P[:, 0]) < CX) & (P[:, 2] > CZ0) & (P[:, 2] < CZ1)
log(f"corridor: {int(corridor.sum())} verts")
Q, nf, it, rel = bilaplacian(P, corridor)
disp = (Q - P) * ALPHA
dmag = np.linalg.norm(disp, axis=1) * UNIT_MM
ov = dmag > CLV_CLAMP_MM
if ov.any():
    disp[ov] *= (CLV_CLAMP_MM / dmag[ov])[:, None]
P2 = P + disp
d = np.linalg.norm(P2 - P, axis=1) * UNIT_MM
log(f"fillet: {nf} verts, CG it={it} rel={rel:.1e}, alpha={ALPHA}, "
    f"moved max {d.max():.2f} mm, median(corridor) {np.median(d[corridor]):.2f} mm, "
    f"clamped {int(ov.sum())}")

me.vertices.foreach_set("co", P2.reshape(-1))
me.update()
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))
probe(P2, "after fillet")
log(f"kink [>6,>12,>20] final: {kink_stats(P2)}")
dT = np.linalg.norm(P2 - co, axis=1) * UNIT_MM
log(f"TOTAL displacement: max {dT.max():.2f} mm, p99 {np.percentile(dT,99):.3f} mm, "
    f"{int((dT>0.05).sum())} verts moved >0.05 mm")

# restore materials BEFORE saving (see header note) and save
for i, ms in enumerate(ob.material_slots):
    ms.material = orig_mats[i]
bpy.ops.wm.save_as_mainfile(filepath=OUT)
log(f"WROTE {OUT} with materials {[m.name if m else None for m in orig_mats]}")
print("FIN_DONE")