__init__.py

#script to make rigify compatible with unity humanoid
#HOWTO: right after generating rig using rigify
#	press armature -> Rigify To Unity Converter -> (Prepare rig for unity) button
bl_info = {
    "name": "Rigify to Unity",
    "category": "Rigging",
    "description": "Change Rigify rig into Mecanim-ready rig for Unity",
    "location": "At the bottom of Rigify rig data/armature tab",
    "blender":(2,80,0)
}

import bpy
import re


BASE_REQUIRED_BONES = (
    'DEF-shoulder.L',
    'DEF-shoulder.R',
    'DEF-spine.003',
    'DEF-upper_arm.L',
    'DEF-upper_arm.L.001',
    'DEF-upper_arm.R',
    'DEF-upper_arm.R.001',
    'DEF-forearm.L',
    'DEF-forearm.L.001',
    'DEF-forearm.R',
    'DEF-forearm.R.001',
    'DEF-hand.L',
    'DEF-hand.R',
    'DEF-spine',
    'DEF-thigh.L',
    'DEF-thigh.L.001',
    'DEF-thigh.R',
    'DEF-thigh.R.001',
    'DEF-shin.L',
    'DEF-shin.L.001',
    'DEF-shin.R',
    'DEF-shin.R.001',
    'DEF-foot.L',
    'DEF-foot.R',
)

TAIL_BONE_PATTERN = re.compile(r'^DEF-tail(?:\.(\d+))?$')
EAR_BONE_PATTERN = re.compile(r'^DEF-ear\.(L|R)(?:\.(\d+))?$')
RITTY_BODY_BONE_PARENTS = (
    ('DEF-belly.C', 'DEF-spine.002'),
    ('DEF-Breast.C', 'DEF-spine.003'),
)


def find_numbered_chain(bones, pattern, suffix_group=1):
    """Return matching bone names with the unnumbered bone first."""
    matches = []

    for bone in bones:
        match = pattern.fullmatch(bone.name)
        if match is None:
            continue

        suffix = match.group(suffix_group)
        sort_key = (suffix is not None, int(suffix) if suffix is not None else -1)
        matches.append((sort_key, bone.name))

    matches.sort(key=lambda item: item[0])
    return [name for _, name in matches]


def find_ear_chains(bones):
    chains = {'L': [], 'R': []}

    for bone in bones:
        match = EAR_BONE_PATTERN.fullmatch(bone.name)
        if match is None:
            continue

        side = match.group(1)
        suffix = match.group(2)
        sort_key = (suffix is not None, int(suffix) if suffix is not None else -1)
        chains[side].append((sort_key, bone.name))

    for side in chains:
        chains[side].sort(key=lambda item: item[0])
        chains[side] = [name for _, name in chains[side]]

    return chains


def reparent_def_chain(edit_bones, chain_names, root_parent):
    if not chain_names:
        return

    parent = root_parent
    for name in chain_names:
        bone = edit_bones.get(name)
        if bone is None:
            continue

        # Parenting is all Unity needs. Keeping the bones disconnected prevents
        # Blender from snapping their heads to a new parent's tail.
        bone.use_connect = False
        bone.parent = parent
        parent = bone


def get_org_bone_name(def_bone_name):
    if def_bone_name.startswith('DEF-'):
        return 'ORG-' + def_bone_name[4:]
    return 'ORG-' + def_bone_name


def ensure_copy_transforms_constraint(ob, bone_name, target_bone_name):
    pose_bone = ob.pose.bones.get(bone_name)
    if pose_bone is None:
        return False

    for constraint in pose_bone.constraints:
        if (
            constraint.type == 'COPY_TRANSFORMS'
            and constraint.target == ob
            and constraint.subtarget == target_bone_name
        ):
            return False

    constraint = pose_bone.constraints.new('COPY_TRANSFORMS')
    constraint.name = 'Ritty Copy Transforms'
    constraint.target = ob
    constraint.subtarget = target_bone_name
    return True


def find_missing_base_bones(ob):
    return [name for name in BASE_REQUIRED_BONES if name not in ob.data.bones]


def validate_ritty_anchors(ob):
    tail_names = find_numbered_chain(ob.data.bones, TAIL_BONE_PATTERN)
    ear_chains = find_ear_chains(ob.data.bones)

    if tail_names and 'root' not in ob.data.bones:
        return "Tail bones were found, but the required 'root' bone is missing"

    has_ears = any(ear_chains.values())
    has_head = 'DEF-head' in ob.data.bones or 'DEF-spine.006' in ob.data.bones
    if has_ears and not has_head:
        return "Ear bones were found, but neither 'DEF-head' nor 'DEF-spine.006' exists"

    for bone_name, parent_name in RITTY_BODY_BONE_PARENTS:
        if bone_name in ob.data.bones and parent_name not in ob.data.bones:
            return "Bone '{}' requires parent bone '{}'".format(bone_name, parent_name)

        org_name = get_org_bone_name(bone_name)
        if bone_name in ob.data.bones and org_name not in ob.data.bones:
            return "Bone '{}' requires constraint target bone '{}'".format(bone_name, org_name)

    return None


