# Stage 5: repaint the garment out of the textures (basecolor + normal + roughness/metallic),
# on the sculpted hires body.
#
#   blender --background --python 05_texture.py -- <06_final.blend> <masks.npz>
#            <00_welded.blend> <out.blend>
#
# The repaint mask is rebuilt from the ORIGINAL geometry (00_welded) because fabric verts are
# no longer rough after the sculpt: garment ∪ hemband ∪ key_raw ∪ rough strap corridor ∪ the
# bow-excision window ∪ the crotch box — the union of every region whose geometry was
# replaced, whose paint is fabric.
#
# Same principles that fixed the game-density bake, at hires scale:
#   - 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;
#   - grain is transplanted from real skin tiles so the fill is not a smooth decal
#     (per-channel high-pass — the channel-mixing blur bug is not repeated here);
#   - normal map goes flat (128,128,255) and rm matches median skin over the same texels,
#     so the fabric weave stops shading through after the paint is gone.
import bpy, sys, time
import numpy as np
from mathutils import Vector
from mathutils.kdtree import KDTree

argv = sys.argv[sys.argv.index("--") + 1:]
BLEND, MASKS, WELDED, OUT = argv[0], argv[1], argv[2], argv[3]
# 'patch' mode (v02 variant): repaint ONLY the garment texels, keep every other texel of the
# original skin — no rosy-chest mask, fill sources include the rosy skin so the patches blend
# with HER tones, and the fill is mirror-averaged so one side's blush can't splash one cup.
PATCH_ONLY = len(argv) > 4 and argv[4] == "patch"
t0 = time.time()
GRAIN_T = 16
FEATHER = 4


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


# ---- build the repaint mask from ORIGINAL geometry ----
bpy.ops.wm.open_mainfile(filepath=WELDED)
ob0 = max([o for o in bpy.data.objects if o.type == 'MESH'],
          key=lambda o: len(o.data.vertices))
me0 = ob0.data
n_v = len(me0.vertices)
co0 = np.empty(n_v * 3, dtype=np.float64)
me0.vertices.foreach_get("co", co0)
co0 = co0.reshape(-1, 3)

ev0 = np.empty(len(me0.edges) * 2, dtype=np.int32)
me0.edges.foreach_get("vertices", ev0)
ev0 = ev0.reshape(-1, 2)
o_r = np.concatenate([ev0[:, 0], ev0[:, 1]])
n_r = np.concatenate([ev0[:, 1], ev0[:, 0]])
s_r = np.argsort(o_r, kind="stable")
o_rs = o_r[s_r]
n_rs = n_r[s_r]
ptr_r = np.searchsorted(o_rs, np.arange(n_v + 1))
cnt_r = np.maximum(np.diff(ptr_r), 1)
sm_r = co0.copy()
for _ in range(8):
    su = np.add.reduceat(sm_r[n_rs], ptr_r[:-1], axis=0)
    emp = np.diff(ptr_r) == 0
    su[emp] = sm_r[emp]
    sm_r = su / cnt_r[:, None]
rough0 = np.linalg.norm(co0 - sm_r, axis=1)

M = np.load(MASKS)
corridor0 = (co0[:, 2] > 0.64) & (co0[:, 2] < 0.86) \
            & (np.abs(co0[:, 0]) > 0.02) & (np.abs(co0[:, 0]) < 0.145)
strap_rough = corridor0 & (rough0 > 0.0005)
bow_win = (co0[:, 1] < 0) & (np.abs(co0[:, 0]) < 0.080) \
          & (co0[:, 2] > 0.630) & (co0[:, 2] < 0.802)
crotch_box = (np.abs(co0[:, 0]) < 0.075) & (co0[:, 2] > 0.340) & (co0[:, 2] < 0.480)
garment = (M["garment"] | M["hemband"] | M["key_raw"]
           | strap_rough | bow_win | crotch_box)
