"""
Spike T-1089 / Synty Sidekick intake — step 4: 11-body batch fit.

Takes the transplanted torso garment (fitted to average_m, bound to our shared
armature = the average_m bind skeleton) through a multi-body refit, following
the fit_outfits_to_bodies.py Surface-Deform approach, adapted for the measured
reality that the 11 bodies do NOT share topology (03_recon_bodies.py:
every body is independently authored; thin_*/heavy_* even lack seg_hips).

Per target body:
  1. Import the transplanted garment GLB; capture reference (average_m)
     landmarks from its armature; detach + freeze the garment mesh.
  2. Import the average_m reference segments, join -> ref_body.
  3. Import the target body segments (skipping missing ones), keep the
     seg_torso armature as the target bind skeleton, join -> tgt_body.
  4. Landmark affine (same machinery as 02): scale XY by shoulder ratio,
     scale Z by pelvis->upperarm span ratio, translate pelvis->pelvis.
     Applied to BOTH the garment and ref_body.
  5. BVH nearest-surface warp (replaces Surface Deform + Shrinkwrap —
     first attempt showed SD driven by a shrinkwrapped driver mesh
     spikes 14-28x on fold discontinuities, and SD bind itself is
     unreliable against the multi-shell joined segment mesh):
       for each garment vertex, find the nearest point on ref_body,
       decompose the offset into (height along surface normal +
       tangential residual), re-evaluate the same surface point on
       tgt_body via its own nearest-surface lookup, and rebuild the
       vertex at the target surface with the offset preserved.  The
       per-vertex displacement field is then Laplacian-smoothed over
       the garment mesh connectivity to kill nearest-point-map
       discontinuities while keeping rigid details (back device,
       plating) coherent.
  6. Wipe vgroups, Data Transfer POLYINTERP_NEAREST weights from tgt_body,
     normalize, bind to the target body's embedded armature (same bind pose
     as the body segments -> garment deforms identically to the body under
     any runtime skeleton pose).
  7. Export armature + garment GLB (export_skins=True) to
     out/bodies/SK_SCFI_CIVL_09_10TORS_HU01_<body>.glb
  8. Metrics (out/bodies_fit_log.json): zero-weight verts, AABB,
     edge-length distortion vs post-affine baseline (collapse/fold
     detector), signed-distance stats garment vs body surface
     (penetration detector).

Run:
    tooling/blender --background --python \
        spikes/synty-intake/scripts/04_batch_fit_bodies.py
"""

import bpy
import json
import math
import os
from mathutils import Matrix, Vector
from mathutils.bvhtree import BVHTree

REPO = "/var/mnt/data/projects/settled-reach"
GARMENT_GLB = os.path.join(REPO, "spikes/synty-intake/out",
                           "SK_SCFI_CIVL_09_10TORS_HU01_quaternius.glb")
BODIES_DIR = os.path.join(REPO, "client/assets/characters/bodies")
OUT_DIR = os.path.join(REPO, "spikes/synty-intake/out/bodies")
OUT_LOG = os.path.join(REPO, "spikes/synty-intake/out/bodies_fit_log.json")

BODY_TYPES = ["average_m", "average_f", "muscular_m", "muscular_f",
              "thin_m", "thin_f", "heavy_m", "heavy_f",
              "teen_m", "teen_f", "child"]

SEGS = ["seg_torso", "seg_hips", "seg_neck",
        "seg_arm_upper_l", "seg_arm_upper_r",
        "seg_leg_upper_l", "seg_leg_upper_r"]

PENETRATION_MM = 3.0   # garment vertex deeper than this inside the body -> flag


def deselect_all():
    bpy.ops.object.select_all(action='DESELECT')


def set_active(obj):
    bpy.context.view_layer.objects.active = obj


def clear_scene():
    bpy.ops.object.select_all(action='SELECT')
    bpy.ops.object.delete(use_global=False)
    for block_list in (bpy.data.meshes, bpy.data.armatures, bpy.data.images,
                       bpy.data.materials):
        for block in list(block_list):
            if block.users == 0:
                block_list.remove(block)


def import_new(op, path, **kw):
    before = set(bpy.data.objects)
    op(filepath=path, **kw)
    return list(set(bpy.data.objects) - before)


