# Stage 5b: global rosy-tone grade. Hard vertex-mask repaints leave visible boundaries
# (lighter V patch vs darker shoulders, pink neck above the z-cap). Instead: per-texel,
# compute the SMOOTHED red-green excess over the belly reference and subtract it with a
# z-tapered weight. Smooth field in, smooth field out — no boundaries; the high-frequency
# detail (pores, grain) rides on top untouched. Face is excluded by the taper (w=0 above
# z 0.91); lips/cheeks keep their red.
#   blender --background --python 05b_tone_grade.py -- <07_textured.blend> <out.blend>
import bpy, sys, time
import numpy as np

argv = sys.argv[sys.argv.index("--") + 1:]
BLEND, OUT = argv[0], argv[1]
t0 = time.time()

REF_RG = None          # belly reference; measured from the mesh below if None
STRENGTH = 0.90
Z_UP0, Z_UP1 = 0.26, 0.30      # taper in above the feet
Z_DN0, Z_DN1 = 0.88, 0.91      # taper out below the face
BLUR_R = 12


def log(m):
    print(f"[grade {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)

base = next(i for i in bpy.data.images if "basecolor" in i.name.lower())
w, h = base.size
buf = np.empty(w * h * 4, dtype=np.float32)
base.pixels.foreach_get(buf)
rgb = buf.reshape(h, w, 4)
out_rgb = rgb[:, :, :3].astype(np.float64)

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)
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)

# belly reference from the mesh (sample texels under belly verts)
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 = out_rgb[ly[first_loop].astype(int), lx[first_loop].astype(int)]
belly = (co[:, 1] < 0) & (co[:, 2] > 0.45) & (co[:, 2] < 0.60) & (np.abs(co[:, 0]) < 0.08)
ref = REF_RG if REF_RG is not None else float(np.median(vcol[belly, 0] - vcol[belly, 1]))
log(f"belly ref rg: {ref:.3f}")


def wz(z):
    a = np.clip((z - Z_UP0) / (Z_UP1 - Z_UP0), 0, 1)
    b = np.clip((Z_DN1 - z) / (Z_DN1 - Z_DN0), 0, 1)
    t = np.minimum(a, b)
    return t * t * (3 - 2 * t)


# rasterise the per-texel weight from vertex z over ALL 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)
vw = wz(co[:, 2])
W = np.zeros((h, w))
for fi in range(n_f):
    s, t = l_start[fi], l_tot[fi]
    li = np.arange(s, s + min(t, 3))
    V = loops_v[li]
    if vw[V].max() <= 0:
        continue
    P = np.stack([lx[li], ly[li]], axis=1)
    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
    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
    wv = a[ins] * vw[V[0]] + b[ins] * vw[V[1]] + c[ins] * vw[V[2]]
    W[gy[ins], gx[ins]] = np.maximum(W[gy[ins], gx[ins]], np.clip(wv, 0, 1))
log(f"weight rasterised: {(W > 0).sum()} texels")


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)


# smoothed rg field, weighted by W so face/eye texels can't bleed into the blur across
# island borders more than locally
rg_s = box_blur1(out_rgb[:, :, 0], BLUR_R) - box_blur1(out_rgb[:, :, 1], BLUR_R)
delta = np.maximum(0.0, rg_s - ref) * W * STRENGTH
out_rgb[:, :, 0] -= delta * 0.8
out_rgb[:, :, 1] += delta * 0.2
log(f"graded: {(delta > 0.005).sum()} texels above 0.005, max delta {delta.max():.3f}")

buf4 = rgb.copy()
buf4[:, :, :3] = np.clip(out_rgb, 0, 1).astype(np.float32)
base.pixels.foreach_set(buf4.reshape(-1))
base.pack()
bpy.context.preferences.filepaths.save_version = 0  # no .blend1 autosave
bpy.ops.wm.save_as_mainfile(filepath=OUT)
log(f"WROTE {OUT}")
print("GRADE_DONE")