def convert_base_rig(ob):
    """Run the add-on's original Humanoid conversion."""
    bpy.ops.object.mode_set(mode='OBJECT')

    if 'DEF-breast.L' in ob.data.bones:
        ob.data.bones['DEF-breast.L'].use_deform = False
    if 'DEF-breast.R' in ob.data.bones:
        ob.data.bones['DEF-breast.R'].use_deform = False

    if 'DEF-pelvis.L' in ob.data.bones:
        ob.data.bones['DEF-pelvis.L'].use_deform = False
    if 'DEF-pelvis.R' in ob.data.bones:
        ob.data.bones['DEF-pelvis.R'].use_deform = False

    bpy.ops.object.mode_set(mode='EDIT')

    ob.data.edit_bones['DEF-shoulder.L'].parent = ob.data.edit_bones['DEF-spine.003']
    ob.data.edit_bones['DEF-shoulder.R'].parent = ob.data.edit_bones['DEF-spine.003']

    ob.data.edit_bones['DEF-upper_arm.L'].parent = ob.data.edit_bones['DEF-shoulder.L']
    ob.data.edit_bones['DEF-upper_arm.R'].parent = ob.data.edit_bones['DEF-shoulder.R']

    ob.data.edit_bones['DEF-thigh.L'].parent = ob.data.edit_bones['DEF-spine']
    ob.data.edit_bones['DEF-thigh.R'].parent = ob.data.edit_bones['DEF-spine']

    ob.data.edit_bones['DEF-upper_arm.L'].tail = ob.data.edit_bones['DEF-upper_arm.L.001'].tail
    ob.data.edit_bones['DEF-forearm.L'].tail = ob.data.edit_bones['DEF-forearm.L.001'].tail
    ob.data.edit_bones['DEF-forearm.L'].parent = ob.data.edit_bones['DEF-upper_arm.L.001'].parent
    ob.data.edit_bones['DEF-hand.L'].parent = ob.data.edit_bones['DEF-forearm.L.001'].parent
    ob.data.edit_bones.remove(ob.data.edit_bones['DEF-upper_arm.L.001'])
    ob.data.edit_bones.remove(ob.data.edit_bones['DEF-forearm.L.001'])

    ob.data.edit_bones['DEF-upper_arm.R'].tail = ob.data.edit_bones['DEF-upper_arm.R.001'].tail
    ob.data.edit_bones['DEF-forearm.R'].tail = ob.data.edit_bones['DEF-forearm.R.001'].tail
    ob.data.edit_bones['DEF-forearm.R'].parent = ob.data.edit_bones['DEF-upper_arm.R.001'].parent
    ob.data.edit_bones['DEF-hand.R'].parent = ob.data.edit_bones['DEF-forearm.R.001'].parent
    ob.data.edit_bones.remove(ob.data.edit_bones['DEF-upper_arm.R.001'])
    ob.data.edit_bones.remove(ob.data.edit_bones['DEF-forearm.R.001'])

    ob.data.edit_bones['DEF-thigh.L'].tail = ob.data.edit_bones['DEF-thigh.L.001'].tail
    ob.data.edit_bones['DEF-shin.L'].tail = ob.data.edit_bones['DEF-shin.L.001'].tail
    ob.data.edit_bones['DEF-shin.L'].parent = ob.data.edit_bones['DEF-thigh.L.001'].parent
    ob.data.edit_bones['DEF-foot.L'].parent = ob.data.edit_bones['DEF-shin.L.001'].parent
    ob.data.edit_bones.remove(ob.data.edit_bones['DEF-thigh.L.001'])
    ob.data.edit_bones.remove(ob.data.edit_bones['DEF-shin.L.001'])

    ob.data.edit_bones['DEF-thigh.R'].tail = ob.data.edit_bones['DEF-thigh.R.001'].tail
    ob.data.edit_bones['DEF-shin.R'].tail = ob.data.edit_bones['DEF-shin.R.001'].tail
    ob.data.edit_bones['DEF-shin.R'].parent = ob.data.edit_bones['DEF-thigh.R.001'].parent
    ob.data.edit_bones['DEF-foot.R'].parent = ob.data.edit_bones['DEF-shin.R.001'].parent
    ob.data.edit_bones.remove(ob.data.edit_bones['DEF-thigh.R.001'])
    ob.data.edit_bones.remove(ob.data.edit_bones['DEF-shin.R.001'])

    if 'DEF-pelvis.L' in ob.data.edit_bones:
        ob.data.edit_bones.remove(ob.data.edit_bones['DEF-pelvis.L'])
    if 'DEF-pelvis.R' in ob.data.edit_bones:
        ob.data.edit_bones.remove(ob.data.edit_bones['DEF-pelvis.R'])

    if 'DEF-breast.L' in ob.data.edit_bones:
        ob.data.edit_bones.remove(ob.data.edit_bones['DEF-breast.L'])
    if 'DEF-breast.R' in ob.data.edit_bones:
        ob.data.edit_bones.remove(ob.data.edit_bones['DEF-breast.R'])

    bpy.ops.object.mode_set(mode='OBJECT')

    namelist = [('DEF-spine.006', 'DEF-head'), ('DEF-spine.005', 'DEF-neck')]

    for name, newname in namelist:
        pb = ob.pose.bones.get(name)
        if pb is None:
            continue
        pb.name = newname


