Papay-Shooter/tools/rig_map.py

#!/usr/bin/env python3
"""
Work out which bone on one rig means the same thing as which bone on another.

This is the piece that lets us STOP throwing away a character's own skeleton.
The old pipeline discarded any foreign rig (tools/strip_rig.py) because
merge_animations.py retargeted by exact bone NAME, so a rig that named things
differently produced a rest-pose statue. Rebinding the mesh to the library
skeleton then destroyed everything an artist had done: per-part weights, skirt
and hair bone chains, limb twist bones. Solving the naming problem here is what
makes keeping the original rig possible.

Roles are resolved STRUCTURALLY wherever a name would lie:

  * `hips` is the lowest common ancestor of both thighs and the head, not
    whatever is called "hips". Rigify calls it `DEF-spine`; Mixamo calls it
    `mixamorig:Hips`; both land on the same bone this way.
  * The spine chain is walked from the hips upward, refusing to turn down a
    limb or a cosmetic chain (hair/skirt/face). The bone it ends on is the
    head. That matters because a stock Rigify rig has NO bone with "head" in
    its name — the head is `DEF-spine.006`.
  * Chains of different length are matched by normalised position along the
    chain, so a 4-bone torso drives a 3-bone one and vice versa.

Only bones the CLIPS need are mapped. Everything else on the character rig —
skirt chains, hair chains, twist bones, face bones — is deliberately left
unmapped so it rests relative to its parent and is free to be driven by
secondary motion at runtime. That is the whole point: the locomotion library
animates the body, physics animates the cloth.
"""
import re

# Names that are never part of the body proper. Walking the spine must not turn
# down one of these, and they must never claim a limb role.
#
# Matched as whole NAME TOKENS, never as substrings. A plain `"ear" in name`
# test marks every `DEF-forearm.L` cosmetic — which silently cost both rigs
# their forearms and is exactly the class of bug this file exists to avoid.
COSMETIC = ("hair", "skirt", "cloth", "ribbon", "tail", "cape", "coat",
            "scarf", "sleeve", "breast", "bust", "ear", "horn", "wing",
            "face", "cheek", "nose", "mouth", "eye", "brow", "jaw", "tongue",
            "teeth", "tooth", "lip", "chin", "accessory", "prop", "weapon",
            "bell", "strap", "belt", "buckle", "feather", "antenna")

# role -> ordered alternative stems. Ordered because "leg" must not win before
# "upleg" has had its chance: Mixamo's LeftUpLeg is a thigh and its LeftLeg is
# a shin, so the more specific spelling has to be tested first.
LIMB_ROLES = {
    "thigh": ("upleg", "upperleg", "thigh", "leg_upper", "upper_leg", "hip"),
    "shin":  ("lowerleg", "lowleg", "shin", "calf", "knee", "leg_lower", "leg"),
    "foot":  ("foot", "ankle"),
    "toe":   ("toebase", "toe", "ball"),
    "shoulder":  ("shoulder", "clavicle", "collar"),
    "upper_arm": ("upperarm", "upper_arm", "arm_upper", "armupper", "arm"),
    "forearm":   ("forearm", "lowerarm", "lowarm", "arm_lower", "elbow"),
    "hand":      ("hand", "wrist"),
}
# Longest-first inside each role, so "upperarm" is tried before "arm".
LIMB_ORDER = ["toe", "foot", "shin", "thigh", "shoulder", "hand", "forearm", "upper_arm"]

_PREFIXES = re.compile(
    r"^(def[-_]|org[-_]|mch[-_]|ctrl[-_]|mixamorig\d*[:_]|bip\d*[-_ ]|"
    r"j_bip_[clr]_|j_sec_[clr]_|valvebiped\.|bone_|b_)", re.I)
# The glTF importer appends _<node index> to every bone name; strip it so
# `DEF-spine.006_2` reads as `DEF-spine.006`.
_GLTF_SUFFIX = re.compile(r"_\d+$")


def strip_gltf_suffix(name: str) -> str:
    return _GLTF_SUFFIX.sub("", name)


