# Stage 36: re-author the breast + underwear region of the NEW atlas.
#
#   blender --background --python 36_refill.py -- <reatlased.blend> <welded_orig.blend>
#                                                 <masks.npz> <out_dir> <out.glb>
#
# Jeremy's brief: keep her original texture everywhere, replace only the breast and underwear
# area. The approved turnaround (`exchange/INBOX/lena-base-turnaround/`) specifies what the
# replacement should look like: featureless — soft volume, no nipples, no crease or garment
# detail, smooth crotch.
#
# WHY THIS IS NOT 05_texture.py AGAIN. That pass filled the region with an inverse-distance
# average of nearby skin and sprinkled grain on top. An IDW average does not agree with the skin
# it abuts, so the fill landed as a flat decal with a tone step at the old bra/briefs outline —
# the ghost Jeremy has been seeing. Here the low-frequency tone is a HARMONIC solve on the mesh
# with Dirichlet boundaries: the fill is *defined* to equal her real skin at the mask edge, so a
# step is impossible rather than merely small. Grain is then transplanted back so it is not
# plastic.
#
# WHY ON THE MESH AND NOT IN THE ATLAS. Neighbouring texels in an atlas are not neighbours on the
# body. Diffusing in atlas space leaks across chart edges and stops at them; the crotch alone
# straddles the torso/leg charts. Solving per-vertex makes the fill continuous across every seam
# by construction. The mesh carries the low frequencies (which is all a featureless fill has);
# grain goes on afterwards in texture space, where resolution actually matters.
#
# The solve is conjugate-gradient, not Jacobi. This lane has already been bitten once by harmonic
# diffusion needing ~width^2 relaxation passes and leaving the interior at its seed tone (the
# "pale panty ghost"); CG converges in ~width iterations and reports its residual.
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, WELDED, MASKS, OUTDIR, OUTGLB = argv[0], argv[1], argv[2], os.path.abspath(argv[3]), \
    os.path.abspath(argv[4])
os.makedirs(OUTDIR, exist_ok=True)
GROW = 3          # edge-dilation rings on the decimated mesh, to clear the old garment rim
GRAIN_T = 16
t0 = time.time()


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


# ---------------------------------------------------------------- the mask, from the hires lane
M = np.load(MASKS)
orig_mask = (M["garment"] | M["cups"] | M["briefs"] | M["hemband"] | M["key_raw"])
log(f"masks.npz: {orig_mask.sum()} of {len(orig_mask)} original verts are garment")

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))
n0 = len(ob0.data.vertices)
assert n0 == len(orig_mask), f"masks.npz is for {len(orig_mask)} verts, welded has {n0}"
co0 = np.empty(n0 * 3)
ob0.data.vertices.foreach_get("co", co0)
co0 = co0.reshape(-1, 3)
log(f"welded original: {n0} verts")

# ---------------------------------------------------------------- the body we are painting
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)
co = np.empty(n_v * 3); me.vertices.foreach_get("co", co); co = co.reshape(-1, 3)
log(f"target body: {n_v} verts, {n_f} faces")

# transfer the mask by nearest original vertex (the sculpt has moved a few mm since masks.npz
# was built, which the GROW rings below absorb)
kd = KDTree(n0)
for i in range(n0):
    kd.insert(Vector(co0[i]), i)
kd.balance()
mask = np.zeros(n_v, dtype=bool)
dists = np.empty(n_v)
for i in range(n_v):
    _, j, d = kd.find(Vector(co[i]))
    mask[i] = orig_mask[j]
    dists[i] = d
UNITM = 1.777 / (co[:, 2].max() - co[:, 2].min())
log(f"nearest-original distance: p50 {np.percentile(dists,50)*UNITM*1000:.2f} mm, "
    f"p99 {np.percentile(dists,99)*UNITM*1000:.2f} mm")
log(f"transferred mask: {int(mask.sum())} verts ({100.0*mask.mean():.1f}%)")

# the crotch was replaced by a donor surface after masks.npz was built, so it is not in those
# masks; add it explicitly (same box 05_texture.py used, in body-frame units)
z = co[:, 2]
crotch = (np.abs(co[:, 0]) < 0.075) & (z > 0.340) & (z < 0.480)
log(f"crotch box adds {int((crotch & ~mask).sum())} verts")
mask |= crotch

ev = np.empty(len(me.edges) * 2, dtype=np.int32); me.edges.foreach_get("vertices", ev)
ev = ev.reshape(-1, 2)
for _ in range(GROW):
    hit = mask[ev[:, 0]] | mask[ev[:, 1]]
    mask[ev[hit, 0]] = True
    mask[ev[hit, 1]] = True
