# Stage 19: measure the seam's cross-section, then trial the heal at several band widths.
#
#   blender --background --python 19_wide_heal.py -- <in.blend> <review_root> [widths]
#
# WHY WIDTH IS THE WHOLE QUESTION. Established so far: the lines are not cracks (16c), not
# painted into the custom normals (18 — corner-vs-vertex deviation is 0.008 deg mean), and a
# 5-vertex-wide collar-fixed membrane moves 12.6k verts without changing the render (17).
# The remaining reading is that a Tripo panel border is a STEP — the reconstruction's two charts
# meet with a sub-millimetre offset, a C0 discontinuity — rather than a ridge sitting on smooth
# skin. A narrow band cannot fix a step, because the fixed collar lands on the step's own
# shoulders and the interpolant faithfully reproduces the offset it is pinned to. Removing a step
# means spreading it over a wide enough neighbourhood that the residual curvature falls below
# visibility.
#
# So this stage MEASURES first: |offset from a broadly smoothed surface| as a function of ring
# distance from the seam. That profile says how wide the disturbance really is, and therefore how
# wide the band must be. Then it renders the heal at several widths so the choice is made from
# pictures rather than from theory. Nothing is saved — this is an experiment; the winning width
# gets applied in the next stage.
import bpy, sys, os, math, time
import numpy as np
from mathutils import Vector

argv = sys.argv[sys.argv.index("--") + 1:]
BLEND, ROOT = argv[0], argv[1]
WIDTHS = [int(x) for x in argv[2].split(",")] if len(argv) > 2 else [4, 8, 12]
t0 = time.time()

UNIT_MM = 1815.0
KINK_DEG = 6.0
COLLAR = 3
Z_LO, Z_HI = 0.04, 0.90
X_MAX = 0.36