def find_armature(objs):
    for o in objs:
        if o.type == 'ARMATURE':
            return o
    return None


def bone_head_world(arm, name):
    b = arm.data.bones.get(name)
    return (arm.matrix_world @ b.head_local).copy() if b else None


def delete_obj(obj):
    deselect_all()
    obj.select_set(True)
    set_active(obj)
    bpy.ops.object.delete(use_global=False)


def freeze(obj):
    """Unparent keeping transform, strip modifiers, bake transform to data."""
    deselect_all()
    set_active(obj)
    obj.select_set(True)
    mw = obj.matrix_world.copy()
    obj.parent = None
    obj.matrix_world = mw
    for mod in list(obj.modifiers):
        obj.modifiers.remove(mod)
    bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)


def import_body(body, keep_armature):
    """Import + join a body's garment-adjacent segments.

    Returns (joined_mesh, armature_or_None, segments_used).
    """
    seg_meshes = []
    kept_arm = None
    used = []
    for seg in SEGS:
        path = os.path.join(BODIES_DIR, body, seg + ".glb")
        if not os.path.exists(path):
            continue
        objs = import_new(bpy.ops.import_scene.gltf, path)
        picked = None
        for o in objs:
            if o.type == 'MESH' and len(o.vertex_groups) > 0:
                picked = o
            elif o.type == 'MESH':
                delete_obj(o)  # debris (bounds helpers etc.)
        arm = find_armature(objs)
        if picked is not None:
            pmw = picked.matrix_world.copy()
            picked.parent = None
            picked.matrix_world = pmw
        if arm is not None:
            if keep_armature and kept_arm is None:
                kept_arm = arm
            else:
                delete_obj(arm)
        if picked is None:
            print(f"    WARNING: no weighted mesh in {body}/{seg}")
            continue
        for mod in list(picked.modifiers):
            picked.modifiers.remove(mod)
        seg_meshes.append(picked)
        used.append(seg)
    deselect_all()
    for m in seg_meshes:
        m.select_set(True)
    set_active(seg_meshes[0])
    if len(seg_meshes) > 1:
        bpy.ops.object.join()
    joined = bpy.context.view_layer.objects.active
    bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
    return joined, kept_arm, used


def edge_lengths(mesh_obj):
    verts = mesh_obj.data.vertices
    return [(verts[e.vertices[0]].co - verts[e.vertices[1]].co).length
            for e in mesh_obj.data.edges]


def aabb(mesh_obj):
    xs = [v.co for v in mesh_obj.data.vertices]
    lo = Vector((min(v.x for v in xs), min(v.y for v in xs),
                 min(v.z for v in xs)))
    hi = Vector((max(v.x for v in xs), max(v.y for v in xs),
                 max(v.z for v in xs)))
    return lo, hi


def warp_garment(garment, ref_body, tgt_body, depsgraph,
                 smooth_iterations=20, smooth_lambda=0.5):
    """Nearest-surface warp: move each garment vertex from its offset
    relative to ref_body onto the equivalent offset relative to tgt_body,
    then Laplacian-smooth the displacement field over the garment mesh."""
    bvh_ref = BVHTree.FromObject(ref_body, depsgraph)
    bvh_tgt = BVHTree.FromObject(tgt_body, depsgraph)

    verts = garment.data.vertices
    n = len(verts)
    disp = [None] * n
    for i, v in enumerate(verts):
        co_r, n_r, _idx, _d = bvh_ref.find_nearest(v.co)
        if co_r is None:
            disp[i] = Vector((0, 0, 0))
            continue
        delta = v.co - co_r
        h = delta.dot(n_r)
        tang = delta - h * n_r
        co_t, n_t, _idx2, _d2 = bvh_tgt.find_nearest(co_r)
        if co_t is None:
            disp[i] = Vector((0, 0, 0))
            continue
        new_co = co_t + h * n_t + tang
        disp[i] = new_co - v.co

    # adjacency from edges
    adj = [[] for _ in range(n)]
    for e in garment.data.edges:
        a, b = e.vertices
        adj[a].append(b)
        adj[b].append(a)

    # Laplacian smoothing of the displacement field (not the geometry)
    for _ in range(smooth_iterations):
        new_disp = [None] * n
        for i in range(n):
            if not adj[i]:
                new_disp[i] = disp[i]
                continue
            avg = Vector((0, 0, 0))
            for j in adj[i]:
                avg += disp[j]
            avg /= len(adj[i])
            new_disp[i] = disp[i].lerp(avg, smooth_lambda)
        disp = new_disp

    max_disp = 0.0
    for i, v in enumerate(verts):
        v.co = v.co + disp[i]
        if disp[i].length > max_disp:
            max_disp = disp[i].length
    garment.data.update()
    return {"max_displacement_mm": round(max_disp * 1000, 1),
            "smooth_iterations": smooth_iterations}