for _ in range(3):     # grow so the rasterised fill overlaps every replaced-geometry rim
    hit = garment[ev0[:, 0]] | garment[ev0[:, 1]]
    g2 = garment.copy()
    g2[ev0[:, 0]] |= hit
    g2[ev0[:, 1]] |= hit
    garment = g2
log(f"repaint mask: {garment.sum()} of {n_v} "
    f"(strap_rough {strap_rough.sum()}, bow {bow_win.sum()}, crotch {crotch_box.sum()})")

# ---- now open the FINAL sculpted body and repaint on it ----
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
assert len(me.vertices) == n_v, "vertex count changed between welded and final"
log(f"loaded {n_v}v, garment {garment.sum()}")

co = np.empty(n_v * 3, dtype=np.float64)
me.vertices.foreach_get("co", co)
co = co.reshape(-1, 3)

imgs = {}
for i in bpy.data.images:
    nm = i.name.lower()
    if "basecolor" in nm:
        imgs["base"] = i
    elif "normal" in nm:
        imgs["normal"] = i
    elif "_rm" in nm or nm.endswith("rm.jpg"):
        imgs["rm"] = i
log(f"images: { {k: v.name for k, v in imgs.items()} }")
base = imgs["base"]
w, h = base.size
buf = np.empty(w * h * 4, dtype=np.float32)
base.pixels.foreach_get(buf)
rgb = buf.reshape(h, w, 4)

loops_v = np.empty(len(me.loops), dtype=np.int32)
me.loops.foreach_get("vertex_index", loops_v)
uv = np.empty(len(me.loops) * 2, dtype=np.float64)
me.uv_layers.active.data.foreach_get("uv", uv)
uv = uv.reshape(-1, 2)
lx = np.clip(uv[:, 0], 0, 1) * (w - 1)
ly = np.clip(uv[:, 1], 0, 1) * (h - 1)

first_loop = np.full(n_v, len(loops_v), dtype=np.int64)
np.minimum.at(first_loop, loops_v, np.arange(len(loops_v), dtype=np.int64))
first_loop = np.minimum(first_loop, len(loops_v) - 1)
vcol = rgb[ly[first_loop].astype(int), lx[first_loop].astype(int), :3]

# ---- rosy paint: the original body has a sunburn-like blush V on the upper chest/throat.
# It reads as a tan line on the nude (uniform-skin decision) and it poisons the 3D-nearest
# fill (asymmetric pink cups). Detect it per-vertex, repaint it, and never sample from it.
band = (co[:, 2] > 0.30) & (co[:, 2] < 0.905)
rg = vcol[:, 0] - vcol[:, 1]
# reference tone is the BELLY, not the band median — the whole upper chest is rosy, so a
# band median is itself rosy-biased and lets the blush field through (measured: belly rg
# 0.239, upper chest 0.30-0.33; a med+0.05 cut only caught the extreme pink core)
belly = ~garment & (co[:, 1] < 0) & (co[:, 2] > 0.45) & (co[:, 2] < 0.60) & (np.abs(co[:, 0]) < 0.08)
med_rg = np.median(rg[belly])
rosy = band & (rg > med_rg + 0.045)
if PATCH_ONLY:
    log(f"patch mode: garment texels only, mask {garment.sum()}")
else:
    rosy_chest = rosy & (co[:, 1] < 0.01) & (co[:, 2] > 0.55)
    garment = garment | rosy_chest
    log(f"rosy: {rosy.sum()} total, chest repaint {rosy_chest.sum()} "
        f"(belly med_rg {med_rg:.3f}) -> mask {garment.sum()}")

# ---- per-vertex fill colour: K nearest skin verts in 3D ----
if PATCH_ONLY:
    skin = ~garment & (co[:, 2] > 0.30) & (co[:, 2] < 0.86)
else:
    skin = ~garment & ~rosy & (co[:, 2] > 0.30) & (co[:, 2] < 0.86)
