#!/usr/bin/env python3
"""
Grow skirt and hair BONE CHAINS on a model that shipped without any.

A character whose source arrived unrigged goes through tools/autorig.py, which
fits a body skeleton and weights everything to it. The body then animates
correctly and the costume does not move at all: hair is welded to the skull and a
skirt is welded to the hips, because there is nothing there to move them. Miku
ships exactly like this — 0 cloth chains against Taila's 35 — so her twin tails
hang off her head like a helmet.

Nothing downstream can fix that. characters/spring_bones.gd simulates cloth
BONES, and if the garment has none there is nothing to simulate. This builds
them, which is the job a technical artist does by hand on a model like this.

The chains it grows are the same shape the Hoyoverse-class rigs use, because that
shape is dictated by what the geometry is:

  HAIR    one chain per strand, found as connected islands of hair geometry,
          each running from the scalp down the strand's own curve.
  SKIRT   a radial grid — the garment is split into panels around the body's up
          axis and each panel gets a chain from the waistband to the hem, which
          is what lets a skirt open around a leg instead of swinging as a bell.

Vertices are then re-weighted onto the new chain with a falloff that keeps the
anchored end (scalp, waistband) on the body, so nothing detaches.

Usage:
  blender --background --python tools/cloth_bones.py -- \
      <rigged.glb> <output.glb> [--hair-segments 3] [--skirt-segments 4]
      [--skirt-panels 12] [--classes hair,skirt]

Run this BEFORE tools/retarget.py. retarget.py's describe_rig() then finds the
chains by name exactly as it would an artist's, and writes them to the sidecar.
"""
import bpy
import bmesh
import json
import math
import os
import sys
from collections import defaultdict, deque
from mathutils import Matrix, Vector

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import gltf_fix
import rig_map
from rig_map import RigRoles

argv = sys.argv
argv = argv[argv.index("--") + 1:] if "--" in argv else []
if len(argv) < 2:
    print(__doc__)
    sys.exit(1)

SOURCE, OUTPUT = argv[0], argv[1]


def _opt(flag, default):
    return type(default)(argv[argv.index(flag) + 1]) if flag in argv else default


HAIR_SEGMENTS = _opt("--hair-segments", 3)
SKIRT_SEGMENTS = _opt("--skirt-segments", 4)
SKIRT_PANELS = _opt("--skirt-panels", 12)
CLASSES = _opt("--classes", "hair,skirt").split(",")

# A strand shorter than this is a fringe or an ornament, not something that
# should swing. Simulating them costs the same as simulating a ponytail and only
# ever produces jitter around the face.
MIN_STRAND_LENGTH = 0.06
MIN_STRAND_VERTS = 12
# How much of the anchored end stays welded to the body. A strand's first
# fraction blends from "entirely the body bone" to "entirely the chain", so the
# scalp and the waistband never separate from the character.
ROOT_BLEND = 0.22


def clear_scene():
    bpy.ops.wm.read_factory_settings(use_empty=True)


def find_rig():
    arms = [o for o in bpy.data.objects if o.type == "ARMATURE"]
    if not arms:
        print("ERROR: no armature — run the auto-rig first")
        sys.exit(1)
    arm = max(arms, key=lambda a: len(a.data.bones))
    meshes = [o for o in bpy.data.objects if o.type == "MESH"
              and any(m.type == "ARMATURE" and m.object == arm for m in o.modifiers)]
    if not meshes:
        meshes = [o for o in bpy.data.objects if o.type == "MESH"]
    return arm, meshes


def cosmetic_slots(mesh, wanted):
    """Material slot indices whose name says what that geometry IS.

    The artist already answered this. A slot called "hair" is hair; guessing from
    position or from how far something hangs would be a guess, and on a joined
    mesh — which is what the auto-rig leaves behind — the material slot is the
    ONLY separation left.
    """
    out = {}
    for i, slot in enumerate(mesh.material_slots):
        name = (slot.name or "").lower()
        for cls in wanted:
            if cls in name:
                out[i] = cls
    return out


def islands(mesh, vert_ids):
    """Connected components of `vert_ids`, over the mesh's own edges.

    A strand of hair is a connected piece of surface. Clustering by position
    instead would merge two ponytails that pass near each other and split a
    single one that bends.
    """
    adj = defaultdict(list)
    keep = set(vert_ids)
    for e in mesh.data.edges:
        a, b = e.vertices
        if a in keep and b in keep:
            adj[a].append(b)
            adj[b].append(a)
    seen, out = set(), []
    for v in vert_ids:
        if v in seen:
            continue
        comp, q = [], deque([v])
        seen.add(v)
        while q:
            n = q.popleft()
            comp.append(n)
            for m in adj[n]:
                if m not in seen:
                    seen.add(m)
                    q.append(m)
        out.append(comp)
    return out


