# Stage 29: level the repainted patches' tone using ON-BODY sampling (the only kind that works).
#
#   blender --background --python 29_tone3d.py -- <in.blend> <out.blend> [orig.blend] [strength]
#
# WHY STAGE 21 FAILED AND WAS REJECTED (REJECTED_28_atlas_tone.blend)
# Stage 21 solved a Poisson correction whose boundary condition was the mismatch against texels
# ADJACENT IN THE ATLAS. That is the one mistake 05_texture.py's header explicitly warns about:
# "fill tone comes from skin NEAREST ON THE BODY (KD over skin verts in 3D), never from atlas
# neighbourhoods — atlas-local fills gave wrong tones and island seams". UV adjacency is not body
# adjacency: a patch border texel's atlas neighbours are frequently a different body part or empty
# gutter, so the boundary mismatch was garbage and the harmonic solve spread it over the whole
# patch. Result: the bra and briefs read as pale grey panels, far worse than the faint tone step we
# started from. The roughness pass compounded it (corrections up to +0.19 turned her plasticky).
#
# THIS STAGE DOES IT ON THE MESH.
#   * a vertex is "garment" if its texel was repainted (mask = current vs pristine basecolor);
#   * for each garment vertex, the target tone is an inverse-distance average of the 8 nearest
#     SKIN vertices in 3D — real neighbours on the body, across UV seams, never a gutter;
#   * correction = target - current, per vertex, then SMOOTHED over the mesh graph so only the
#     low-frequency level is carried and the transplanted grain survives untouched;
#   * the smoothed correction is rasterised over garment faces and added.
# Because the correction is smooth and vanishes where current tone already equals nearby skin, a
# well-matched region is left alone and only a genuine offset is removed.
#
# STRENGTH is deliberately conservative (default 0.75): the baseline is already close, and the
# failure mode of this whole family of fixes is overshoot.
import bpy, sys, os, time
import numpy as np
from mathutils import Vector
from mathutils.kdtree import KDTree

argv = sys.argv[sys.argv.index("--") + 1:]
BLEND, OUT = argv[0], argv[1]
ORIG_BLEND = argv[2] if len(argv) > 2 else "00_welded.blend"
STRENGTH = float(argv[3]) if len(argv) > 3 else 0.75
SMOOTH_ARG = int(argv[4]) if len(argv) > 4 else None
t0 = time.time()

DIFF_T = 0.02
SMOOTH_CORR = 120        # graph-smoothing passes on the correction field (low-frequency only)
KNN = 8
Z_LO, Z_HI = 0.25, 0.895     # sample skin on the body, never the face/lips/eyes
FEATHER = 3


def log(m):
    print(f"[t3d {time.time()-t0:6.1f}s] {m}", flush=True)


def getpx(img):
    w, h = img.size
    b = np.empty(w * h * 4, dtype=np.float32)
    img.pixels.foreach_get(b)
    return b.reshape(h, w, 4)


# ---- pristine originals ----
bpy.ops.wm.open_mainfile(filepath=ORIG_BLEND)
CACHE = {}
for i in bpy.data.images:
    nm = i.name.lower()
    k = "base" if "basecolor" in nm else ("rm" if "_rm" in nm else None)
    if k and k not in CACHE:
        CACHE[k] = (getpx(i)[:, :, :3].astype(np.float32), tuple(i.size))
log(f"cached pristine: { {k: v[1] for k, v in CACHE.items()} }")

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)
wired = {}
for ms in ob.material_slots:
    if not ms.material or not ms.material.node_tree:
        continue
    for n in ms.material.node_tree.nodes:
        if n.type != 'TEX_IMAGE' or not n.image:
            continue
        for o in n.outputs:
            for lk in o.links:
                tn = lk.to_node.name.lower()
                if "principled" in tn:
                    wired["base"] = n.image
                elif "separate" in tn:
                    wired["rm"] = n.image
log(f"body {n_v}v, wired { {k: v.name for k, v in wired.items()} }")
if "base" not in wired:
    raise SystemExit("[t3d] FATAL: no wired basecolor")

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)

base = wired["base"]
B4 = getpx(base)
B = B4[:, :, :3].astype(np.float64)
h, w = B.shape[:2]
Orig = CACHE["base"][0]
mask = np.abs(B - Orig).max(axis=2) > DIFF_T
log(f"repainted texels: {int(mask.sum())} ({100.0*mask.sum()/(w*h):.2f}%)")

# ---- per-vertex UV -> texel, colour, and garment flag ----
lv = np.empty(len(me.loops), dtype=np.int32)
me.loops.foreach_get("vertex_index", lv)
uv = np.empty(len(me.loops) * 2)
me.uv_layers.active.data.foreach_get("uv", uv)
uv = uv.reshape(-1, 2)
px = np.clip(uv[:, 0], 0, 1) * (w - 1)
py = np.clip(uv[:, 1], 0, 1) * (h - 1)
pxi = px.astype(np.int32)
pyi = py.astype(np.int32)

first = np.full(n_v, -1, dtype=np.int64)
np.maximum.at(first, lv, np.arange(len(lv), dtype=np.int64))
has = first >= 0
vcol = np.zeros((n_v, 3))
vcol[has] = B[pyi[first[has]], pxi[first[has]]]
# a vertex is garment if ANY of its loops lands on a repainted texel
gv = np.zeros(n_v, dtype=bool)
np.logical_or.at(gv, lv, mask[pyi, pxi])
log(f"garment verts: {int(gv.sum())} of {n_v}")

zone = (co[:, 2] > Z_LO) & (co[:, 2] < Z_HI)
skin = (~gv) & zone & has
skin_idx = np.nonzero(skin)[0][::4]
log(f"skin sample for KD: {len(skin_idx)}")
kd = KDTree(len(skin_idx))
for j, i in enumerate(skin_idx):
    kd.insert(Vector(co[i]), j)