def signed_distance_stats(garment, body_obj, depsgraph):
    """Signed distance of each garment vertex to the body surface.

    Positive = outside the body (along surface normal), negative = inside.
    """
    bvh = BVHTree.FromObject(body_obj, depsgraph)
    dists = []
    for v in garment.data.vertices:
        co, normal, _idx, _d = bvh.find_nearest(v.co)
        if co is None:
            continue
        dists.append((v.co - co).dot(normal))
    dists.sort()
    n = len(dists)
    inside = [d for d in dists if d < -PENETRATION_MM / 1000.0]
    return {
        "verts_sampled": n,
        "min_signed_mm": round(dists[0] * 1000, 2),
        "p05_signed_mm": round(dists[max(0, int(n * 0.05) - 1)] * 1000, 2),
        "median_signed_mm": round(dists[n // 2] * 1000, 2),
        "max_signed_mm": round(dists[-1] * 1000, 2),
        f"verts_inside_gt_{PENETRATION_MM:g}mm": len(inside),
        "pct_inside": round(100.0 * len(inside) / n, 2),
    }


log = {"garment_glb": GARMENT_GLB, "bodies": {}}
os.makedirs(OUT_DIR, exist_ok=True)

for body in BODY_TYPES:
    print(f"\n{'='*60}\n  BODY: {body}\n{'='*60}")
    entry = {}
    clear_scene()

    # --- 1. garment + reference landmarks ---
    g_objs = import_new(bpy.ops.import_scene.gltf, GARMENT_GLB)
    ref_arm = find_armature(g_objs)
    # the GLB also carries a bounds icosphere — the garment is the skinned
    # mesh (has vertex groups); the set-diff order is nondeterministic, so
    # pick explicitly and delete the rest
    g_meshes = [o for o in g_objs if o.type == 'MESH']
    garment = max((m for m in g_meshes if len(m.vertex_groups) > 0),
                  key=lambda m: len(m.data.vertices))
    ref_pelvis = bone_head_world(ref_arm, "pelvis")
    ref_upperarm = bone_head_world(ref_arm, "upperarm_l")
    freeze(garment)
    for m in g_meshes:
        if m is not garment:
            delete_obj(m)
    delete_obj(ref_arm)
    entry["garment_verts"] = len(garment.data.vertices)

    # --- 2. reference body (average_m) ---
    ref_body, _, _ = import_body("average_m", keep_armature=False)
    ref_body.name = "ref_body"

    # --- 3. target body ---
    tgt_body, tgt_arm, used = import_body(body, keep_armature=True)
    tgt_body.name = "tgt_body"
    entry["segments_used"] = used
    if tgt_arm is None:
        entry["result"] = "FAIL: no target armature"
        log["bodies"][body] = entry
        continue
    tgt_pelvis = bone_head_world(tgt_arm, "pelvis")
    tgt_upperarm = bone_head_world(tgt_arm, "upperarm_l")

    # --- 4. landmark affine ref -> tgt ---
    sxy = tgt_upperarm.x / ref_upperarm.x
    sz = (tgt_upperarm.z - tgt_pelvis.z) / (ref_upperarm.z - ref_pelvis.z)
    S = Matrix.Diagonal(Vector((sxy, sxy, sz, 1.0)))
    scaled_pelvis = Vector((ref_pelvis.x * sxy, ref_pelvis.y * sxy,
                            ref_pelvis.z * sz))
    t = tgt_pelvis - scaled_pelvis
    M = Matrix.Translation(t) @ S
    garment.data.transform(M)
    garment.data.update()
    ref_body.data.transform(M)
    ref_body.data.update()
    entry["affine"] = {"scale_xy": round(sxy, 5), "scale_z": round(sz, 5),
                       "translate": [round(v, 5) for v in t]}

    # post-affine baseline for distortion metrics
    base_edges = edge_lengths(garment)
    lo, hi = aabb(garment)
    entry["aabb_post_affine"] = {"min": [round(v, 4) for v in lo],
                                 "max": [round(v, 4) for v in hi]}

    # --- 5. BVH nearest-surface warp ref_body -> tgt_body ---
    depsgraph = bpy.context.evaluated_depsgraph_get()
    entry["warp"] = warp_garment(garment, ref_body, tgt_body, depsgraph)

    delete_obj(ref_body)

    # --- 7. weights from the target body ---
    for vg in list(garment.vertex_groups):
        garment.vertex_groups.remove(vg)
    deselect_all()
    set_active(garment)
    garment.select_set(True)
    dt = garment.modifiers.new(name="WeightTransfer", type='DATA_TRANSFER')
    dt.object = tgt_body
    dt.use_vert_data = True
    dt.data_types_verts = {'VGROUP_WEIGHTS'}
    dt.vert_mapping = 'POLYINTERP_NEAREST'
    dt.layers_vgroup_select_src = 'ALL'
    dt.layers_vgroup_select_dst = 'NAME'
    bpy.ops.object.datalayout_transfer(modifier=dt.name)
    bpy.ops.object.modifier_apply(modifier=dt.name)
    bpy.ops.object.mode_set(mode='WEIGHT_PAINT')
    bpy.ops.object.vertex_group_normalize_all(lock_active=False)
    bpy.ops.object.mode_set(mode='OBJECT')

    zero = sum(1 for v in garment.data.vertices
               if sum(g.weight for g in v.groups) < 1e-4)
    entry["zero_weight_verts"] = zero
    entry["vgroups"] = len(garment.vertex_groups)

    # --- 8. metrics ---
    post_edges = edge_lengths(garment)
    ratios = [p / b for p, b in zip(post_edges, base_edges) if b > 1e-9]
    entry["edge_distortion"] = {
        "max_stretch": round(max(ratios), 3),
        "max_shrink": round(min(ratios), 3),
        "edges_gt_2x": sum(1 for r in ratios if r > 2.0),
        "edges_lt_0.5x": sum(1 for r in ratios if r < 0.5),
        "edges_total": len(ratios),
    }
    lo, hi = aabb(garment)
    entry["aabb_post_fit"] = {"min": [round(v, 4) for v in lo],
                              "max": [round(v, 4) for v in hi]}
    depsgraph = bpy.context.evaluated_depsgraph_get()
    entry["signed_distance"] = signed_distance_stats(garment, tgt_body,
                                                     depsgraph)

    # --- 9. bind + export ---
    garment.parent = tgt_arm
    garment.matrix_parent_inverse = tgt_arm.matrix_world.inverted()
    am = garment.modifiers.new(name="Armature", type='ARMATURE')
    am.object = tgt_arm
    delete_obj(tgt_body)

    out_glb = os.path.join(OUT_DIR,
                           f"SK_SCFI_CIVL_09_10TORS_HU01_{body}.glb")
    deselect_all()
    tgt_arm.select_set(True)
    garment.select_set(True)
    set_active(tgt_arm)
    bpy.ops.export_scene.gltf(
        filepath=out_glb,
        export_format='GLB',
        use_selection=True,
        export_apply=False,
        export_animations=False,
        export_skins=True,
        export_morph=False,
        export_extras=False,
    )
    entry["export"] = {"path": out_glb,
                       "size_kb": os.path.getsize(out_glb) // 1024}
    entry["result"] = "OK"
    log["bodies"][body] = entry

with open(OUT_LOG, "w") as f:
    json.dump(log, f, indent=2)

print("\n=== BATCH FIT SUMMARY ===")
for body, e in log["bodies"].items():
    sd_stats = e.get("signed_distance", {})
    print(f"  {body:12s} {e.get('result','?'):14s} "
          f"zero_w={e.get('zero_weight_verts','-')} "
          f"inside%={sd_stats.get('pct_inside','-')} "
          f"stretch={e.get('edge_distortion',{}).get('max_stretch','-')}")