def radial_panels(mesh, vert_ids, centre, up, count):
    """Split a garment into `count` wedges around the body's up axis.

    This is the skirt grid. A panel is a wedge rather than a connected island
    because a skirt IS one connected surface — islands would return the whole
    thing as a single piece, which is the bell-shaped failure.
    """
    ref = Vector((1.0, 0.0, 0.0))
    ref = (ref - up * ref.dot(up)).normalized()
    side = up.cross(ref).normalized()
    out = defaultdict(list)
    for v in vert_ids:
        d = mesh.data.vertices[v].co - centre
        d = d - up * d.dot(up)
        if d.length < 1e-6:
            continue
        ang = math.atan2(d.dot(side), d.dot(ref))
        out[int((ang + math.pi) / (2 * math.pi) * count) % count].append(v)
    return [v for v in out.values() if len(v) >= 6]


def fit_polyline(mesh, vert_ids, root_point, segments):
    """A polyline down the middle of a clump of geometry.

    Vertices are binned by distance from the anchored end and each bin's centroid
    becomes a joint, so the chain follows the piece's own CURVE. A straight line
    from root to tip would cut the corner on a ponytail that bends, and every
    vertex on the outside of that bend would then be weighted to a bone travelling
    the wrong way.
    """
    co = [mesh.data.vertices[v].co for v in vert_ids]
    d = [(c - root_point).length for c in co]
    lo, hi = min(d), max(d)
    if hi - lo < 1e-5:
        return []
    pts = [root_point.copy()]
    for s in range(1, segments + 1):
        a = lo + (hi - lo) * (s - 1) / segments
        b = lo + (hi - lo) * s / segments
        bucket = [c for c, dist in zip(co, d) if a <= dist <= b + 1e-9]
        if not bucket:
            # An empty band means the geometry does not reach here; carry the
            # direction on rather than collapsing the bone to zero length, which
            # would drop it out of the chain at runtime.
            if len(pts) >= 2:
                pts.append(pts[-1] + (pts[-1] - pts[-2]))
            continue
        pts.append(sum(bucket, Vector()) / len(bucket))
    return pts if len(pts) >= 2 else []


def nearest_body_group(mesh, vert_ids, arm):
    """Which body bone this clump is currently welded to.

    That bone is the right parent for the new chain, and its weight is what the
    root blend fades out of — so the transition is to exactly what was holding
    this geometry before.
    """
    tally = defaultdict(float)
    gname = {g.index: g.name for g in mesh.vertex_groups}
    bones = set(b.name for b in arm.data.bones)
    for v in vert_ids:
        for g in mesh.data.vertices[v].groups:
            n = gname.get(g.group, "")
            if n in bones:
                tally[n] += g.weight
    return max(tally, key=tally.get) if tally else None


def make_chain(arm, name, points, parent_name):
    """Create one bone chain along `points`. Returns the bone names."""
    bpy.context.view_layer.objects.active = arm
    bpy.ops.object.mode_set(mode="EDIT")
    made = []
    prev = arm.data.edit_bones.get(parent_name)
    for i in range(len(points) - 1):
        bn = name if i == 0 else "%s.seg%d" % (name, i)
        eb = arm.data.edit_bones.new(bn)
        eb.head = points[i]
        eb.tail = points[i + 1]
        if eb.length < 1e-5:
            arm.data.edit_bones.remove(eb)
            continue
        eb.parent = prev
        eb.use_connect = False
        prev = eb
        made.append(bn)
    bpy.ops.object.mode_set(mode="OBJECT")
    return made


def weight_chain(mesh, vert_ids, points, bone_names, anchor_bone):
    """Move this clump's weights off the body and onto its new chain.

    Each vertex is projected onto the polyline and given the two bones either side
    of where it lands, blended by how far between them it is — the same linear
    split a hand-painted chain gets. Near the anchored end the ORIGINAL body
    weight is kept and faded out over ROOT_BLEND, so the scalp stays on the skull
    and the waistband stays on the hips.
    """
    groups = []
    for bn in bone_names:
        groups.append(mesh.vertex_groups.get(bn) or mesh.vertex_groups.new(name=bn))
    anchor = mesh.vertex_groups.get(anchor_bone) if anchor_bone else None

    # Cumulative length along the polyline, so position is measured in metres
    # rather than in segment index — segments are not equal lengths.
    seg_len = [(points[i + 1] - points[i]).length for i in range(len(points) - 1)]
    total = sum(seg_len) or 1.0

    for v in vert_ids:
        co = mesh.data.vertices[v].co
        best, best_d, best_run = 0, 1e18, 0.0
        run = 0.0
        for i in range(len(points) - 1):
            a, b = points[i], points[i + 1]
            ab = b - a
            L2 = ab.length_squared
            t = 0.0 if L2 < 1e-12 else max(0.0, min(1.0, (co - a).dot(ab) / L2))
            p = a + ab * t
            d = (co - p).length
            if d < best_d:
                best_d, best, best_run = d, i, run + seg_len[i] * t
            run += seg_len[i]
        along = best_run / total

        # Clear whatever held this vertex before, then re-add the anchor share.
        for g in mesh.data.vertices[v].groups:
            grp = mesh.vertex_groups[g.group]
            if grp.name not in bone_names:
                grp.remove([v])
        hold = 0.0
        if anchor and along < ROOT_BLEND:
            hold = 1.0 - along / ROOT_BLEND
            anchor.add([v], hold, "REPLACE")

        # Split what is left between the two bones either side of the landing.
        f = best_run / total * len(bone_names)
        i0 = max(0, min(len(bone_names) - 1, int(f)))
        i1 = min(len(bone_names) - 1, i0 + 1)
        frac = f - i0
        groups[i0].add([v], (1.0 - hold) * (1.0 - frac), "REPLACE")
        if i1 != i0:
            groups[i1].add([v], (1.0 - hold) * frac, "REPLACE")