kd.balance()

gidx = np.nonzero(gv & zone & has)[0]
target = np.zeros((n_v, 3))
for i in gidx:
    hits = kd.find_n(Vector(co[i]), KNN)
    wsum = 0.0
    acc = np.zeros(3)
    for (_, j, d) in hits:
        wt = 1.0 / max(d * d, 1e-9)
        acc += wt * vcol[skin_idx[j]]
        wsum += wt
    target[i] = acc / wsum
log("on-body target tones computed")

corr = np.zeros((n_v, 3))
corr[gidx] = target[gidx] - vcol[gidx]
pre = np.abs(corr[gidx]).mean(axis=0)
log(f"raw correction magnitude per channel: {pre.round(4)}")

# ---- smooth the correction over the mesh graph: keep only the level, not the detail ----
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
if SMOOTH_ARG is not None:
    SMOOTH_CORR = SMOOTH_ARG
for _ in range(SMOOTH_CORR):
    a = np.add.reduceat(corr[n_s], ptr[:-1], axis=0)
    a[empty] = corr[empty]
    corr = a / cnt[:, None]
corr *= STRENGTH
log(f"smoothed x{SMOOTH_CORR}, strength {STRENGTH}: "
    f"mean |corr| on garment {np.abs(corr[gidx]).mean(axis=0).round(4)}, "
    f"max {np.abs(corr[gidx]).max():.4f}")

# ---- on-body verification metric: garment interior tone vs the skin ring around it, in 3D ----
def grow(m, k):
    x = m.copy()
    for _ in range(k):
        hit = x[ev[:, 0]] | x[ev[:, 1]]
        y = x.copy()
        y[ev[:, 0]] |= hit
        y[ev[:, 1]] |= hit
        x = y
    return x


inner_v = gv & ~grow(~gv, 6)
ring_v = grow(gv, 8) & ~gv & zone
if inner_v.sum() and ring_v.sum():
    a0 = vcol[inner_v].mean(axis=0)
    r0 = vcol[ring_v].mean(axis=0)
    a1 = (vcol[inner_v] + corr[inner_v]).mean(axis=0)
    log(f"ON-BODY tone gap (garment interior - surrounding skin):")
    log(f"   before {(a0-r0).round(4)}  |gap| luma {abs(a0.mean()-r0.mean()):.4f}")
    log(f"   after  {(a1-r0).round(4)}  |gap| luma {abs(a1.mean()-r0.mean()):.4f}")

# ---- rasterise the correction over garment faces ----
n_f = len(me.polygons)
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)
face_g = np.add.reduceat(gv[lv].astype(np.int32), ls) > 0
log(f"garment faces: {int(face_g.sum())}")

out = B.copy()
painted = np.zeros((h, w), dtype=bool)
for fi in np.nonzero(face_g)[0]:
    s, t = int(ls[fi]), int(lt[fi])
    li = np.arange(s, min(s + t, s + 3))
    P = np.stack([px[li], py[li]], axis=1)
    V = lv[li]
    if len(P) < 3:
        continue
    x0, x1 = int(P[:, 0].min()), int(np.ceil(P[:, 0].max()))
    y0, y1 = int(P[:, 1].min()), int(np.ceil(P[:, 1].max()))
    if x1 - x0 > 128 or y1 - y0 > 128 or x1 < x0 or y1 < y0:
        continue
    d = ((P[1, 1] - P[2, 1]) * (P[0, 0] - P[2, 0]) +
         (P[2, 0] - P[1, 0]) * (P[0, 1] - P[2, 1]))
    if abs(d) < 1e-12:
        continue
    gx, gy = np.meshgrid(np.arange(x0, min(x1, w - 1) + 1),
                         np.arange(y0, min(y1, h - 1) + 1))
    a = ((P[1, 1] - P[2, 1]) * (gx - P[2, 0]) + (P[2, 0] - P[1, 0]) * (gy - P[2, 1])) / d
    b = ((P[2, 1] - P[0, 1]) * (gx - P[2, 0]) + (P[0, 0] - P[2, 0]) * (gy - P[2, 1])) / d
    c = 1.0 - a - b
    ins = (a >= -0.02) & (b >= -0.02) & (c >= -0.02)
    if not ins.any():
        continue
    e = (a[ins, None] * corr[V[0]] + b[ins, None] * corr[V[1]] + c[ins, None] * corr[V[2]])
    out[gy[ins], gx[ins]] += e
    painted[gy[ins], gx[ins]] = True
log(f"rasterised correction over {int(painted.sum())} texels")

# only correct where the atlas was actually repainted, feathered at the rim
apply_m = painted & mask
alpha = apply_m.astype(np.float64)
edge = apply_m.copy()
for k in range(FEATHER):
    g = edge.copy()
    g[1:, :] |= edge[:-1, :]
    g[:-1, :] |= edge[1:, :]
    g[:, 1:] |= edge[:, :-1]
    g[:, :-1] |= edge[:, 1:]
    ring = g & ~edge
    alpha[ring] = 1.0 - (k + 1) / (FEATHER + 1.0)
    edge = g
final = np.clip(B * (1 - alpha[:, :, None]) + out * alpha[:, :, None], 0, 1)
log(f"applied to {int(apply_m.sum())} texels; "
    f"mean shift {np.abs(final-B)[apply_m].mean():.5f}, max {np.abs(final-B).max():.4f}")
B4[:, :, :3] = final.astype(np.float32)
base.pixels.foreach_set(B4.reshape(-1))
base.pack()
log("basecolor written + packed")

bpy.ops.wm.save_as_mainfile(filepath=OUT)
log(f"WROTE {OUT}")
print("T3D_DONE")