# Stage 40: the OTHER half of the speckle. 39_despeckle.py cut the geometric folds (892 -> 361
# edges over 70 deg) and the clay render is now clean, but the textured render still shows flecks:
# they are painted into the albedo — scan noise from the Tripo source, faithfully resampled into
# the new atlas. Clay vs textured is what separates the two, and both had to be fixed.
#
#   blender --background --python 40_texspeckle.py -- <in.blend> <out.blend> [out.glb]
#
# Detection is a local-contrast outlier test: a fleck is a small blob that is much darker (or
# brighter) than the skin immediately around it. Two guards keep it from eating her:
#   - BLOB SIZE. Only components under MAX_BLOB texels are touched, so lips, brows, eyes and
#     nostrils — which are large, coherent regions — are never candidates.
#   - THE HEAD IS EXCLUDED OUTRIGHT. Her face is legitimately high-contrast and is the one place
#     a contrast test cannot be trusted. Faces above the neck landmark are rasterised into a
#     protect mask first.
import bpy, sys, os, time
import numpy as np

argv = sys.argv[sys.argv.index("--") + 1:]
BLEND, OUT = argv[0], argv[1]
GLB = next((a for a in argv[2:] if a.lower().endswith(".glb")), "")
t0 = time.time()
THRESH = 0.040        # local-contrast deviation that counts as a fleck
MAX_BLOB = 400        # texels; above this it is a feature, not a fleck
NECK_U = 0.825        # same landmark 24_seams.py measured


