# Stage 33: bust-size variants guided by the approved concept turnaround.
#
#   blender --background --python 33_bust_variants.py -- <in.blend> <out_root> [k1,k2,...]
#
# WHY VARIANTS AND NOT ONE "MATCHED" NUMBER
# 32_ref_measure put the reference (exchange/INBOX/lena-base-turnaround, GPT Image 2, approved by
# Can 2026-08-06) beside the current sculpt. Proportions already agree — 6.4 vs 6.1 heads,
# waist/hip 0.504 vs 0.505, widths within 5% — so this is the same character, not a new one. The
# one real gap is the bust: projection +0.0156 H on the reference against +0.0210 H on the mesh.
# But that reference number is soft: in a side view a T-pose arm crosses the chest silhouette, so
# the front-most pixel at bust height may be forearm rather than breast, and the source is
# generated concept art rather than an orthographic render. Chasing it to three decimals would be
# false precision. This lane already has the right convention for exactly this decision — v01's
# bust size was picked by rendering 0.75 / 1.0 / 1.25 for Jeremy — so: render the variants.
#
# HOW THE REDUCTION WORKS
# Re-running the original sculpt at a smaller size is not an option: that restarts from 00_welded
# and throws away every heal from today (lines, flipped faces, cleavage fillet, tone). Instead the
# breast forms are rescaled in place:
#   wall  = bi-harmonic membrane across the bust region with the collar fixed — the chest wall as
#           it would be with no breasts on it, derived from the ribcage boundary;
#   P_k   = P + W * (k-1) * (P - wall)
# so k=1 is untouched, k=0.7 keeps 70% of the forms' projection off that wall, and W is a smooth
# region weight so the edit vanishes at the region edge instead of stepping. Shape character is
# preserved because the whole displacement field is scaled, not re-sculpted — the cleavage fillet
# from stage 26 scales with it and stays consistent.
#
# Nothing is saved. This emits renders framed like the reference (front + side, full body) plus the
# measured projection per k, so the choice is made from pictures against the concept art. The
# chosen k gets applied and registered in a follow-up.
import bpy, sys, os, math, time
import numpy as np
from mathutils import Vector

argv = sys.argv[sys.argv.index("--") + 1:]
BLEND, ROOT = argv[0], argv[1]
KS = [float(x) for x in argv[2].split(",")] if len(argv) > 2 else [1.0, 0.8, 0.65]
os.makedirs(ROOT, exist_ok=True)
t0 = time.time()

# bust region, in mesh units on this 0.9792-tall sculpt
Z0, Z1 = 0.620, 0.790
Z_FADE = 0.022
X_MAX = 0.135
X_FADE = 0.022
Y_FRONT = 0.010
COLLAR = 3


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


def smoothstep(x):
    x = np.clip(x, 0.0, 1.0)
    return x * x * (3.0 - 2.0 * x)


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)
orig_mats = [ms.material for ms in ob.material_slots]
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)
H = co[:, 2].max() - co[:, 2].min()
log(f"in: {n_v}v, height {H:.4f}")

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)


def grow(mask, rings):
    m = mask.copy()
    for _ in range(rings):
        hit = m[ev[:, 0]] | m[ev[:, 1]]
        m2 = m.copy()
        m2[ev[:, 0]] |= hit
        m2[ev[:, 1]] |= hit
        m = m2
    return m


W = (smoothstep((co[:, 2] - (Z0 - Z_FADE)) / Z_FADE)
     * smoothstep(((Z1 + Z_FADE) - co[:, 2]) / Z_FADE)
     * smoothstep(((X_MAX + X_FADE) - np.abs(co[:, 0])) / X_FADE)
     * smoothstep((Y_FRONT - co[:, 1]) / 0.030))
region = W > 0.02
log(f"region: {int(region.sum())} verts (W>0.5: {int((W>0.5).sum())})")


def bilaplacian(P, free_m, maxit=8000):
    collar = grow(free_m, COLLAR) & ~free_m
    S = np.nonzero(free_m | collar)[0]
    in_S = np.zeros(n_v, dtype=bool)
    in_S[S] = True
    glb = np.full(n_v, -1, dtype=np.int64)
    glb[S] = np.arange(len(S))
    se = ev[in_S[ev].all(axis=1)]
    a_, b_ = glb[se[:, 0]], glb[se[:, 1]]
    deg = np.zeros(len(S))
    np.add.at(deg, a_, 1.0)
    np.add.at(deg, b_, 1.0)
    free = free_m[S]

    def Ls(X):
        out = deg[:, None] * X
        np.add.at(out, a_, -X[b_])
        np.add.at(out, b_, -X[a_])
        return out

    def A_op(U):
        X = np.zeros((len(S), 3))
        X[free] = U
        return Ls(Ls(X))[free]

    Xc = np.zeros((len(S), 3))
    Xc[~free] = P[S[~free]]
    rhs = -Ls(Ls(Xc))[free]
    U = P[S[free]].copy()
    r = rhs - A_op(U)
    p = r.copy()
    rs = (r * r).sum()
    rs0 = max(rs, 1e-30)
    it = 0
    for it in range(maxit):
        Ap = A_op(p)
        den = (p * Ap).sum()
        if abs(den) < 1e-30:
            break
        al = rs / den
        U += al * p
        r -= al * Ap
        rs2 = (r * r).sum()
        if rs2 < 1e-20 or rs2 < rs0 * 1e-13:
            rs = rs2
            break
        p = r + (rs2 / rs) * p
        rs = rs2
    Q = P.copy()
    Q[S[free]] = U
    return Q, it, rs / rs0


