"""
FAILED (Trellis clothing pipeline): Auto-rig a Trellis-generated clothing mesh
onto a Quaternius character skeleton.

The weight transfer algorithm works correctly, but the Trellis-generated clothing
meshes themselves are unsuitable for rigging: non-manifold geometry, inconsistent
vertex density, and topology that does not deform well under skinning. The
auto-rigging pipeline is technically sound but the INPUT meshes are the problem.

Conclusion: Trellis CANNOT generate riggable clothing. Use hand-authored or
Quaternius-pack clothing instead.

This script is preserved as documentation of the approach and the robust weight
transfer implementation (which may be reusable for other mesh sources).

Uses robust weight transfer with Laplacian inpainting for unmatched vertices,
based on the SIGGRAPH Asia 2023 algorithm (MIT reference implementation).

Run via:
    reach blender run \\
        spikes/quaternius-aesthetic/scripts/blender/auto_rig_clothing.py \\
        -- body.gltf clothing.glb output.glb

Pipeline: Import -> Cleanup -> Scale/Position -> Weight Transfer -> Normalize -> Export
"""

import bpy
import site
import sys
import os

# Blender flatpak installs --user packages outside the default path
sys.path.insert(0, site.getusersitepackages())

import numpy as np
from mathutils import Vector
import igl
import scipy.sparse as sp
import robust_laplacian


# --- Parse args ---

argv = sys.argv
argv = argv[argv.index("--") + 1:] if "--" in argv else []
if len(argv) < 3:
    print("Usage: -- <body.gltf> <clothing.glb> <output.glb>")
    sys.exit(1)

body_path, clothing_path, output_path = argv[0], argv[1], argv[2]

print("=== Auto-rig clothing (robust weight transfer) ===")
print(f"  Body:     {body_path}")
print(f"  Clothing: {clothing_path}")
print(f"  Output:   {output_path}")


# --- Helpers ---

def clear_scene():
    bpy.ops.object.select_all(action='SELECT')
    bpy.ops.object.delete()
    for col in list(bpy.data.collections):
        bpy.data.collections.remove(col)


def import_glb(path):
    before = set(bpy.data.objects)
    bpy.ops.import_scene.gltf(filepath=path)
    return list(set(bpy.data.objects) - before)


def find_armature(objects):
    for obj in objects:
        if obj.type == 'ARMATURE':
            return obj
    return None


def find_meshes(objects):
    return [obj for obj in objects if obj.type == 'MESH']


def get_bounds(obj):
    corners = [obj.matrix_world @ Vector(c) for c in obj.bound_box]
    mins = Vector((min(c.x for c in corners), min(c.y for c in corners), min(c.z for c in corners)))
    maxs = Vector((max(c.x for c in corners), max(c.y for c in corners), max(c.z for c in corners)))
    return mins, maxs


def mesh_to_numpy(obj):
    """Extract world-space vertices, faces, and normals as numpy arrays."""
    mesh = obj.data
    mesh.calc_loop_triangles()
    mw = obj.matrix_world

    verts = np.array([mw @ v.co for v in mesh.vertices], dtype=np.float64)
    faces = np.array([[lt.vertices[i] for i in range(3)] for lt in mesh.loop_triangles], dtype=np.int64)
    normals = np.array([mw.to_3x3() @ v.normal for v in mesh.vertices], dtype=np.float64)

    return verts, faces, normals


def get_bone_weights(obj, bone_names):
    """Extract per-vertex bone weights as a (n_verts x n_bones) matrix."""
    n_verts = len(obj.data.vertices)
    n_bones = len(bone_names)
    weights = np.zeros((n_verts, n_bones), dtype=np.float64)

    # Map vertex group names to bone indices
    vg_to_bone = {}
    for vg in obj.vertex_groups:
        if vg.name in bone_names:
            vg_to_bone[vg.index] = bone_names.index(vg.name)

    for v in obj.data.vertices:
        for g in v.groups:
            if g.group in vg_to_bone:
                weights[v.index, vg_to_bone[g.group]] = g.weight

    return weights


def set_bone_weights(obj, bone_names, weights):
    """Set per-vertex bone weights from a (n_verts x n_bones) matrix."""
    # Create vertex groups
    for name in bone_names:
        if name not in obj.vertex_groups:
            obj.vertex_groups.new(name=name)

    vg_map = {name: obj.vertex_groups[name] for name in bone_names if name in obj.vertex_groups}

    for vi in range(len(obj.data.vertices)):
        for bi, name in enumerate(bone_names):
            w = weights[vi, bi]
            if w > 0.001:
                vg_map[name].add([vi], w, 'REPLACE')


# --- Pipeline steps ---

def cleanup_trellis_mesh(obj):
    """Fix common Trellis output issues."""
    print("\n  Cleanup:")
    verts_before = len(obj.data.vertices)
    bpy.context.view_layer.objects.active = obj
    obj.select_set(True)
    bpy.ops.object.mode_set(mode='EDIT')

    bpy.ops.mesh.select_all(action='SELECT')
    bpy.ops.mesh.remove_doubles(threshold=0.001)
    bpy.ops.mesh.normals_make_consistent(inside=False)
    bpy.ops.mesh.select_all(action='DESELECT')
    bpy.ops.mesh.select_loose()
    bpy.ops.mesh.delete(type='VERT')

    bpy.ops.object.mode_set(mode='OBJECT')
    print(f"    {verts_before} → {len(obj.data.vertices)} verts")