def canon(name: str) -> str:
    """Bone name reduced to a comparable stem: no rig prefix, no separators."""
    n = strip_gltf_suffix(name).lower()
    while True:
        stripped = _PREFIXES.sub("", n)
        if stripped == n:
            break
        n = stripped
    return re.sub(r"[^a-z0-9]", "", n)


def side_of(name: str):
    """'L', 'R' or None. Checked on the ORIGINAL name so `.L` survives."""
    n = strip_gltf_suffix(name).lower()
    if re.search(r"(^|[._\- ])l($|[._\- 0-9])", n) or "left" in n:
        return "L"
    if re.search(r"(^|[._\- ])r($|[._\- 0-9])", n) or "right" in n:
        return "R"
    return None


def tokens(name: str):
    """Name split into alphabetic words: `DEF-hair.L.001` -> def, hair, l."""
    return [t for t in re.split(r"[^a-z]+", strip_gltf_suffix(name).lower()) if t]


def is_cosmetic(name: str) -> bool:
    """Whole-token match, plus a short positional suffix — see COSMETIC.

    A bare token match misses the very common habit of gluing a position onto
    the word: DANDADAN's hair bones are HairFL / HairFR / HairF_Top (front-left,
    front-right, front-top), which tokenise to "hairfl" and match nothing, so
    her hair was not detected as hair at all and none of it moved.

    Only a SHORT remainder counts — two characters at most, which covers l/r/f/b
    and the fl/fr/bl/br pairs and nothing else. That is what keeps the original
    rule intact: "forearm" still does not begin with any cosmetic stem, and
    "earring" has a four-character remainder and is not swept in by accident.
    """
    for t in tokens(name):
        if t in COSMETIC:
            return True
        for c in COSMETIC:
            if len(t) - len(c) <= 2 and t.startswith(c) and len(t) > len(c):
                return True
    return False


def is_segment_of(name: str, all_names: set) -> bool:
    """True for twist/segment bones like `DEF-upper_arm.L.001`.

    Rigify subdivides a limb into a main bone plus numbered twist bones. Those
    must never claim the limb's role — the clip would drive the twist bone and
    the real limb would stay put. Detected by construction rather than by a
    name list: strip a trailing `.001` and see whether the parent spelling is
    itself a bone on this rig.
    """
    base = strip_gltf_suffix(name)
    m = re.match(r"^(.*)\.(\d{3})$", base)
    if not m or m.group(2) == "000":
        return False
    stem = m.group(1)
    return any(strip_gltf_suffix(n) == stem for n in all_names)