# The wall solve is the expensive part (~8 min of CG over ~69k free verts), and it depends only on
# the input mesh and the region constants — not on k. Cache it so trying more sizes is seconds.
CACHE = os.path.join(os.path.dirname(os.path.abspath(ROOT)), "_bust_wall_cache.npy")
if os.path.exists(CACHE):
    wall = np.load(CACHE)
    if wall.shape == co.shape:
        log(f"wall: loaded cache {CACHE}")
    else:
        wall = None
        log("wall: cache shape mismatch, re-solving")
else:
    wall = None
if wall is None:
    log("solving the chest wall (bust removed) ...")
    wall, it, rel = bilaplacian(co, W > 0.5)
    np.save(CACHE, wall)
    log(f"wall: CG it={it} rel={rel:.1e}, cached -> {CACHE}")
d_wall = np.linalg.norm(co - wall, axis=1)
log(f"wall: forms stand up to {d_wall.max()*1815:.1f} mm off it")


def bust_projection(P):
    """Front-most torso y at the bust vs at the underbust, normalised by height — the same
    statistic 32_ref_measure applied to the reference silhouette."""
    z0 = P[:, 2].min()
    Hh = P[:, 2].max() - z0
    u = (P[:, 2] - z0) / Hh
    NB = 220
    uu = np.linspace(0, 1, NB)
    fm = np.full(NB, np.nan)
    dm = np.full(NB, np.nan)
    for i, uc in enumerate(uu):
        sel = np.abs(u - uc) < (0.5 / NB) * 1.6
        if sel.sum() < 8:
            continue
        sl = P[sel]
        t = sl[np.abs(sl[:, 0]) < 0.16]
        if len(t) >= 8:
            fm[i] = t[:, 1].min() / Hh
            dm[i] = (t[:, 1].max() - t[:, 1].min()) / Hh
    band = (uu > 0.66) & (uu < 0.80) & np.isfinite(dm)
    ib = int(np.nanargmax(np.where(band, dm, -np.inf)))
    band2 = (uu > 0.60) & (uu < 0.70) & np.isfinite(dm)
    iu = int(np.nanargmin(np.where(band2, dm, np.inf)))
    return fm[iu] - fm[ib], uu[ib]


# ---------- render rig, framed like the reference turnaround ----------
scn = bpy.context.scene
wd = bpy.data.worlds.new("W")
wd.color = (0.20, 0.20, 0.21)
scn.world = wd
for rot, e in (((55, 0, -35), 2.2), ((120, 0, 150), 1.0), ((85, 0, 35), 0.8)):
    ld = bpy.data.lights.new("s", 'SUN')
    ld.energy = e
    ld.use_shadow = False
    lo = bpy.data.objects.new("s", ld)
    lo.rotation_euler = tuple(math.radians(a) for a in rot)
    bpy.context.collection.objects.link(lo)
cam = bpy.data.objects.new("Cam", bpy.data.cameras.new("Cam"))
cam.data.lens = 70
bpy.context.collection.objects.link(cam)
scn.camera = cam
scn.render.engine = 'BLENDER_EEVEE' if bpy.app.version >= (4, 2) else 'BLENDER_EEVEE_NEXT'
scn.render.resolution_x = 700
scn.render.resolution_y = 1000
clay = bpy.data.materials.new("Clay")
clay.use_nodes = True
cn = clay.node_tree.nodes["Principled BSDF"]
cn.inputs["Base Color"].default_value = (0.80, 0.56, 0.42, 1)
cn.inputs["Roughness"].default_value = 0.55

CTR = (0.0, 0.0, 0.50)


def shoot(outdir, tag, yaw_deg, span, ctr=CTR, use_clay=True):
    os.makedirs(outdir, exist_ok=True)
    for i, ms in enumerate(ob.material_slots):
        ms.material = clay if use_clay else orig_mats[i]
    y = math.radians(yaw_deg)
    dist = span * 3.0
    cam.location = Vector(ctr) + Vector((math.sin(y) * dist, -math.cos(y) * dist, 0.0))
    cam.rotation_euler = (Vector(ctr) - cam.location).to_track_quat('-Z', 'Y').to_euler()
    scn.render.filepath = os.path.abspath(os.path.join(outdir, f"{tag}.png"))
    bpy.ops.render.render(write_still=True)


for k in KS:
    P = co + (W * (k - 1.0))[:, None] * (co - wall)
    me.vertices.foreach_set("co", P.reshape(-1))
    me.update()
    if me.has_custom_normals:
        vn = np.empty(n_v * 3, dtype=np.float32)
        me.vertices.foreach_get("normal", vn)
        me.normals_split_custom_set_from_vertices(vn.reshape(-1, 3))
    proj, ub = bust_projection(P)
    moved = np.linalg.norm(P - co, axis=1) * 1815
    tag = f"k{int(round(k*100)):03d}"
    log(f"{tag}: bust projection {proj:+.4f} H at u={ub:.3f} "
        f"(reference +0.0156) | max move {moved.max():.1f} mm")
    out = os.path.join(ROOT, tag)
    shoot(out, "front", 0, 0.55)
    shoot(out, "side", 90, 0.55)
    shoot(out, "chest", 0, 0.20, (0.0, 0.0, 0.70))
    shoot(out, "chest_side", 90, 0.20, (0.0, 0.0, 0.70))
    log(f"{tag}: rendered -> {out}")

me.vertices.foreach_set("co", co.reshape(-1))
me.update()
print("BUST_DONE")