1#====================== BEGIN GPL LICENSE BLOCK ======================
2#
3#  This program is free software; you can redistribute it and/or
4#  modify it under the terms of the GNU General Public License
5#  as published by the Free Software Foundation; either version 2
6#  of the License, or (at your option) any later version.
7#
8#  This program is distributed in the hope that it will be useful,
9#  but WITHOUT ANY WARRANTY; without even the implied warranty of
10#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11#  GNU General Public License for more details.
12#
13#  You should have received a copy of the GNU General Public License
14#  along with this program; if not, write to the Free Software Foundation,
15#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
16#
17#======================= END GPL LICENSE BLOCK ========================
18
19# <pep8 compliant>
20
21import bpy
22import re
23import time
24import traceback
25import sys
26from rna_prop_ui import rna_idprop_ui_prop_get
27
28from .utils import MetarigError, new_bone, get_rig_type
29from .utils import ORG_PREFIX, MCH_PREFIX, DEF_PREFIX, WGT_PREFIX, ROOT_NAME, make_original_name
30from .utils import RIG_DIR
31from .utils import create_root_widget, ensure_widget_collection
32from .utils import random_id
33from .utils import copy_attributes
34from .rig_ui_template import UI_SLIDERS, layers_ui, UI_REGISTER
35from .rig_ui_pitchipoy_template import UI_P_SLIDERS, layers_P_ui, UI_P_REGISTER
36
37
38RIG_MODULE = "rigs"
39ORG_LAYER = [n == 31 for n in range(0, 32)]  # Armature layer that original bones should be moved to.
40MCH_LAYER = [n == 30 for n in range(0, 32)]  # Armature layer that mechanism bones should be moved to.
41DEF_LAYER = [n == 29 for n in range(0, 32)]  # Armature layer that deformation bones should be moved to.
42ROOT_LAYER = [n == 28 for n in range(0, 32)]  # Armature layer that root bone should be moved to.
43
44
45class Timer:
46    def __init__(self):
47        self.timez = time.time()
48
49    def tick(self, string):
50        t = time.time()
51        print(string + "%.3f" % (t - self.timez))
52        self.timez = t
53
54
55# TODO: generalize to take a group as input instead of an armature.
56def generate_rig(context, metarig):
57    """ Generates a rig from a metarig.
58
59    """
60    t = Timer()
61
62    # Random string with time appended so that
63    # different rigs don't collide id's
64    rig_id = random_id(16)
65
66    # Initial configuration
67    # mode_orig = context.mode  # UNUSED
68    rest_backup = metarig.data.pose_position
69    metarig.data.pose_position = 'REST'
70
71    bpy.ops.object.mode_set(mode='OBJECT')
72
73    scene = context.scene
74    view_layer = context.view_layer
75    collection = context.collection
76    layer_collection = context.layer_collection
77
78    #------------------------------------------
79    # Create/find the rig object and set it up
80
81    # Check if the generated rig already exists, so we can
82    # regenerate in the same object.  If not, create a new
83    # object to generate the rig in.
84    print("Fetch rig.")
85    try:
86        name = metarig["rig_object_name"]
87    except KeyError:
88        name = "rig"
89
90    try:
91        obj = scene.objects[name]
92    except KeyError:
93        obj = bpy.data.objects.new(name, bpy.data.armatures.new(name))
94        obj.display_type = 'WIRE'
95        collection.objects.link(obj)
96
97    obj.data.pose_position = 'POSE'
98
99    # Get rid of anim data in case the rig already existed
100    print("Clear rig animation data.")
101    obj.animation_data_clear()
102    obj.data.animation_data_clear()
103
104    # Select generated rig object
105    metarig.select_set(False)
106    obj.select_set(True)
107    view_layer.objects.active = obj
108
109    # Remove all bones from the generated rig armature.
110    bpy.ops.object.mode_set(mode='EDIT')
111    for bone in obj.data.edit_bones:
112        obj.data.edit_bones.remove(bone)
113    bpy.ops.object.mode_set(mode='OBJECT')
114
115    # Create temporary duplicates for merging
116    temp_rig_1 = metarig.copy()
117    temp_rig_1.data = metarig.data.copy()
118    collection.objects.link(temp_rig_1)
119
120    temp_rig_2 = metarig.copy()
121    temp_rig_2.data = obj.data
122    collection.objects.link(temp_rig_2)
123
124    # Select the temp rigs for merging
125    for objt in view_layer.objects:
126        objt.select_set(False)  # deselect all objects
127    temp_rig_1.select_set(True)
128    temp_rig_2.select_set(True)
129    view_layer.objects.active = temp_rig_2
130
131    # Merge the temporary rigs
132    bpy.ops.object.join()
133
134    # Delete the second temp rig
135    bpy.ops.object.delete()
136
137    # Select the generated rig
138    for objt in view_layer.objects:
139        objt.select_set(False) # deselect all objects
140    obj.select_set(True)
141    view_layer.objects.active = obj
142
143    # Copy over bone properties
144    for bone in metarig.data.bones:
145        bone_gen = obj.data.bones[bone.name]
146
147        # B-bone stuff
148        bone_gen.bbone_segments = bone.bbone_segments
149        bone_gen.bbone_easein = bone.bbone_easein
150        bone_gen.bbone_easeout = bone.bbone_easeout
151
152    # Copy over the pose_bone properties
153    for bone in metarig.pose.bones:
154        bone_gen = obj.pose.bones[bone.name]
155
156        # Rotation mode and transform locks
157        bone_gen.rotation_mode = bone.rotation_mode
158        bone_gen.lock_rotation = tuple(bone.lock_rotation)
159        bone_gen.lock_rotation_w = bone.lock_rotation_w
160        bone_gen.lock_rotations_4d = bone.lock_rotations_4d
161        bone_gen.lock_location = tuple(bone.lock_location)
162        bone_gen.lock_scale = tuple(bone.lock_scale)
163
164        # rigify_type and rigify_parameters
165        bone_gen.rigify_type = bone.rigify_type
166        for prop in dir(bone_gen.rigify_parameters):
167            if (not prop.startswith("_")) \
168            and (not prop.startswith("bl_")) \
169            and (prop != "rna_type"):
170                try:
171                    setattr(bone_gen.rigify_parameters, prop, \
172                            getattr(bone.rigify_parameters, prop))
173                except AttributeError:
174                    print("FAILED TO COPY PARAMETER: " + str(prop))
175
176        # Custom properties
177        for prop in bone.keys():
178            try:
179                bone_gen[prop] = bone[prop]
180            except KeyError:
181                pass
182
183        # Constraints
184        for con1 in bone.constraints:
185            con2 = bone_gen.constraints.new(type=con1.type)
186            copy_attributes(con1, con2)
187
188            # Set metarig target to rig target
189            if "target" in dir(con2):
190                if con2.target == metarig:
191                    con2.target = obj
192
193    # Copy drivers
194    if metarig.animation_data:
195        for d1 in metarig.animation_data.drivers:
196            d2 = obj.driver_add(d1.data_path)
197            copy_attributes(d1, d2)
198            copy_attributes(d1.driver, d2.driver)
199
200            # Remove default modifiers, variables, etc.
201            for m in d2.modifiers:
202                d2.modifiers.remove(m)
203            for v in d2.driver.variables:
204                d2.driver.variables.remove(v)
205
206            # Copy modifiers
207            for m1 in d1.modifiers:
208                m2 = d2.modifiers.new(type=m1.type)
209                copy_attributes(m1, m2)
210
211            # Copy variables
212            for v1 in d1.driver.variables:
213                v2 = d2.driver.variables.new()
214                copy_attributes(v1, v2)
215                for i in range(len(v1.targets)):
216                    copy_attributes(v1.targets[i], v2.targets[i])
217                    # Switch metarig targets to rig targets
218                    if v2.targets[i].id == metarig:
219                        v2.targets[i].id = obj
220
221                    # Mark targets that may need to be altered after rig generation
222                    tar = v2.targets[i]
223                    # If a custom property
224                    if v2.type == 'SINGLE_PROP' \
225                    and re.match('^pose.bones\["[^"\]]*"\]\["[^"\]]*"\]$', tar.data_path):
226                        tar.data_path = "RIGIFY-" + tar.data_path
227
228            # Copy key frames
229            for i in range(len(d1.keyframe_points)):
230                d2.keyframe_points.add()
231                k1 = d1.keyframe_points[i]
232                k2 = d2.keyframe_points[i]
233                copy_attributes(k1, k2)
234
235    t.tick("Duplicate rig: ")
236    #----------------------------------
237    # Make a list of the original bones so we can keep track of them.
238    original_bones = [bone.name for bone in obj.data.bones]
239
240    # Add the ORG_PREFIX to the original bones.
241    bpy.ops.object.mode_set(mode='OBJECT')
242    for i in range(0, len(original_bones)):
243        obj.data.bones[original_bones[i]].name = make_original_name(original_bones[i])
244        original_bones[i] = make_original_name(original_bones[i])
245
246    # Create a sorted list of the original bones, sorted in the order we're
247    # going to traverse them for rigging.
248    # (root-most -> leaf-most, alphabetical)
249    bones_sorted = []
250    for name in original_bones:
251        bones_sorted += [name]
252    bones_sorted.sort()  # first sort by names
253    bones_sorted.sort(key=lambda bone: len(obj.pose.bones[bone].parent_recursive))  # then parents before children
254
255    t.tick("Make list of org bones: ")
256    #----------------------------------
257    # Create the root bone.
258    bpy.ops.object.mode_set(mode='EDIT')
259    root_bone = new_bone(obj, ROOT_NAME)
260    obj.data.edit_bones[root_bone].head = (0, 0, 0)
261    obj.data.edit_bones[root_bone].tail = (0, 1, 0)
262    obj.data.edit_bones[root_bone].roll = 0
263    bpy.ops.object.mode_set(mode='OBJECT')
264    obj.data.bones[root_bone].layers = ROOT_LAYER
265    # Put the rig_name in the armature custom properties
266    rna_idprop_ui_prop_get(obj.data, "rig_id", create=True)
267    obj.data["rig_id"] = rig_id
268
269    # Create/find widget collection
270    ensure_widget_collection(context)
271
272    t.tick("Create root bone: ")
273    #----------------------------------
274    try:
275        # Collect/initialize all the rigs.
276        rigs = []
277        for bone in bones_sorted:
278            bpy.ops.object.mode_set(mode='EDIT')
279            rigs += get_bone_rigs(obj, bone)
280        t.tick("Initialize rigs: ")
281
282        # Generate all the rigs.
283        ui_scripts = []
284        for rig in rigs:
285            # Go into editmode in the rig armature
286            bpy.ops.object.mode_set(mode='OBJECT')
287            context.view_layer.objects.active = obj
288            obj.select_set(True)
289            bpy.ops.object.mode_set(mode='EDIT')
290            scripts = rig.generate()
291            if scripts is not None:
292                ui_scripts += [scripts[0]]
293        t.tick("Generate rigs: ")
294    except Exception as e:
295        # Cleanup if something goes wrong
296        print("Rigify: failed to generate rig.")
297        metarig.data.pose_position = rest_backup
298        obj.data.pose_position = 'POSE'
299        bpy.ops.object.mode_set(mode='OBJECT')
300
301        # Continue the exception
302        raise e
303
304    #----------------------------------
305    bpy.ops.object.mode_set(mode='OBJECT')
306
307    # Get a list of all the bones in the armature
308    bones = [bone.name for bone in obj.data.bones]
309
310    # Parent any free-floating bones to the root.
311    bpy.ops.object.mode_set(mode='EDIT')
312    for bone in bones:
313        if obj.data.edit_bones[bone].parent is None:
314            obj.data.edit_bones[bone].use_connect = False
315            obj.data.edit_bones[bone].parent = obj.data.edit_bones[root_bone]
316    bpy.ops.object.mode_set(mode='OBJECT')
317
318    # Lock transforms on all non-control bones
319    r = re.compile("[A-Z][A-Z][A-Z]-")
320    for bone in bones:
321        if r.match(bone):
322            pb = obj.pose.bones[bone]
323            pb.lock_location = (True, True, True)
324            pb.lock_rotation = (True, True, True)
325            pb.lock_rotation_w = True
326            pb.lock_scale = (True, True, True)
327
328    # Every bone that has a name starting with "DEF-" make deforming.  All the
329    # others make non-deforming.
330    for bone in bones:
331        if obj.data.bones[bone].name.startswith(DEF_PREFIX):
332            obj.data.bones[bone].use_deform = True
333        else:
334            obj.data.bones[bone].use_deform = False
335
336    # Alter marked driver targets
337    if obj.animation_data:
338        for d in obj.animation_data.drivers:
339            for v in d.driver.variables:
340                for tar in v.targets:
341                    if tar.data_path.startswith("RIGIFY-"):
342                        temp, bone, prop = tuple([x.strip('"]') for x in tar.data_path.split('["')])
343                        if bone in obj.data.bones \
344                        and prop in obj.pose.bones[bone].keys():
345                            tar.data_path = tar.data_path[7:]
346                        else:
347                            tar.data_path = 'pose.bones["%s"]["%s"]' % (make_original_name(bone), prop)
348
349    # Move all the original bones to their layer.
350    for bone in original_bones:
351        obj.data.bones[bone].layers = ORG_LAYER
352
353    # Move all the bones with names starting with "MCH-" to their layer.
354    for bone in bones:
355        if obj.data.bones[bone].name.startswith(MCH_PREFIX):
356            obj.data.bones[bone].layers = MCH_LAYER
357
358    # Move all the bones with names starting with "DEF-" to their layer.
359    for bone in bones:
360        if obj.data.bones[bone].name.startswith(DEF_PREFIX):
361            obj.data.bones[bone].layers = DEF_LAYER
362
363    #  Create root bone widget
364    create_root_widget(obj, "root")
365
366    # Assign shapes to bones
367    # Object's with name WGT-<bone_name> get used as that bone's shape.
368    for bone in bones:
369        wgt_name = (WGT_PREFIX + obj.data.bones[bone].name)[:63]  # Object names are limited to 63 characters... arg
370        if wgt_name in context.scene.objects:
371            # Weird temp thing because it won't let me index by object name
372            for ob in context.scene.objects:
373                if ob.name == wgt_name:
374                    obj.pose.bones[bone].custom_shape = ob
375                    break
376            # This is what it should do:
377            # obj.pose.bones[bone].custom_shape = context.scene.objects[wgt_name]
378    # Reveal all the layers with control bones on them
379    vis_layers = [False for n in range(0, 32)]
380    for bone in bones:
381        for i in range(0, 32):
382            vis_layers[i] = vis_layers[i] or obj.data.bones[bone].layers[i]
383    for i in range(0, 32):
384        vis_layers[i] = vis_layers[i] and not (ORG_LAYER[i] or MCH_LAYER[i] or DEF_LAYER[i])
385    obj.data.layers = vis_layers
386
387    # Ensure the collection of layer names exists
388    for i in range(1 + len(metarig.data.rigify_layers), 29):
389        metarig.data.rigify_layers.add()
390
391    # Create list of layer name/row pairs
392    layer_layout = []
393    for l in metarig.data.rigify_layers:
394        print( l.name )
395        layer_layout += [(l.name, l.row)]
396
397
398    if isPitchipoy(metarig):
399
400        # Generate the UI Pitchipoy script
401        if "rig_ui.py" in bpy.data.texts:
402            script = bpy.data.texts["rig_ui.py"]
403            script.clear()
404        else:
405            script = bpy.data.texts.new("rig_ui.py")
406        script.write(UI_P_SLIDERS % rig_id)
407        for s in ui_scripts:
408            script.write("\n        " + s.replace("\n", "\n        ") + "\n")
409        script.write(layers_P_ui(vis_layers, layer_layout))
410        script.write(UI_P_REGISTER)
411        script.use_module = True
412
413    else:
414        # Generate the UI script
415        if "rig_ui.py" in bpy.data.texts:
416            script = bpy.data.texts["rig_ui.py"]
417            script.clear()
418        else:
419            script = bpy.data.texts.new("rig_ui.py")
420        script.write(UI_SLIDERS % rig_id)
421        for s in ui_scripts:
422            script.write("\n        " + s.replace("\n", "\n        ") + "\n")
423        script.write(layers_ui(vis_layers, layer_layout))
424        script.write(UI_REGISTER)
425        script.use_module = True
426
427    # Run UI script
428    exec(script.as_string(), {})
429
430    t.tick("The rest: ")
431    #----------------------------------
432    # Deconfigure
433    bpy.ops.object.mode_set(mode='OBJECT')
434    metarig.data.pose_position = rest_backup
435    obj.data.pose_position = 'POSE'
436
437    #----------------------------------
438    # Restore active collection
439    view_layer.active_layer_collection = layer_collection
440
441
442def get_bone_rigs(obj, bone_name, halt_on_missing=False):
443    """ Fetch all the rigs specified on a bone.
444    """
445    rigs = []
446    rig_type = obj.pose.bones[bone_name].rigify_type
447    rig_type = rig_type.replace(" ", "")
448
449    if rig_type == "":
450        pass
451    else:
452        # Gather parameters
453        params = obj.pose.bones[bone_name].rigify_parameters
454
455        # Get the rig
456        try:
457            rig = get_rig_type(rig_type).Rig(obj, bone_name, params)
458        except ImportError:
459            message = "Rig Type Missing: python module for type '%s' not found (bone: %s)" % (rig_type, bone_name)
460            if halt_on_missing:
461                raise MetarigError(message)
462            else:
463                print(message)
464                print('print_exc():')
465                traceback.print_exc(file=sys.stdout)
466        else:
467            rigs += [rig]
468    return rigs
469
470
471def param_matches_type(param_name, rig_type):
472    """ Returns True if the parameter name is consistent with the rig type.
473    """
474    if param_name.rsplit(".", 1)[0] == rig_type:
475        return True
476    else:
477        return False
478
479
480def param_name(param_name, rig_type):
481    """ Get the actual parameter name, sans-rig-type.
482    """
483    return param_name[len(rig_type) + 1:]
484
485def isPitchipoy(metarig):
486    """ Returns True if metarig is type pitchipoy.
487    """
488    pbones=metarig.pose.bones
489    for pb in pbones:
490        words = pb.rigify_type.partition('.')
491        if  words[0] == 'pitchipoy':
492            return True
493    return False
494