skin_idx = np.nonzero(skin)[0][::3]          # 1-in-3 sample is plenty at this density
kd = KDTree(len(skin_idx))
for j, i in enumerate(skin_idx):
    kd.insert(Vector(co[i]), j)
kd.balance()
log(f"skin KD: {len(skin_idx)} verts")

gidx = np.nonzero(garment)[0]
fill_c = vcol.copy()


def idw_at(p):
    hits = kd.find_n(Vector(p), 8)
    wsum = 0.0
    acc = np.zeros(3)
    for (_, j, dist) in hits:
        wgt = 1.0 / max(dist * dist, 1e-9)
        acc += wgt * vcol[skin_idx[j]]
        wsum += wgt
    return acc / wsum


for i in gidx:
    c1 = idw_at(co[i])
    if PATCH_ONLY:
        # mirror-average so asymmetric blush near one cup cannot tint only that cup
        c2 = idw_at([-co[i][0], co[i][1], co[i][2]])
        fill_c[i] = 0.5 * (c1 + c2)
    else:
        fill_c[i] = c1
log("per-vertex fill colours done")

# ---- rasterise fill over garment faces ----
n_f = len(me.polygons)
l_tot = np.empty(n_f, dtype=np.int32)
me.polygons.foreach_get("loop_total", l_tot)
l_start = np.empty(n_f, dtype=np.int32)
me.polygons.foreach_get("loop_start", l_start)
gv = np.zeros(n_v, dtype=bool)
gv[gidx] = True
face_g = np.zeros(n_f, dtype=bool)
# a face is garment if ANY corner is (covers the rim); loop over faces via numpy reduceat
face_flag = np.add.reduceat(gv[loops_v].astype(np.int32), l_start)
face_g = face_flag > 0
log(f"garment faces: {face_g.sum()}")

mask_px = np.zeros((h, w), dtype=bool)
out_rgb = rgb[:, :, :3].astype(np.float64)
for fi in np.nonzero(face_g)[0]:
    s, t = l_start[fi], l_tot[fi]
    li = np.arange(s, s + t)
    P = np.stack([lx[li], ly[li]], axis=1)
    V = loops_v[li]
    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 > 256 or y1 - y0 > 256 or x1 < x0 or y1 < y0:
        continue
    if t != 3:
        P = P[:3]
        V = V[:3]
    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.03) & (b >= -0.03) & (c >= -0.03)
    if not ins.any():
        continue
    col = (a[ins, None] * fill_c[V[0]] + b[ins, None] * fill_c[V[1]]
           + c[ins, None] * fill_c[V[2]])
    out_rgb[gy[ins], gx[ins]] = col
    mask_px[gy[ins], gx[ins]] = True
log(f"rasterised fill: {mask_px.sum()} texels")

# ---- absorb mask-rim slivers (white stitch piping extends past the colour key, and UV
# island borders leave old texels between rasterised faces): dilate 8 px, propagate fill
# colours outward into the ring so no legacy pixel survives inside the dilated mask.
def dil1(m):
    g = m.copy()
    g[1:, :] |= m[:-1, :]
    g[:-1, :] |= m[1:, :]
    g[:, 1:] |= m[:, :-1]
    g[:, :-1] |= m[:, 1:]
    return g


dil = mask_px.copy()
for _ in range(8):
    dil = dil1(dil)
ring = dil & ~mask_px
C = out_rgb.copy()
have = mask_px.copy()
for _ in range(10):
    if not (ring & ~have).any():
        break
    Wf = have.astype(np.float64)
    acc = np.zeros_like(C)
    wacc = 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))
        wacc += np.roll(Wf, (dy, dx), axis=(0, 1))
    newly = ring & ~have & (wacc > 0)
    C[newly] = acc[newly] / wacc[newly, None]
    have |= newly
out_rgb[ring & have] = C[ring & have]
mask_px |= ring & have
log(f"rim absorb: +{(ring & have).sum()} ring texels -> mask {mask_px.sum()}")


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