def grow(arm, meshes, roles):
    head = roles.head
    trunk = roles.spine[0] if roles.spine else roles.hips
    made_chains = 0
    made_bones = 0

    for mesh in meshes:
        slots = cosmetic_slots(mesh, CLASSES)
        if not slots:
            continue
        # Vertices per class, taken from the polygons that use each slot.
        per_class = defaultdict(set)
        for poly in mesh.data.polygons:
            cls = slots.get(poly.material_index)
            if cls:
                per_class[cls].update(poly.vertices)

        for cls, verts in per_class.items():
            verts = list(verts)
            if len(verts) < MIN_STRAND_VERTS:
                continue
            world = mesh.matrix_world
            if cls == "hair":
                anchor_default = head
                clumps = islands(mesh, verts)
            else:
                anchor_default = trunk
                centre = sum((mesh.data.vertices[v].co for v in verts),
                             Vector()) / len(verts)
                clumps = radial_panels(mesh, verts, centre,
                                       Vector((0.0, 0.0, 1.0)), SKIRT_PANELS)

            idx = 0
            for clump in clumps:
                if len(clump) < MIN_STRAND_VERTS:
                    continue
                anchor = nearest_body_group(mesh, clump, arm) or anchor_default
                if anchor is None or anchor not in arm.data.bones:
                    continue
                # The anchored end is the end nearest the bone currently holding
                # it — the scalp for hair, the waistband for a skirt.
                bone_head = (arm.matrix_world
                             @ arm.data.bones[anchor].matrix_local).translation
                bone_head = world.inverted() @ bone_head
                root = min((mesh.data.vertices[v].co for v in clump),
                           key=lambda c: (c - bone_head).length)
                segs = HAIR_SEGMENTS if cls == "hair" else SKIRT_SEGMENTS
                pts = fit_polyline(mesh, clump, root, segs)
                if not pts:
                    continue
                span = sum((pts[i + 1] - pts[i]).length for i in range(len(pts) - 1))
                if span < MIN_STRAND_LENGTH:
                    continue
                name = "DEF-%s.%03d" % (cls, idx)
                idx += 1
                names = make_chain(arm, name, [world @ p for p in pts], anchor)
                if not names:
                    continue
                weight_chain(mesh, clump, pts, names, anchor)
                made_chains += 1
                made_bones += len(names)
            print("  %-6s %-22s %d clumps -> %d chains"
                  % (cls, mesh.name, len(clumps), idx))
    return made_chains, made_bones


def main():
    clear_scene()
    bpy.ops.import_scene.gltf(filepath=gltf_fix.prepare(SOURCE,
                                                        os.path.dirname(OUTPUT)))
    arm, meshes = find_rig()
    roles = RigRoles(arm)
    missing = roles.missing_core()
    if missing:
        print("ERROR: could not identify %s on this rig" % missing)
        sys.exit(1)

    before = len(arm.data.bones)
    chains, bones = grow(arm, meshes, roles)
    print("Grew %d cloth chains (%d bones); armature %d -> %d bones"
          % (chains, bones, before, len(arm.data.bones)))
    if chains == 0:
        print("NOTE: no cosmetic material slots matched %s — nothing to grow. "
              "Check the model's material names." % CLASSES)

    os.makedirs(os.path.dirname(os.path.abspath(OUTPUT)), exist_ok=True)
    bpy.ops.export_scene.gltf(
        filepath=OUTPUT,
        export_format="GLB",
        export_yup=True,
        export_apply=False,
        export_skins=True,
        export_animations=False,
    )
    print("Wrote", OUTPUT)


main()