log(f"after {GROW} grow rings: {int(mask.sum())} verts ({100.0*mask.mean():.1f}%)")

# ---------------------------------------------------------------- current per-vertex colour
src = {}
for slot in ob.material_slots:
    mat = slot.material
    if not mat or not mat.node_tree:
        continue
    for node in mat.node_tree.nodes:
        if node.type != 'BSDF_PRINCIPLED':
            continue
        for sock, key in (("Base Color", "base"), ("Normal", "normal"), ("Roughness", "rm")):
            if sock not in node.inputs or not node.inputs[sock].links:
                continue
            nd = node.inputs[sock].links[0].from_node
            seen = set()
            while nd and nd.type != 'TEX_IMAGE' and id(nd) not in seen:
                seen.add(id(nd))
                nxt = None
                for i in nd.inputs:
                    if i.links:
                        nxt = i.links[0].from_node
                        break
                nd = nxt
            if nd and nd.type == 'TEX_IMAGE' and nd.image:
                src[key] = nd.image
base = src["base"]
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}")

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)
l_start = np.empty(n_f, dtype=np.int32); me.polygons.foreach_get("loop_start", l_start)
l_tot = np.empty(n_f, dtype=np.int32); me.polygons.foreach_get("loop_total", l_tot)
li = l_start[l_tot == 3]

px = np.clip(np.round(uv[:, 0] * (W - 1)).astype(np.int32), 0, W - 1)
py = np.clip(np.round(uv[:, 1] * (H - 1)).astype(np.int32), 0, H - 1)
lc = rgb[py, px]
vc = np.zeros((n_v, 3))
for c in range(3):
    vc[:, c] = np.bincount(loops_v, weights=lc[:, c], minlength=n_v)
vn = np.maximum(np.bincount(loops_v, minlength=n_v), 1)
vc /= vn[:, None]

# ---------------------------------------------------------------- harmonic fill (CG)
o_ = np.concatenate([ev[:, 0], ev[:, 1]])
n_ = np.concatenate([ev[:, 1], ev[:, 0]])
deg = np.bincount(o_, minlength=n_v).astype(np.float64)
deg = np.maximum(deg, 1.0)
mf = mask.astype(np.float64)


def A_mul(x):
    """graph Laplacian restricted to the unknown (masked) set"""
    xm = x * mf
    return (deg * xm - np.bincount(o_, weights=xm[n_], minlength=n_v)) * mf


fill = vc.copy()
for c in range(3):
    known = vc[:, c] * (1.0 - mf)
    b = np.bincount(o_, weights=known[n_], minlength=n_v) * mf
    x = np.zeros(n_v)
    r = b - A_mul(x)
    p = r.copy()
    rs = float(r @ r)
    r0 = rs
    it = 0
    for it in range(4000):
        if rs <= max(r0 * 1e-12, 1e-20):
            break
        Ap = A_mul(p)
        denom = float(p @ Ap)
        if abs(denom) < 1e-30:
            break
        al = rs / denom
        x += al * p
        r -= al * Ap
        rs2 = float(r @ r)
        p = r + (rs2 / rs) * p
        rs = rs2
    fill[mask, c] = x[mask]
    log(f"  channel {c}: CG {it+1} iters, residual {np.sqrt(rs/max(r0,1e-30)):.2e}")

# sanity: the fill must agree with real skin at the boundary
bnd = mask & (np.bincount(o_, weights=(1.0 - mf)[n_], minlength=n_v) > 0)
if bnd.any():
    step = np.abs(fill[bnd] - vc[bnd]).max(axis=1)
    log(f"boundary vertices {int(bnd.sum())}: |fill - original| mean {step.mean():.4f} "
        f"p99 {np.percentile(step,99):.4f}  (a step here would be the old ghost)")
log(f"fill tone: mean {fill[mask].mean(axis=0)} vs surrounding skin "
    f"{vc[~mask & (z>0.33) & (z<0.86)].mean(axis=0)}")

# ---------------------------------------------------------------- rasterise into the atlas
IDX = np.stack([li, li + 1, li + 2], axis=1)
V = loops_v[IDX]
face_any = mask[V].any(axis=1)
P = np.stack([uv[IDX][:, :, 0] * (W - 1), uv[IDX][:, :, 1] * (H - 1)], axis=2)
out = rgb.copy()
paint = np.zeros((H, W), dtype=bool)
for f in np.nonzero(face_any)[0]:
    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
    aa, bb, cc = a[ins], b_[ins], c_[ins]
    w = aa * mf[V[f, 0]] + bb * mf[V[f, 1]] + cc * mf[V[f, 2]]
    col = (aa[:, None] * fill[V[f, 0]] + bb[:, None] * fill[V[f, 1]] + cc[:, None] * fill[V[f, 2]])
    yy, xx = gy[ins], gx[ins]
    hard = w > 0.5
    if hard.any():
        out[yy[hard], xx[hard]] = col[hard]
        paint[yy[hard], xx[hard]] = True