def scale_and_position(clothing_obj, body_obj, armature):
    """Scale clothing to match body and center on spine."""
    print("\n  Scale & position:")

    body_mins, body_maxs = get_bounds(body_obj)
    cloth_mins, cloth_maxs = get_bounds(clothing_obj)
    body_size = body_maxs - body_mins
    cloth_size = cloth_maxs - cloth_mins

    # Uniform scale based on body width
    torso_width = body_size.x * 0.85
    scale = torso_width / max(cloth_size.x, 0.001)
    clothing_obj.scale = (scale, scale, scale)
    bpy.context.view_layer.update()

    # Center on spine_02 bone
    bone = armature.data.bones.get("spine_02")
    if bone:
        target = armature.matrix_world @ bone.head_local
    else:
        target = Vector(((body_mins.x + body_maxs.x) / 2, body_mins.y + body_size.y * 0.55, (body_mins.z + body_maxs.z) / 2))

    cloth_mins2, cloth_maxs2 = get_bounds(clothing_obj)
    cloth_center = (cloth_mins2 + cloth_maxs2) / 2
    clothing_obj.location += target - cloth_center

    # Apply transforms
    bpy.ops.object.select_all(action='DESELECT')
    clothing_obj.select_set(True)
    bpy.context.view_layer.objects.active = clothing_obj
    bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)

    cloth_mins3, cloth_maxs3 = get_bounds(clothing_obj)
    print(f"    Scale: {scale:.3f}, bounds: {cloth_mins3.y:.2f}–{cloth_maxs3.y:.2f}")


def robust_weight_transfer(body_obj, clothing_obj, bone_names,
                           dist_threshold=0.1, normal_threshold_deg=90.0):
    """
    Transfer bone weights from body to clothing using closest-point matching
    with Laplacian inpainting for unmatched vertices.

    Based on: "Robust Skin Weights Transfer via Weight Inpainting"
    (Abdrashitov et al., SIGGRAPH Asia 2023, MIT reference implementation)
    """
    print("\n  Robust weight transfer:")

    # Extract mesh data
    src_verts, src_faces, src_normals = mesh_to_numpy(body_obj)
    tgt_verts, tgt_faces, tgt_normals = mesh_to_numpy(clothing_obj)

    # Get source weights
    src_weights = get_bone_weights(body_obj, bone_names)

    print(f"    Source: {len(src_verts)} verts, {len(src_faces)} faces")
    print(f"    Target: {len(tgt_verts)} verts, {len(tgt_faces)} faces")
    print(f"    Bones: {len(bone_names)}")

    # Step 1: Find closest point on source surface for each target vertex
    sqr_dist, face_idx, closest_pts = igl.point_mesh_squared_distance(
        tgt_verts, src_verts, src_faces
    )
    distances = np.sqrt(sqr_dist)

    # Step 2: Compute barycentric coordinates for interpolation
    tgt_weights = np.zeros((len(tgt_verts), len(bone_names)), dtype=np.float64)
    matched = np.zeros(len(tgt_verts), dtype=bool)

    normal_threshold = np.cos(np.radians(normal_threshold_deg))

    for vi in range(len(tgt_verts)):
        if distances[vi] > dist_threshold:
            continue

        fi = face_idx[vi]
        tri_verts = src_faces[fi]

        # Check normal compatibility
        src_normal = np.mean(src_normals[tri_verts], axis=0)
        src_normal /= max(np.linalg.norm(src_normal), 1e-10)
        tgt_normal = tgt_normals[vi]
        tgt_normal /= max(np.linalg.norm(tgt_normal), 1e-10)
        dot = np.dot(src_normal, tgt_normal)

        if dot < normal_threshold:
            continue

        # Barycentric interpolation of weights
        p = closest_pts[vi]
        a, b, c = src_verts[tri_verts[0]], src_verts[tri_verts[1]], src_verts[tri_verts[2]]

        # Compute barycentric coords
        v0, v1, v2 = b - a, c - a, p - a
        d00, d01, d11 = np.dot(v0, v0), np.dot(v0, v1), np.dot(v1, v1)
        d20, d21 = np.dot(v2, v0), np.dot(v2, v1)
        denom = d00 * d11 - d01 * d01
        if abs(denom) < 1e-10:
            continue
        bary_v = (d11 * d20 - d01 * d21) / denom
        bary_w = (d00 * d21 - d01 * d20) / denom
        bary_u = 1.0 - bary_v - bary_w

        # Interpolate source weights
        tgt_weights[vi] = (bary_u * src_weights[tri_verts[0]] +
                           bary_v * src_weights[tri_verts[1]] +
                           bary_w * src_weights[tri_verts[2]])
        matched[vi] = True

    n_matched = np.sum(matched)
    n_unmatched = len(tgt_verts) - n_matched
    print(f"    Matched: {n_matched}/{len(tgt_verts)} ({100*n_matched/max(len(tgt_verts),1):.0f}%)")
    print(f"    Unmatched: {n_unmatched} (will inpaint)")

    # Step 3: Laplacian inpainting for unmatched vertices
    if n_unmatched > 0 and n_matched > 0:
        print("    Computing Laplacian inpainting...")
        L, M = robust_laplacian.mesh_laplacian(tgt_verts, tgt_faces)

        # Solve per bone: minimize ||L @ w||^2 subject to matched vertices = known values
        unmatched_idx = np.where(~matched)[0]
        matched_idx = np.where(matched)[0]

        for bi in range(len(bone_names)):
            known_weights = tgt_weights[matched_idx, bi]

            # Build system: L[unmatched, unmatched] @ w_unknown = -L[unmatched, matched] @ w_known
            L_uu = L[np.ix_(unmatched_idx, unmatched_idx)]
            L_um = L[np.ix_(unmatched_idx, matched_idx)]

            rhs = -L_um @ known_weights

            if L_uu.shape[0] > 0:
                try:
                    result = sp.linalg.spsolve(L_uu, rhs)
                    tgt_weights[unmatched_idx, bi] = np.clip(result, 0.0, 1.0)
                except Exception:
                    # Fallback: use nearest matched vertex weight
                    for ui in unmatched_idx:
                        dists_to_matched = np.linalg.norm(tgt_verts[matched_idx] - tgt_verts[ui], axis=1)
                        nearest = matched_idx[np.argmin(dists_to_matched)]
                        tgt_weights[ui, bi] = tgt_weights[nearest, bi]

        print("    Inpainting complete")

    # Step 4: Normalize weights per vertex
    row_sums = tgt_weights.sum(axis=1, keepdims=True)
    row_sums[row_sums < 1e-10] = 1.0  # avoid division by zero
    tgt_weights /= row_sums

    # Step 5: Limit to 4 bones per vertex (game engine constraint)
    for vi in range(len(tgt_verts)):
        w = tgt_weights[vi]
        if np.count_nonzero(w > 0.001) > 4:
            top4 = np.argsort(w)[-4:]
            mask = np.zeros_like(w)
            mask[top4] = w[top4]
            mask /= max(mask.sum(), 1e-10)
            tgt_weights[vi] = mask

    # Apply to clothing mesh
    set_bone_weights(clothing_obj, bone_names, tgt_weights)

    active_bones = np.sum(tgt_weights.max(axis=0) > 0.01)
    print(f"    Active bones: {active_bones}/{len(bone_names)}")