class RigRoles:
    """Resolved anatomy of one armature."""

    def __init__(self, arm):
        self.arm = arm
        self.bones = {b.name: b for b in arm.data.bones}
        self.names = set(self.bones)
        self.limb = {}       # ("thigh", "L") -> bone name
        self.spine = []      # hips-exclusive chain, ordered, ending on the head
        self.hips = None
        self.head = None
        self.neck = None
        self._resolve_limbs()
        self._fill_limbs_by_anatomy()
        self._resolve_spine()

    # ---------------------------------------------------------------- limbs

    def _resolve_limbs(self):
        """Claim limb roles by the MOST SPECIFIC spelling that matches.

        Not by walking LIMB_ORDER and taking the first role that matches at all.
        That let a catch-all stem in an early role beat an exact one in a later
        role, and the result depended on the order bones happened to arrive in:

            J_Bip_L_UpperLeg   canon "upperleg"
                shin  matches "leg"        <- claimed it, because shin is first
                thigh matches "upperleg"   <- never consulted
            J_Bip_L_LowerLeg   canon "lowerleg"
                shin already taken; nothing in thigh matches -> unclaimed

        so a VRoid rig resolved with no thighs at all and the build stopped. The
        same trap catches Mixamo, whose LeftUpLeg is a thigh and LeftLeg a shin —
        it stayed hidden only because every character imported so far used the
        Rigify DEF- spellings, where each role has an exact stem of its own.

        Every plausible claim is collected, then granted longest-stem first, so
        "upperleg" always beats "leg" whatever order the bones are in.
        """
        claims = []
        for name in self.names:
            if is_cosmetic(name) or is_segment_of(name, self.names):
                continue
            side = side_of(name)
            if side is None:
                continue
            stem = canon(name)
            for role in LIMB_ORDER:
                hits = [alt.replace("_", "") for alt in LIMB_ROLES[role]
                        if alt.replace("_", "") in stem]
                if hits:
                    claims.append((len(max(hits, key=len)), role, side, name))
        # Longest stem wins; LIMB_ORDER breaks ties, so a bone matching two roles
        # equally well still resolves the way the list intends.
        order = {role: i for i, role in enumerate(LIMB_ORDER)}
        claims.sort(key=lambda c: (-c[0], order[c[1]]))
        used = set()
        for _, role, side, name in claims:
            if (role, side) in self.limb or name in used:
                continue
            self.limb[(role, side)] = name
            used.add(name)

    # A leg is thigh -> shin -> foot, whatever the bones are called. Walk it.
    #
    # Names alone cannot settle this and it is not a matter of adding more
    # spellings. A bare "leg" is the SHIN on a Mixamo rig (whose thigh is
    # "UpLeg") and the THIGH on a rig whose shin is called "knee" — the same
    # token means opposite bones, and both rigs are common. What does not vary
    # is the skeleton: the foot's parent is the shin and the shin's parent is
    # the thigh.
    #
    # Only ever FILLS IN what the names could not resolve; a confident name
    # match is never overridden. Twist bones are stepped over, since a rig may
    # put one between the thigh and the shin.
    def _fill_limbs_by_anatomy(self):
        for side in ("L", "R"):
            chain = []
            b = self.bones.get(self.limb.get(("foot", side), ""))
            while b is not None and len(chain) < 6:
                b = b.parent
                if b is None or b.name not in self.names:
                    break
                if is_cosmetic(b.name) or is_segment_of(b.name, self.names):
                    continue          # a twist segment is not a joint of its own
                chain.append(b.name)
            # chain is now [shin, thigh, hips, ...] going up from the foot.
            for role, up in (("shin", 0), ("thigh", 1)):
                if (role, side) in self.limb or up >= len(chain):
                    continue
                name = chain[up]
                # Never claim a bone another role already owns, and never claim
                # the trunk — a two-bone leg would otherwise take the hips.
                if name in self.limb.values():
                    continue
                if side_of(name) != side:
                    continue
                self.limb[(role, side)] = name

    # ---------------------------------------------------------------- spine

    def _resolve_spine(self):
        """Find the trunk as the LONGEST chain of non-limb, non-cosmetic bones.

        Deliberately not "lowest common ancestor of the two thighs", which is
        the obvious thing and is wrong on the rigs we actually get. A Rigify
        DEF-rig exports its chain ROOTS parented straight to the armature root
        (Rigify drives them by constraint, not by hierarchy), so on Taila the
        thighs, shoulders, skirt and hair all hang off `rootJoint` and the LCA
        is that root — which makes the hips the root bone and the spine one
        bone long. Taking the longest trunk chain instead reads the same rig
        correctly whether or not its hierarchy survived export.
        """
        blocked = set(self.limb.values())

        def candidate(name):
            if name in blocked or is_cosmetic(name):
                return False
            return not any(t in ("root", "master", "armature", "scene", "rootjoint")
                           for t in tokens(name))

        def walk(start):
            chain = [start]
            cur = self.bones[start]
            while True:
                nxt = [c for c in cur.children if candidate(c.name)]
                if not nxt:
                    break
                # If the walk forks, follow whichever branch reaches highest —
                # the torso continues upward, a stray nub does not.
                cur = max(nxt, key=self._branch_height)
                chain.append(cur.name)
            return chain

        best = []
        for name in self.names:
            if not candidate(name):
                continue
            parent = self.bones[name].parent
            if parent is not None and candidate(parent.name):
                continue  # not the base of a chain
            chain = walk(name)
            if len(chain) > len(best):
                best = chain
        if not best:
            return

        self.hips = best[0]
        self.spine = best[1:]
        if self.spine:
            self.head = self.spine[-1]
            named = [n for n in self.spine if "neck" in tokens(n)]
            self.neck = named[0] if named else (
                self.spine[-2] if len(self.spine) > 1 else None)

    def _branch_height(self, bone):
        """How high this branch reaches, in WORLD space.

        Measured on bone HEADS. glTF has no concept of a bone tail — joints are
        just nodes — so the tails Blender's importer shows are synthesised, and
        on Taila every skirt bone comes back with an identical 0.78 m tail.
        Heads are the only authored positions here.

        World space, not armature-local: the importer leaves the Y-up-to-Z-up
        correction on the armature OBJECT, so a local-space test picked the
        pelvis over the spine and the walk stopped one bone in.
        """
        mw = self.arm.matrix_world
        best = (mw @ bone.matrix_local.translation).z
        for c in bone.children:
            best = max(best, self._branch_height(c))
        return best

    # --------------------------------------------------------------- report

    def missing_core(self):
        need = [("thigh", "L"), ("thigh", "R"), ("shin", "L"), ("shin", "R"),
                ("foot", "L"), ("foot", "R"), ("upper_arm", "L"), ("upper_arm", "R"),
                ("forearm", "L"), ("forearm", "R"), ("hand", "L"), ("hand", "R")]
        miss = [f"{r}.{s}" for r, s in need if (r, s) not in self.limb]
        if not self.hips:
            miss.append("hips")
        if not self.head:
            miss.append("head")
        return miss

    def describe(self):
        lines = [f"  hips  {self.hips}", f"  spine {' -> '.join(self.spine)}"]
        for key in sorted(self.limb):
            lines.append(f"  {key[0]}.{key[1]:<10s} {self.limb[key]}")
        return "\n".join(lines)