# ---- grain transplant (per-channel high-pass, tile-based) ----
src_ok = ~mask_px
grain = np.stack([out_rgb[:, :, c] - box_blur1(out_rgb[:, :, c], 5) for c in range(3)], axis=2)
# candidate tiles must be clean AND skin-toned (the atlas also holds eyes/lips whose
# high-contrast grain would streak the fill)
skin_tone = np.median(out_rgb[mask_px], axis=0) if mask_px.any() else np.array([0.6, 0.45, 0.38])
cand = []
for ty in range(0, h - GRAIN_T, GRAIN_T):
    for tx in range(0, w - GRAIN_T, GRAIN_T):
        if not src_ok[ty:ty + GRAIN_T, tx:tx + GRAIN_T].all():
            continue
        tmean = out_rgb[ty:ty + GRAIN_T, tx:tx + GRAIN_T].reshape(-1, 3).mean(axis=0)
        if np.abs(tmean - skin_tone).max() < 0.13:
            cand.append((ty, tx))
rng = np.random.RandomState(77)
covered = 0
for ty in range(0, h - GRAIN_T + 1, GRAIN_T):
    for tx in range(0, w - GRAIN_T + 1, GRAIN_T):
        tm = mask_px[ty:ty + GRAIN_T, tx:tx + GRAIN_T]
        if not tm.any():
            continue
        sy, sx = cand[rng.randint(len(cand))]
        blk = out_rgb[ty:ty + GRAIN_T, tx:tx + GRAIN_T]
        blk[tm] += grain[sy:sy + GRAIN_T, sx:sx + GRAIN_T][tm] * 0.85
        covered += int(tm.sum())
log(f"grain: {covered} texels from {len(cand)} source tiles")

# feather rim
a_ = np.ones((h, w))
edge = mask_px.copy()
for k in range(FEATHER):
    grown = edge.copy()
    grown[1:-1, 1:-1] |= (edge[:-2, 1:-1] | edge[2:, 1:-1] | edge[1:-1, :-2] | edge[1:-1, 2:])
    ring = grown & ~edge
    a_[ring] = (k + 1) / (FEATHER + 1.0)
    edge = grown
blend = np.where(mask_px, 1.0, 1.0 - a_)[:, :, None]
orig = rgb[:, :, :3].astype(np.float64)
final = np.clip(out_rgb * blend + orig * (1 - blend), 0, 1)
buf4 = rgb.copy()
buf4[:, :, :3] = final.astype(np.float32)
base.pixels.foreach_set(buf4.reshape(-1))
base.pack()
log("basecolor updated + packed")

# ---- normal + rm: neutralise over the same texels ----
for key_, flatval in (("normal", None), ("rm", None)):
    if key_ not in imgs:
        continue
    im = imgs[key_]
    iw, ih = im.size
    b2 = np.empty(iw * ih * 4, dtype=np.float32)
    im.pixels.foreach_get(b2)
    arr = b2.reshape(ih, iw, 4)
    if (iw, ih) != (w, h):
        log(f"  {key_}: size {iw}x{ih} != base — skipping")
        continue
    if key_ == "normal":
        arr[mask_px, 0] = 0.5
        arr[mask_px, 1] = 0.5
        arr[mask_px, 2] = 1.0
    else:
        med = np.median(arr[src_ok][:, :3], axis=0)
        arr[mask_px, 0] = med[0]
        arr[mask_px, 1] = med[1]
        arr[mask_px, 2] = med[2]
    im.pixels.foreach_set(arr.reshape(-1))
    im.pack()
    log(f"  {key_}: neutralised {mask_px.sum()} texels + packed")

bpy.context.preferences.filepaths.save_version = 0  # no .blend1 autosave
bpy.ops.wm.save_as_mainfile(filepath=OUT)
log(f"WROTE {OUT}")
print("TEXTURE_DONE")