def log(m):
    print(f"[tspk {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_l, n_f = len(me.vertices), len(me.loops), len(me.polygons)

base = None
for slot in ob.material_slots:
    if not slot.material or not slot.material.node_tree:
        continue
    for node in slot.material.node_tree.nodes:
        if node.type == 'BSDF_PRINCIPLED' and node.inputs["Base Color"].links:
            nd = node.inputs["Base Color"].links[0].from_node
            if nd.type == 'TEX_IMAGE':
                base = nd.image
W, H = base.size
buf = np.empty(W * H * 4, dtype=np.float32)
base.pixels.foreach_get(buf)
tex = buf.reshape(H, W, 4)
rgb = tex[:, :, :3].astype(np.float64)
log(f"atlas '{base.name}' {W}x{H}")

# ---- protect the head, and know where the atlas actually has skin ----
co = np.empty(n_v * 3); me.vertices.foreach_get("co", co); co = co.reshape(-1, 3)
u = (co[:, 2] - co[:, 2].min()) / (co[:, 2].max() - co[:, 2].min())
loops_v = np.empty(n_l, dtype=np.int32); me.loops.foreach_get("vertex_index", loops_v)
uv = np.empty(n_l * 2); me.uv_layers.active.data.foreach_get("uv", uv); uv = uv.reshape(-1, 2)
ls = np.empty(n_f, dtype=np.int32); me.polygons.foreach_get("loop_start", ls)
lt = np.empty(n_f, dtype=np.int32); me.polygons.foreach_get("loop_total", lt)
li = ls[lt == 3]
IDX = np.stack([li, li + 1, li + 2], axis=1)
V = loops_v[IDX]
P = np.stack([np.clip(uv[IDX][:, :, 0], 0, 1) * (W - 1),
              np.clip(uv[IDX][:, :, 1], 0, 1) * (H - 1)], axis=2)

protect = np.zeros((H, W), dtype=bool)
covered = np.zeros((H, W), dtype=bool)
is_head = u > NECK_U
for f in range(len(P)):
    p3 = P[f]
    x0, x1 = int(p3[:, 0].min()), int(np.ceil(p3[:, 0].max()))
    y0, y1 = int(p3[:, 1].min()), int(np.ceil(p3[:, 1].max()))
    if x1 < x0 or y1 < y0 or x1 - x0 > 512 or y1 - y0 > 512:
        continue
    det = ((p3[1, 1] - p3[2, 1]) * (p3[0, 0] - p3[2, 0])
           + (p3[2, 0] - p3[1, 0]) * (p3[0, 1] - p3[2, 1]))
    if abs(det) < 1e-12:
        continue
    gx, gy = np.meshgrid(np.arange(max(x0, 0), min(x1, W - 1) + 1),
                         np.arange(max(y0, 0), min(y1, H - 1) + 1))
    if gx.size == 0:
        continue
    a = ((p3[1, 1] - p3[2, 1]) * (gx - p3[2, 0]) + (p3[2, 0] - p3[1, 0]) * (gy - p3[2, 1])) / det
    b = ((p3[2, 1] - p3[0, 1]) * (gx - p3[2, 0]) + (p3[0, 0] - p3[2, 0]) * (gy - p3[2, 1])) / det
    c = 1.0 - a - b
    ins = (a >= -0.02) & (b >= -0.02) & (c >= -0.02)
    if not ins.any():
        continue
    covered[gy[ins], gx[ins]] = True
    if is_head[V[f]].any():
        protect[gy[ins], gx[ins]] = True
log(f"atlas coverage {100.0*covered.mean():.1f}%, head protected {100.0*protect.mean():.1f}%")


def box(a, r):
    def b1(x, ax):
        pad = [(0, 0)] * x.ndim
        pad[ax] = (r, r)
        cs = np.cumsum(np.pad(x, pad, mode="edge"), axis=ax)
        return (np.take(cs, np.arange(2 * r, cs.shape[ax]), axis=ax)
                - np.take(cs, np.arange(0, cs.shape[ax] - 2 * r), axis=ax)) / (2 * r)
    return b1(b1(a, 0), 1)


lum = rgb.mean(axis=2)
bg = box(lum, 6)
dev = bg - lum
cand = covered & ~protect & (np.abs(dev) > THRESH)
log(f"local-contrast outliers: {int(cand.sum())} texels ({100.0*cand.mean():.3f}%)")

# blob-size filter: keep only small ones
lab = np.zeros((H, W), dtype=np.int32)
seen = np.zeros((H, W), dtype=bool)
from collections import deque
spots = np.zeros((H, W), dtype=bool)
ys, xs = np.nonzero(cand)
big_kept = 0
for y0, x0 in zip(ys, xs):
    if seen[y0, x0]:
        continue
    q = deque([(y0, x0)])
    seen[y0, x0] = True
    cells = []
    while q:
        y, x = q.popleft()
        cells.append((y, x))
        if len(cells) > MAX_BLOB:
            break
        for dy, dx in ((1, 0), (-1, 0), (0, 1), (0, -1)):
            yy, xx = y + dy, x + dx
            if 0 <= yy < H and 0 <= xx < W and cand[yy, xx] and not seen[yy, xx]:
                seen[yy, xx] = True
                q.append((yy, xx))
    while q:
        y, x = q.popleft()
        seen[y, x] = True
    if len(cells) <= MAX_BLOB:
        for y, x in cells:
            spots[y, x] = True
    else:
        big_kept += 1
log(f"flecks (blobs <= {MAX_BLOB} px): {int(spots.sum())} texels; "
    f"{big_kept} larger regions left alone as features")

# grow slightly so the fleck's soft edge goes too
g = spots.copy()
for _ in range(2):
    n2 = g.copy()
    n2[1:, :] |= g[:-1, :]; n2[:-1, :] |= g[1:, :]
    n2[:, 1:] |= g[:, :-1]; n2[:, :-1] |= g[:, 1:]
    g = n2
spots = g & covered & ~protect
log(f"after grow: {int(spots.sum())} texels ({100.0*spots.mean():.3f}% of the atlas)")

# ---- inpaint: neighbour-average diffusion from clean skin around each fleck ----
C = rgb.copy()
have = covered & ~spots
for _ in range(24):
    todo = spots & ~have
    if not todo.any():
        break
    Wf = have.astype(np.float64)
    acc = np.zeros_like(C)
    wac = np.zeros((H, W))
    for dy, dx in ((1, 0), (-1, 0), (0, 1), (0, -1)):
        acc += np.roll(C * Wf[:, :, None], (dy, dx), axis=(0, 1))
        wac += np.roll(Wf, (dy, dx), axis=(0, 1))
    new = todo & (wac > 0)
    if not new.any():
        break
    C[new] = acc[new] / wac[new, None]
    have |= new
# a couple of smoothing passes confined to the repaired texels, so the patch is not blocky
for _ in range(2):
    sm = np.stack([box(C[:, :, c], 2) for c in range(3)], axis=2)
    C[spots] = sm[spots]
log(f"inpainted {int((spots & have).sum())} texels")

b4 = tex.copy()
b4[:, :, :3] = np.clip(C, 0, 1).astype(np.float32)
base.pixels.foreach_set(b4.reshape(-1))
base.pack()
# name the loose sidecar after the output blend: a fixed name made every run overwrite the
# previous version's copy and left both blends pointing at the same path
stem = os.path.splitext(os.path.basename(OUT))[0]
p = os.path.join(os.path.dirname(os.path.abspath(OUT)), f"{stem}_base.jpg")
base.file_format = 'JPEG'
base.filepath_raw = p
base.save(filepath=p)
log(f"wrote {p}")

bpy.ops.wm.save_as_mainfile(filepath=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("TEXSPECKLE_DONE")