def export_result(clothing_obj, armature, path):
    """Export clothing + armature as GLB."""
    print(f"\n  Exporting to {path}...")

    for obj in bpy.data.objects:
        obj.hide_set(True)
    clothing_obj.hide_set(False)
    armature.hide_set(False)

    bpy.ops.object.select_all(action='DESELECT')
    clothing_obj.select_set(True)
    armature.select_set(True)

    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    bpy.ops.export_scene.gltf(
        filepath=path,
        export_format='GLB',
        use_selection=True,
        export_apply=False,
        export_animations=False,
        export_skins=True,
    )
    print(f"    Size: {os.path.getsize(path)} bytes")


# --- Main ---

clear_scene()

print("\nStep 1: Import body...")
body_objects = import_glb(body_path)
armature = find_armature(body_objects)
body_meshes = find_meshes(body_objects)
body_mesh = max(body_meshes, key=lambda m: len(m.data.vertices))
bone_names = [b.name for b in armature.data.bones]
print(f"  Armature: {armature.name} ({len(bone_names)} bones)")
print(f"  Body mesh: {body_mesh.name} ({len(body_mesh.data.vertices)} verts)")

print("\nStep 2: Import clothing...")
clothing_objects = import_glb(clothing_path)
clothing_meshes = find_meshes(clothing_objects)

if len(clothing_meshes) > 1:
    bpy.ops.object.select_all(action='DESELECT')
    for m in clothing_meshes:
        m.select_set(True)
    bpy.context.view_layer.objects.active = clothing_meshes[0]
    bpy.ops.object.join()

clothing_mesh = clothing_meshes[0]
print(f"  Clothing mesh: {clothing_mesh.name} ({len(clothing_mesh.data.vertices)} verts)")

print("\nStep 3: Cleanup...")
cleanup_trellis_mesh(clothing_mesh)

print("\nStep 4: Scale & position...")
scale_and_position(clothing_mesh, body_mesh, armature)

print("\nStep 5: Robust weight transfer...")
# Parent to armature first
clothing_mesh.parent = armature
clothing_mesh.matrix_parent_inverse = armature.matrix_world.inverted()
arm_mod = clothing_mesh.modifiers.new(name="Armature", type='ARMATURE')
arm_mod.object = armature

robust_weight_transfer(body_mesh, clothing_mesh, bone_names,
                       dist_threshold=0.15, normal_threshold_deg=120.0)

print("\nStep 6: Export...")
export_result(clothing_mesh, armature, output_path)

print("\n=== Done ===")