def _match_chains(src_chain, tgt_chain):
    """Pair up two ordered chains of possibly different length.

    Each target bone takes the source bone nearest it in NORMALISED position,
    so a 4-bone torso can be driven by a 3-bone one. Because the retarget
    copies ABSOLUTE world orientation rather than composing local rotations,
    two target bones sharing one source bone simply end up parallel — the chain
    still finishes where the source says it does, it does not double the bend.
    """
    pairs = []
    if not src_chain or not tgt_chain:
        return pairs
    for i, tgt in enumerate(tgt_chain):
        t = (i + 0.5) / len(tgt_chain)
        j = min(range(len(src_chain)),
                key=lambda k: abs((k + 0.5) / len(src_chain) - t))
        pairs.append((src_chain[j], tgt))
    return pairs


def build_map(src: RigRoles, tgt: RigRoles, overrides: dict = None) -> dict:
    """target bone name -> source bone name, for the core body only."""
    mapping = {}

    if src.hips and tgt.hips:
        mapping[tgt.hips] = src.hips

    # Split both spines at the neck so a head never drives a chest, then match
    # torso-to-torso and neck-to-neck by position.
    def split(roles):
        chain = roles.spine
        if not chain:
            return [], [], None
        head = chain[-1]
        rest = chain[:-1]
        if roles.neck and roles.neck in rest:
            i = rest.index(roles.neck)
            return rest[:i], rest[i:], head
        return rest, [], head

    s_torso, s_neck, s_head = split(src)
    t_torso, t_neck, t_head = split(tgt)
    for a, b in _match_chains(s_torso, t_torso):
        mapping[b] = a
    for a, b in _match_chains(s_neck or s_torso[-1:], t_neck):
        mapping[b] = a
    if s_head and t_head:
        mapping[t_head] = s_head

    for key, tgt_name in tgt.limb.items():
        src_name = src.limb.get(key)
        if src_name:
            mapping[tgt_name] = src_name

    # Fingers and anything else that happens to share a spelling: map by
    # canonical stem + side. Cheap, and it makes a shared trigger-finger pose
    # come across when both rigs have fingers.
    src_by_stem = {}
    for n in src.names:
        src_by_stem.setdefault((canon(n), side_of(n)), n)
    for n in tgt.names:
        if n in mapping or is_cosmetic(n) or is_segment_of(n, tgt.names):
            continue
        hit = src_by_stem.get((canon(n), side_of(n)))
        if hit:
            mapping[n] = hit

    if overrides:
        for tgt_name, src_name in overrides.items():
            if src_name is None:
                mapping.pop(tgt_name, None)
            else:
                mapping[tgt_name] = src_name
    return mapping