def convert_ritty_extras(ob):
    """Build Ritty's custom export hierarchy and remove its central pelvis bone."""
    bpy.ops.object.mode_set(mode='EDIT')
    try:
        edit_bones = ob.data.edit_bones

        tail_names = find_numbered_chain(edit_bones, TAIL_BONE_PATTERN)
        ear_chains = find_ear_chains(edit_bones)
        body_bone_names = []
        pelvis_removed = False

        if tail_names:
            reparent_def_chain(edit_bones, tail_names, edit_bones.get('root'))

        head = edit_bones.get('DEF-head') or edit_bones.get('DEF-spine.006')
        if head is not None:
            for chain_names in ear_chains.values():
                reparent_def_chain(edit_bones, chain_names, head)

        for bone_name, parent_name in RITTY_BODY_BONE_PARENTS:
            bone = edit_bones.get(bone_name)
            if bone is not None:
                bone.use_connect = False
                bone.parent = edit_bones.get(parent_name)
                body_bone_names.append(bone_name)

        pelvis = edit_bones.get('DEF-pelvis.C')
        pelvis_found = pelvis is not None
        if pelvis_found:
            spine = edit_bones.get('DEF-spine')
            for child in list(pelvis.children):
                child.use_connect = False
                child.parent = spine

            edit_bones.remove(pelvis)
            pelvis_removed = True
    finally:
        bpy.ops.object.mode_set(mode='OBJECT')

    body_constraint_count = 0
    for bone_name in body_bone_names:
        if ensure_copy_transforms_constraint(ob, bone_name, get_org_bone_name(bone_name)):
            body_constraint_count += 1

    return {
        'tail_count': len(tail_names),
        'ear_count': sum(len(chain) for chain in ear_chains.values()),
        'body_bone_count': len(body_bone_names),
        'body_constraint_count': body_constraint_count,
        'pelvis_found': pelvis_found,
        'pelvis_removed': pelvis_removed,
    }


class UnityMecanim_Panel(bpy.types.Panel):
    bl_label = "Rigify to Unity converter"
    bl_space_type = "PROPERTIES"
    bl_region_type = "WINDOW"
    bl_context = "data"

    @classmethod
    def poll(self, context):
        return context.object is not None and context.object.type == 'ARMATURE'
    
    def draw(self, context):
        self.layout.operator("rig4mec.convert2unity", text="Prepare rig for Unity")
        self.layout.operator("rig4mec.convert_ritty2unity", text="Prepare Ritty rig for Unity")
        
        
class UnityMecanim_Convert2Unity(bpy.types.Operator):
    bl_idname = "rig4mec.convert2unity"
    bl_label = "Prepare rig for unity"
    
    def execute(self, context):
        ob = context.object
        missing = find_missing_base_bones(ob)
        if missing:
            self.report({'ERROR'}, 'Missing required bones: ' + ', '.join(missing))
            return {'CANCELLED'}

        convert_base_rig(ob)
        self.report({'INFO'}, 'Unity ready rig!')
        return {'FINISHED'}


class UnityMecanim_ConvertRitty2Unity(bpy.types.Operator):
    bl_idname = "rig4mec.convert_ritty2unity"
    bl_label = "Prepare Ritty rig for Unity"
    bl_description = "Run the standard conversion and add Ritty custom bone handling"

    def execute(self, context):
        ob = context.object
        missing = find_missing_base_bones(ob)
        if missing:
            self.report(
                {'ERROR'},
                'Use an unconverted generated rig. Missing required bones: ' + ', '.join(missing),
            )
            return {'CANCELLED'}

        anchor_error = validate_ritty_anchors(ob)
        if anchor_error is not None:
            self.report({'ERROR'}, anchor_error)
            return {'CANCELLED'}

        convert_base_rig(ob)
        result = convert_ritty_extras(ob)

        pelvis_status = 'not found'
        if result['pelvis_removed']:
            pelvis_status = 'removed'

        self.report(
            {'INFO'},
            'Ritty rig ready! Tail: {tail}, ears: {ears}, body: {body}, constraints: {constraints}, DEF-pelvis.C: {pelvis}'.format(
                tail=result['tail_count'],
                ears=result['ear_count'],
                body=result['body_bone_count'],
                constraints=result['body_constraint_count'],
                pelvis=pelvis_status,
            ),
        )
        return {'FINISHED'}


def register():
    bpy.utils.register_class(UnityMecanim_Convert2Unity)
    bpy.utils.register_class(UnityMecanim_ConvertRitty2Unity)
    bpy.utils.register_class(UnityMecanim_Panel)
    
    
def unregister():
    bpy.utils.unregister_class(UnityMecanim_Panel)
    bpy.utils.unregister_class(UnityMecanim_ConvertRitty2Unity)
    bpy.utils.unregister_class(UnityMecanim_Convert2Unity)