log(f"repainted {int(paint.sum())} texels ({100.0*paint.mean():.2f}% of the atlas)")

# ---------------------------------------------------------------- grain, so it is not plastic
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)


grain = np.stack([rgb[:, :, c] - box(rgb[:, :, c], 5) for c in range(3)], axis=2)
skin_tone = np.median(out[paint], axis=0) if paint.any() else np.array([0.7, 0.45, 0.3])
cand = []
for ty in range(0, H - GRAIN_T, GRAIN_T):
    for tx in range(0, W - GRAIN_T, GRAIN_T):
        if paint[ty:ty + GRAIN_T, tx:tx + GRAIN_T].any():
            continue
        t = rgb[ty:ty + GRAIN_T, tx:tx + GRAIN_T].reshape(-1, 3)
        if t.min() < 0.02:                       # skip tiles touching empty atlas
            continue
        if np.abs(t.mean(axis=0) - skin_tone).max() < 0.10:
            cand.append((ty, tx))
rng = np.random.RandomState(11)
amp = 0.85
if cand:
    for ty in range(0, H - GRAIN_T + 1, GRAIN_T):
        for tx in range(0, W - GRAIN_T + 1, GRAIN_T):
            tm = paint[ty:ty + GRAIN_T, tx:tx + GRAIN_T]
            if not tm.any():
                continue
            sy, sx = cand[rng.randint(len(cand))]
            out[ty:ty + GRAIN_T, tx:tx + GRAIN_T][tm] += \
                grain[sy:sy + GRAIN_T, sx:sx + GRAIN_T][tm] * amp
log(f"grain: {len(cand)} source tiles, amplitude {amp}")
out = np.clip(out, 0, 1)

# ---------------------------------------------------------------- write maps
def dil(m, k):
    g = m.copy()
    for _ in range(k):
        n2 = g.copy()
        n2[1:, :] |= g[:-1, :]; n2[:-1, :] |= g[1:, :]
        n2[:, 1:] |= g[:, :-1]; n2[:, :-1] |= g[:, 1:]
        g = n2
    return g


newimg = {}
b4 = tex.copy()
b4[:, :, :3] = out.astype(np.float32)
base.pixels.foreach_set(b4.reshape(-1))
base.pack()

# normal flat and rm to surrounding-skin median over the same texels: with the paint gone the
# fabric weave must stop shading through, and flat normals are what guarantee no nipple or
# crotch detail can reappear
soft = dil(paint, 2)
for key in ("normal", "rm"):
    if key not in src:
        continue
    im = src[key]
    iw, ih = im.size
    if (iw, ih) != (W, H):
        log(f"  {key}: {iw}x{ih} != atlas, skipped")
        continue
    a2 = np.empty(iw * ih * 4, dtype=np.float32)
    im.pixels.foreach_get(a2)
    arr = a2.reshape(ih, iw, 4)
    if key == "normal":
        arr[soft, 0] = 0.5; arr[soft, 1] = 0.5; arr[soft, 2] = 1.0
    else:
        med = np.median(arr[~dil(paint, 8)][:, :3], axis=0)
        arr[soft, 0] = med[0]; arr[soft, 1] = med[1]; arr[soft, 2] = med[2]
        log(f"  rm median from surrounding skin: {med}")
    im.pixels.foreach_set(arr.reshape(-1))
    im.pack()
    log(f"  {key}: {int(soft.sum())} texels neutralised")

for k, im in list(src.items()):
    p = os.path.join(OUTDIR, f"lena_nude_refill_{k}.jpg")
    im.file_format = 'JPEG'
    im.filepath_raw = p
    im.save(filepath=p)
np.save(os.path.join(OUTDIR, "refill_mask.npy"), paint)
bpy.ops.wm.save_as_mainfile(filepath=os.path.join(OUTDIR, "refilled.blend"))
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=OUTGLB, export_format='GLB', use_selection=True,
                          export_image_format='AUTO', export_jpeg_quality=95,
                          export_yup=True, export_apply=False)
log(f"EXPORTED {OUTGLB} ({os.path.getsize(OUTGLB)/1e6:.1f} MB)")
print("REFILL_DONE")