# Stage 6 v2: male-donor crotch transfer, full-gusset window.
#
# v1 post-mortem: the female briefs gusset bridges the legs ~3 cm ABOVE the male crotch
# saddle, so with SNAP_MAX 0.020 the lower half of the gusset couldn't reach the donor
# cage — the surface tore along the snapped/rejected boundary and the leftover panel kept
# its gathered-cloth wrinkles. v2 widens the window to the whole gusset, raises SNAP_MAX
# to span the real gap, melts whatever the cage still rejects, and Taubin-polishes the
# window so the cloth gathers on the inner thighs go too.
#
#   blender --background --python 06_crotch_v2.py -- <03_final_geometry.blend> <male.glb> <out.blend>
import bpy, sys, time
from collections import deque
import numpy as np
from mathutils import Vector
from mathutils.bvhtree import BVHTree

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

WIN_Z0, WIN_Z1 = 0.355, 0.470   # absolute: whole briefs gusset + pubic base
WIN_R_X = 0.058                 # into the inner-thigh gathers, not past the thigh walls
SNAP_MAX = 0.035                # the gusset sits ~3 cm above the male crotch
BLEND_RINGS = 10
MELT_ITERS = 60                 # residue verts the cage rejected
TAUBIN_PAIRS = 10


def log(m):
    print(f"[crotch2 {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)
f_height = co[:, 2].max() - co[:, 2].min()

mid_f = (np.abs(co[:, 0]) < 0.01) & (co[:, 2] > 0.38) & (co[:, 2] < 0.50)
saddle_f = co[mid_f][np.argmin(co[mid_f, 2])]
log(f"female: height {f_height:.4f}, saddle {saddle_f}")

# ---- male donor ----
before = set(bpy.data.objects)
bpy.ops.import_scene.gltf(filepath=MALE)
new = [o for o in bpy.data.objects if o not in before]
male = max([o for o in new if o.type == 'MESH'], key=lambda o: len(o.data.vertices))
mme = male.data
nm = len(mme.vertices)
mco = np.empty(nm * 3)
mme.vertices.foreach_get("co", mco)
mco = mco.reshape(-1, 3)
m_height = mco[:, 2].max() - mco[:, 2].min()
mid_m = (np.abs(mco[:, 0]) < 0.02) & (mco[:, 2] > 0.55 * m_height) & (mco[:, 2] < 0.75 * m_height)
if not mid_m.any():
    mid_m = (np.abs(mco[:, 0]) < 0.02)
saddle_m = mco[mid_m][np.argmin(mco[mid_m, 2])]
s = f_height / m_height
log(f"male '{male.name}': {nm}v height {m_height:.3f} m, saddle {saddle_m}, scale {s:.4f}")

# male crotch patch generous enough to cover the whole female window after transform
half_z = max(abs(WIN_Z1 - saddle_f[2]), abs(saddle_f[2] - WIN_Z0))
sel_m = (np.abs(mco[:, 0] - saddle_m[0]) < WIN_R_X / s * 1.8) \
        & (np.abs(mco[:, 2] - saddle_m[2]) < half_z / s * 1.8)
log(f"male patch verts: {sel_m.sum()}")

mco_t = (mco - saddle_m) * s + saddle_f

ml_tot = np.empty(len(mme.polygons), dtype=np.int32)
mme.polygons.foreach_get("loop_total", ml_tot)
ml_start = np.empty(len(mme.polygons), dtype=np.int32)
mme.polygons.foreach_get("loop_start", ml_start)
ml_v = np.empty(len(mme.loops), dtype=np.int32)
mme.loops.foreach_get("vertex_index", ml_v)
faces = []
for fs, ft in zip(ml_start, ml_tot):
    idxs = ml_v[fs:fs + ft]
    if sel_m[idxs].all():
        faces.append(idxs.tolist())
log(f"male patch faces: {len(faces)}")
used = sorted(set(i for f in faces for i in f))
remap = {g: i for i, g in enumerate(used)}
pm = bpy.data.meshes.new("crotch_patch")
pm.from_pydata([Vector(mco_t[i]) for i in used], [],
               [[remap[i] for i in f] for f in faces])
pm.update()
po = bpy.data.objects.new("crotch_patch", pm)
bpy.context.collection.objects.link(po)
sub = po.modifiers.new("s", 'SUBSURF')
sub.levels = 2
bpy.context.view_layer.objects.active = po
bpy.ops.object.modifier_apply(modifier="s")
dme = po.data
dv = np.empty(len(dme.vertices) * 3)
dme.vertices.foreach_get("co", dv)
dv = dv.reshape(-1, 3)
dl_tot = np.empty(len(dme.polygons), dtype=np.int32)
dme.polygons.foreach_get("loop_total", dl_tot)
dl_start = np.empty(len(dme.polygons), dtype=np.int32)
dme.polygons.foreach_get("loop_start", dl_start)
dl_v = np.empty(len(dme.loops), dtype=np.int32)
dme.loops.foreach_get("vertex_index", dl_v)
polys_d = [dl_v[fs:fs + ft].tolist() for fs, ft in zip(dl_start, dl_tot)]
bvh = BVHTree.FromPolygons([Vector(v) for v in dv], polys_d,
                           all_triangles=False, epsilon=0.0)
dv_f = dv.copy()
dv_f[:, 1] = 2 * saddle_f[1] - dv[:, 1]
bvh_f = BVHTree.FromPolygons([Vector(v) for v in dv_f], polys_d,
                             all_triangles=False, epsilon=0.0)
for o in new + [po]:
    bpy.data.objects.remove(o, do_unlink=True)
log("donor cages ready (both orientations)")

# ---- female window ----
win = (np.abs(co[:, 0] - saddle_f[0]) < WIN_R_X) \
      & (co[:, 2] > WIN_Z0) & (co[:, 2] < WIN_Z1)
widx = np.nonzero(win)[0]
log(f"female window: {len(widx)} verts")

n_e = len(me.edges)
ev = np.empty(n_e * 2, dtype=np.int32)
me.edges.foreach_get("vertices", ev)
ev = ev.reshape(-1, 2)
order = np.concatenate([ev[:, 0], ev[:, 1]])
nbr = np.concatenate([ev[:, 1], ev[:, 0]])
srt = np.argsort(order, kind="stable")
o_s = order[srt]
n_s = nbr[srt]
ptr = np.searchsorted(o_s, np.arange(n_v + 1))
depth = np.zeros(n_v, dtype=np.int32)
dq = deque()
seen = np.zeros(n_v, dtype=bool)
for a, b in ev:
    if win[a] != win[b]:
        sv = a if win[a] else b
        if not seen[sv]:
            seen[sv] = True
            depth[sv] = 1
            dq.append(sv)
while dq:
    c = dq.popleft()
    for nb in n_s[ptr[c]:ptr[c + 1]]:
        if win[nb] and not seen[nb]:
            seen[nb] = True
            depth[nb] = depth[c] + 1
            dq.append(nb)
depth[win & ~seen] = BLEND_RINGS + 2
wgt = np.clip(depth[widx] / float(BLEND_RINGS), 0.0, 1.0)
wgt = wgt * wgt * (3 - 2 * wgt)

samp = widx[::37]
def med_d(tree):
    ds = []
    for i in samp:
        h = tree.find_nearest(Vector(co[i]), 0.10)
        ds.append(h[3] if h[0] is not None else 0.10)
    return float(np.median(ds))
d_id, d_fl = med_d(bvh), med_d(bvh_f)
use = bvh if d_id <= d_fl else bvh_f
log(f"orientation: identity {d_id:.4f} vs flipped {d_fl:.4f} -> "
    f"{'identity' if d_id <= d_fl else 'flipped'}")

co_new = co.copy()
moved = 0
rejected = []
for k, i in enumerate(widx):
    hit = use.find_nearest(Vector(co[i]), SNAP_MAX)
    if hit[0] is None:
        rejected.append(i)
        continue
    tgt = np.array(hit[0])
    co_new[i] += (tgt - co[i]) * wgt[k]
    moved += 1
delta = np.linalg.norm(co_new - co, axis=1)
log(f"snapped {moved} verts (rejected {len(rejected)}), max move {delta.max():.4f}")

# ---- melt the rejects toward their neighbours so no torn seam survives ----
def neigh_mean(Q, idx):
    out = np.empty((len(idx), 3))
    for j, i in enumerate(idx):
        nbrs = n_s[ptr[i]:ptr[i + 1]]
        out[j] = Q[nbrs].mean(axis=0) if len(nbrs) else Q[i]
    return out

rej = np.array(rejected, dtype=np.int32)
if len(rej):
    for _ in range(MELT_ITERS):
        co_new[rej] = 0.5 * co_new[rej] + 0.5 * neigh_mean(co_new, rej)
    log(f"melted {len(rej)} rejected verts")

# ---- Taubin polish over the whole window (kills cloth gathers on the inner thighs) ----
lam, mu = 0.5, -0.53
for _ in range(TAUBIN_PAIRS):
    co_new[widx] += lam * (neigh_mean(co_new, widx) - co_new[widx])
    co_new[widx] += mu * (neigh_mean(co_new, widx) - co_new[widx])
log(f"Taubin x{TAUBIN_PAIRS} on window")

me.vertices.foreach_set("co", co_new.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))
bpy.context.preferences.filepaths.save_version = 0  # no .blend1 autosave
bpy.ops.wm.save_as_mainfile(filepath=OUT)
log(f"WROTE {OUT}")
print("CROTCH2_DONE")