# Tailor measurements from a Quaternius-skinned body GLB.
# Slices the mesh at scanned heights, restricted to vertices whose dominant
# vertex group is in an allowed bone set (so arms don't inflate the chest).
# Run: blender --background --python measure_body.py -- <body.glb> <out.json>
import bpy, sys, json, math

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

bpy.ops.wm.read_factory_settings(use_empty=True)
bpy.ops.import_scene.gltf(filepath=SRC)

mesh_obj = max((o for o in bpy.data.objects if o.type == "MESH"),
               key=lambda o: len(o.data.vertices))
arm_obj = next(o for o in bpy.data.objects if o.type == "ARMATURE")

# --- world-space bone head positions (rest pose) ---
bones = {}
for b in arm_obj.data.bones:
    bones[b.name] = arm_obj.matrix_world @ b.head_local

# --- world-space vertex positions + dominant vertex group ---
mw = mesh_obj.matrix_world
group_name = {g.index: g.name for g in mesh_obj.vertex_groups}
verts = []   # (co_world, dominant_group_name)
for v in mesh_obj.data.vertices:
    best, bw = None, 0.0
    for g in v.groups:
        if g.weight > bw:
            bw, best = g.weight, group_name.get(g.group)
    verts.append((mw @ v.co, best))

edges = [(e.vertices[0], e.vertices[1]) for e in mesh_obj.data.edges]

def hull_perimeter(pts):
    """2D convex hull (monotone chain) perimeter."""
    pts = sorted(set(pts))
    if len(pts) < 3:
        return 0.0
    def cross(o, a, b):
        return (a[0]-o[0])*(b[1]-o[1]) - (a[1]-o[1])*(b[0]-o[0])
    lower, upper = [], []
    for p in pts:
        while len(lower) >= 2 and cross(lower[-2], lower[-1], p) <= 0:
            lower.pop()
        lower.append(p)
    for p in reversed(pts):
        while len(upper) >= 2 and cross(upper[-2], upper[-1], p) <= 0:
            upper.pop()
        upper.append(p)
    hull = lower[:-1] + upper[:-1]
    return sum(math.dist(hull[i], hull[(i+1) % len(hull)]) for i in range(len(hull)))

def slice_circ(z, allowed):
    """Perimeter of the mesh cross-section at height z, using edges whose
    both endpoints' dominant group is in `allowed` (prefix match)."""
    def ok(g):
        return g is not None and any(g.startswith(a) for a in allowed)
    pts = []
    for i0, i1 in edges:
        c0, g0 = verts[i0]
        c1, g1 = verts[i1]
        if not (ok(g0) and ok(g1)):
            continue
        if (c0.z - z) * (c1.z - z) < 0:
            t = (z - c0.z) / (c1.z - c0.z)
            pts.append((c0.x + t*(c1.x-c0.x), c0.y + t*(c1.y-c0.y)))
    return hull_perimeter(pts)

def scan(z0, z1, allowed, mode, steps=40):
    """Max or min nonzero circumference in [z0,z1]."""
    best_c, best_z = (0.0, None) if mode == "max" else (1e9, None)
    for i in range(steps + 1):
        z = z0 + (z1 - z0) * i / steps
        c = slice_circ(z, allowed)
        if c <= 0:
            continue
        if (mode == "max" and c > best_c) or (mode == "min" and c < best_c):
            best_c, best_z = c, z
    return best_c, best_z

TORSO = ["pelvis", "spine_", "neck_", "clavicle_", "Breast"]
HIPZONE = ["pelvis", "spine_01", "thigh_"]

z_pelvis = bones["pelvis"].z
z_neck = bones["neck_01"].z
z_thigh = bones["thigh_l"].z
z_calf = bones["calf_l"].z
z_foot = bones["foot_l"].z

zmin = min(c.z for c, _ in verts)
zmax = max(c.z for c, _ in verts)

chest, z_chest = scan(z_pelvis + 0.55*(z_neck - z_pelvis), z_neck - 0.05*(z_neck - z_pelvis), TORSO, "max")
waist, z_waist = scan(z_pelvis + 0.05, z_chest - 0.03, TORSO, "min")
hip, z_hip = scan(z_thigh, z_pelvis + 0.06, HIPZONE, "max")
thigh, z_th = scan(z_thigh - 0.6*(z_thigh - z_calf), z_thigh - 0.05*(z_thigh - z_calf), ["thigh_l"], "max")
neckc, z_nk = scan(z_neck, z_neck + 0.5*(zmax - z_neck), ["neck_", "Head", "spine_03", "clavicle_"], "min")
# upper arm circumference (sleeve): slice perpendicular-ish — arms in T-pose are
# horizontal, so slice by X instead of Z for the left arm.
def slice_circ_x(x, allowed):
    def ok(g):
        return g is not None and any(g.startswith(a) for a in allowed)
    pts = []
    for i0, i1 in edges:
        c0, g0 = verts[i0]
        c1, g1 = verts[i1]
        if not (ok(g0) and ok(g1)):
            continue
        if (c0.x - x) * (c1.x - x) < 0:
            t = (x - c0.x) / (c1.x - c0.x)
            pts.append((c0.y + t*(c1.y-c0.y), c0.z + t*(c1.z-c0.z)))
    return hull_perimeter(pts)

x_ua = bones["upperarm_l"].x
x_la = bones["lowerarm_l"].x
bicep = 0.0
xb = None
for i in range(21):
    x = x_ua + (x_la - x_ua) * (0.15 + 0.5 * i / 20)
    c = slice_circ_x(x, ["upperarm_l"])
    if c > bicep:
        bicep, xb = c, x

d = math.dist
m = {
    "source": SRC,
    "units": "meters (glTF)",
    "height_total": zmax - zmin,
    "chest_circ": chest, "chest_z": z_chest,
    "waist_circ": waist, "waist_z": z_waist,
    "hip_circ": hip, "hip_z": z_hip,
    "thigh_circ": thigh, "thigh_z": z_th,
    "neck_circ": neckc, "neck_z": z_nk,
    "bicep_circ": bicep,
    "shoulder_width": d(bones["upperarm_l"], bones["upperarm_r"]),
    "arm_len_shoulder_to_wrist": d(bones["upperarm_l"], bones["lowerarm_l"]) + d(bones["lowerarm_l"], bones["hand_l"]),
    "nape_to_pelvis": z_neck - z_pelvis,
    "crotch_height": z_thigh,
    "pelvis_height": z_pelvis,
    "knee_height": z_calf,
    "ankle_height": z_foot,
    "shoulder_z": bones["clavicle_l"].z,
}
with open(OUT, "w") as f:
    json.dump({k: (round(v, 4) if isinstance(v, float) else v) for k, v in m.items()}, f, indent=2)
print("WROTE", OUT)