def log(m):
    print(f"[wide {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)
ev = np.empty(len(me.edges) * 2, dtype=np.int32)
me.edges.foreach_get("vertices", ev)
ev = ev.reshape(-1, 2)
log(f"{n_v}v {len(me.polygons)}f")

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


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 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 kink_of(P):
    nrm = np.empty(n_v * 3)
    me.vertices.foreach_get("normal", nrm)
    nrm = nrm.reshape(-1, 3)
    N = nrm.copy()
    for _ in range(5):
        N = nbr_mean(N)
        N /= np.maximum(np.linalg.norm(N, axis=1, keepdims=True), 1e-12)
    return np.degrees(np.arccos(np.clip((nrm * N).sum(axis=1), -1, 1))), N


zone = (co[:, 2] > Z_LO) & (co[:, 2] < Z_HI) & (np.abs(co[:, 0]) < X_MAX)
navel = (np.abs(co[:, 0]) < 0.022) & (co[:, 2] > 0.495) & (co[:, 2] < 0.555) & (co[:, 1] < 0)
ang, N = kink_of(co)
seed = zone & ~navel & (ang > KINK_DEG)
log(f"seed (kink>{KINK_DEG}deg): {int(seed.sum())} verts")

# =============================================================================
# cross-section profile: |offset from broad smooth| vs ring distance from seed
# =============================================================================
sm40 = smooth_n(co, 40)
sm12 = smooth_n(co, 12)
dev40 = ((co - sm40) * N).sum(axis=1) * UNIT_MM
dev12 = ((co - sm12) * N).sum(axis=1) * UNIT_MM

ring = np.full(n_v, -1, dtype=np.int32)
ring[seed] = 0
cur = seed.copy()
for r in range(1, 16):
    nxt = grow(cur, 1) & ~cur & zone
    ring[nxt & (ring < 0)] = r
    cur = cur | nxt
print("\nSEAM CROSS-SECTION (real mm, magnitudes; ring 0 = detected seam centre)")
print("  ring   n        |dev12| med   p90      |dev40| med   p90")
for r in range(0, 15):
    m = ring == r
    if m.sum() < 50:
        continue
    print(f"  {r:4d} {int(m.sum()):8d}     {np.median(np.abs(dev12[m])):7.3f} "
          f"{np.percentile(np.abs(dev12[m]),90):7.3f}     "
          f"{np.median(np.abs(dev40[m])):7.3f} {np.percentile(np.abs(dev40[m]),90):7.3f}")
far = zone & (ring < 0)
if far.sum() > 50:
    print(f"  far  {int(far.sum()):8d}     {np.median(np.abs(dev12[far])):7.3f} "
          f"{np.percentile(np.abs(dev12[far]),90):7.3f}     "
          f"{np.median(np.abs(dev40[far])):7.3f} {np.percentile(np.abs(dev40[far]),90):7.3f}")

# =============================================================================
# solver
# =============================================================================
def bilaplacian(P, free_m, collar_rings=COLLAR, maxit=6000):
    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(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] = P[S[~free]]
    rhs = -Ls(Ls(Xc))[free]
    U = P[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
    Q = P.copy()
    Q[S[free]] = U
    return Q, int(free.sum()), it, rs / rs0


# =============================================================================
# render helper (same framing/lighting as 04_review)
# =============================================================================
scn = bpy.context.scene
wd = bpy.data.worlds.new("W")
wd.color = (0.22, 0.22, 0.24)
scn.world = wd
key = bpy.data.objects.new("Key", bpy.data.lights.new("Key", 'SUN'))
key.data.energy = 3.0
key.data.use_shadow = False
bpy.context.collection.objects.link(key)
fl = bpy.data.objects.new("Fill", bpy.data.lights.new("Fill", 'SUN'))
fl.data.energy = 1.0
fl.data.use_shadow = False
bpy.context.collection.objects.link(fl)
cam = bpy.data.objects.new("Cam", bpy.data.cameras.new("Cam"))
cam.data.lens = 85
bpy.context.collection.objects.link(cam)
scn.camera = cam
scn.render.engine = 'BLENDER_EEVEE' if bpy.app.version >= (4, 2) else 'BLENDER_EEVEE_NEXT'
scn.render.resolution_x = scn.render.resolution_y = 1000
clay = bpy.data.materials.new("Clay")
clay.use_nodes = True
clay.node_tree.nodes["Principled BSDF"].inputs["Base Color"].default_value = (0.62, 0.60, 0.58, 1)
clay.node_tree.nodes["Principled BSDF"].inputs["Roughness"].default_value = 0.45
orig = [ms.material for ms in ob.material_slots]


def shoot(outdir, tag, ctr, span, yaw_deg, use_clay=True):
    os.makedirs(outdir, exist_ok=True)
    for i, ms in enumerate(ob.material_slots):
        ms.material = clay if use_clay else orig[i]
    yaw = math.radians(yaw_deg)
    dist = span * 3.0
    cam.location = Vector(ctr) + Vector((math.sin(yaw) * dist, -math.cos(yaw) * dist, 0.02))
    cam.rotation_euler = (Vector(ctr) - cam.location).to_track_quat('-Z', 'Y').to_euler()
    key.rotation_euler = (math.radians(62), 0, math.radians(35 + yaw_deg))
    fl.rotation_euler = (math.radians(75), 0, math.radians(yaw_deg - 110))
    scn.render.filepath = os.path.abspath(os.path.join(outdir, f"{tag}.png"))
    bpy.ops.render.render(write_still=True)


CHEST = (0.0, 0.0, 0.675)
FULL = (0.0, 0.0, 0.50)
HIP = (0.0, 0.0, 0.53)

for W in WIDTHS:
    band = grow(seed, W) & zone & ~navel
    Q, nf, it, rel = bilaplacian(co, band)
    d = np.linalg.norm(Q - co, axis=1) * UNIT_MM
    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))
    ang2, _ = kink_of(Q)
    torso = (co[:, 2] > 0.28) & (co[:, 2] < 0.90)
    log(f"W={W:2d}: band {nf} verts, CG it={it} rel={rel:.1e}, moved max {d.max():.2f} mm "
        f"median(band) {np.median(d[band]):.3f} mm | kink>6 {int((torso&(ang2>6)).sum())} "
        f">12 {int((torso&(ang2>12)).sum())} >20 {int((torso&(ang2>20)).sum())}")
    out = os.path.join(ROOT, f"w{W:02d}")
    shoot(out, "chest_clay_40", CHEST, 0.22, 40)
    shoot(out, "full_clay_0", FULL, 0.55, 0)
    shoot(out, "hip_clay_0", HIP, 0.22, 0)
    log(f"W={W}: rendered -> {out}")
    me.vertices.foreach_set("co", co.reshape(-1))      # reset for the next width
    me.update()

print("WIDE_DONE")