1import bpy
2from ...utils    import copy_bone
3from ...utils    import strip_org, make_deformer_name, connected_children_names
4from ...utils    import put_bone, create_sphere_widget
5from ...utils    import create_circle_widget
6from ...utils    import MetarigError
7
8class Rig:
9
10    def __init__(self, obj, bone_name, params):
11        self.obj = obj
12        self.org_bones = [bone_name] + connected_children_names(obj, bone_name)
13        self.params = params
14
15        self.copy_rotaion_axes = params.copy_rotaion_axes
16
17        if params.tweak_extra_layers:
18            self.tweak_layers = list( params.tweak_layers )
19        else:
20            self.tweak_layers = None
21
22        if len(self.org_bones) <= 1:
23            raise MetarigError(
24                "RIGIFY ERROR: invalid rig structure" % (strip_org(bone_name))
25            )
26
27
28    def make_controls( self ):
29
30        bpy.ops.object.mode_set(mode ='EDIT')
31        org_bones = self.org_bones
32
33        ctrl_chain = []
34        for i in range( len( org_bones ) ):
35            name = org_bones[i]
36
37            ctrl_bone  = copy_bone(
38                self.obj,
39                name,
40                strip_org(name)
41            )
42
43            ctrl_chain.append( ctrl_bone )
44
45        # Make widgets
46        bpy.ops.object.mode_set(mode ='OBJECT')
47
48        for ctrl in ctrl_chain:
49            create_circle_widget(self.obj, ctrl, radius=0.3, head_tail=0.5)
50
51        return ctrl_chain
52
53
54    def make_tweaks( self ):
55
56        bpy.ops.object.mode_set(mode ='EDIT')
57        eb = self.obj.data.edit_bones
58        org_bones = self.org_bones
59
60        tweak_chain = []
61        for i in range( len( org_bones ) + 1 ):
62            if i == len( org_bones ):
63                # Make final tweak at the tip of the tentacle
64                name = org_bones[i-1]
65            else:
66                name = org_bones[i]
67
68            tweak_bone = copy_bone(
69                self.obj,
70                name,
71                "tweak_" + strip_org(name)
72            )
73
74            tweak_e = eb[ tweak_bone ]
75
76            tweak_e.length /= 2 # Set size to half
77
78            if i == len( org_bones ):
79                # Position final tweak at the tip
80                put_bone( self.obj, tweak_bone, eb[ org_bones[-1]].tail )
81
82            tweak_chain.append( tweak_bone )
83
84        # Make widgets
85        bpy.ops.object.mode_set(mode = 'OBJECT')
86
87        for tweak in tweak_chain:
88            create_sphere_widget( self.obj, tweak )
89
90            tweak_pb = self.obj.pose.bones[ tweak ]
91
92            # Set locks
93            if tweak_chain.index( tweak ) != len( tweak_chain ) - 1:
94                tweak_pb.lock_rotation = (True, False, True)
95                tweak_pb.lock_scale    = (False, True, False)
96            else:
97                tweak_pb.lock_rotation_w = True
98                tweak_pb.lock_rotation   = (True, True, True)
99                tweak_pb.lock_scale      = (True, True, True)
100
101            # Set up tweak bone layers
102            if self.tweak_layers:
103                tweak_pb.bone.layers = self.tweak_layers
104
105        return tweak_chain
106
107
108    def make_deform( self ):
109
110        bpy.ops.object.mode_set(mode ='EDIT')
111        org_bones = self.org_bones
112
113        def_chain = []
114        for i in range( len( org_bones ) ):
115            name = org_bones[i]
116
117            def_bone  = copy_bone(
118                self.obj,
119                name,
120                make_deformer_name(strip_org(name))
121            )
122
123            def_chain.append( def_bone )
124
125        return def_chain
126
127
128    def parent_bones( self, all_bones ):
129
130        bpy.ops.object.mode_set(mode ='EDIT')
131        org_bones = self.org_bones
132        eb        = self.obj.data.edit_bones
133
134        # Parent control bones
135        for bone in all_bones['control'][1:]:
136            previous_index    = all_bones['control'].index( bone ) - 1
137            eb[ bone ].parent = eb[ all_bones['control'][previous_index] ]
138
139        # Parent tweak bones
140        tweaks = all_bones['tweak']
141        for tweak in all_bones['tweak']:
142            parent = ''
143            if tweaks.index( tweak ) == len( tweaks ) - 1:
144                parent = all_bones['control'][ -1 ]
145            else:
146                parent = all_bones['control'][ tweaks.index( tweak ) ]
147
148            eb[ tweak ].parent = eb[ parent ]
149
150        # Parent deform bones
151        for bone in all_bones['deform'][1:]:
152            previous_index = all_bones['deform'].index( bone ) - 1
153
154            eb[ bone ].parent = eb[ all_bones['deform'][previous_index] ]
155            eb[ bone ].use_connect = True
156
157        # Parent org bones ( to tweaks by default, or to the controls )
158        for org, tweak in zip( org_bones, all_bones['tweak'] ):
159            eb[ org ].parent = eb[ tweak ]
160
161
162    def make_constraints( self, all_bones ):
163
164        bpy.ops.object.mode_set(mode ='OBJECT')
165        org_bones = self.org_bones
166        pb        = self.obj.pose.bones
167
168        # Deform bones' constraints
169        ctrls   = all_bones['control']
170        tweaks  = all_bones['tweak'  ]
171        deforms = all_bones['deform' ]
172
173        for deform, tweak, ctrl in zip( deforms, tweaks, ctrls ):
174            con           = pb[deform].constraints.new('COPY_TRANSFORMS')
175            con.target    = self.obj
176            con.subtarget = tweak
177
178            con           = pb[deform].constraints.new('DAMPED_TRACK')
179            con.target    = self.obj
180            con.subtarget = tweaks[ tweaks.index( tweak ) + 1 ]
181
182            con           = pb[deform].constraints.new('STRETCH_TO')
183            con.target    = self.obj
184            con.subtarget = tweaks[ tweaks.index( tweak ) + 1 ]
185
186            # Control bones' constraints
187            if ctrl != ctrls[0]:
188                con = pb[ctrl].constraints.new('COPY_ROTATION')
189                con.target       = self.obj
190                con.subtarget    = ctrls[ ctrls.index(ctrl) - 1 ]
191                for i, prop in enumerate( [ 'use_x', 'use_y', 'use_z' ] ):
192                    if self.copy_rotaion_axes[i]:
193                        setattr( con, prop, True )
194                    else:
195                        setattr( con, prop, False )
196                con.mix_mode     = 'OFFSET'
197                con.target_space = 'LOCAL'
198                con.owner_space  = 'LOCAL'
199
200
201
202    def generate(self):
203        bpy.ops.object.mode_set(mode ='EDIT')
204        eb = self.obj.data.edit_bones
205
206        # Clear all initial parenting
207        for bone in self.org_bones:
208        #    eb[ bone ].parent      = None
209            eb[ bone ].use_connect = False
210
211        # Creating all bones
212        ctrl_chain  = self.make_controls()
213        tweak_chain = self.make_tweaks()
214        def_chain   = self.make_deform()
215
216        all_bones = {
217            'control' : ctrl_chain,
218            'tweak'   : tweak_chain,
219            'deform'  : def_chain
220        }
221
222        self.make_constraints( all_bones )
223        self.parent_bones( all_bones )
224
225
226def add_parameters(params):
227    """ Add the parameters of this rig type to the
228        RigifyParameters PropertyGroup
229    """
230    params.copy_rotaion_axes = bpy.props.BoolVectorProperty(
231        size        = 3,
232        description = "Layers for the tweak controls to be on",
233        default     = tuple( [ i == 0 for i in range(0, 3) ] )
234        )
235
236    # Setting up extra tweak layers
237    params.tweak_extra_layers = bpy.props.BoolProperty(
238        name        = "tweak_extra_layers",
239        default     = True,
240        description = ""
241        )
242
243    params.tweak_layers = bpy.props.BoolVectorProperty(
244        size        = 32,
245        description = "Layers for the tweak controls to be on",
246        default     = tuple( [ i == 1 for i in range(0, 32) ] )
247        )
248
249
250def parameters_ui(layout, params):
251    """ Create the ui for the rig parameters.
252    """
253
254    r = layout.row()
255    col = r.column(align=True)
256    row = col.row(align=True)
257    for i,axis in enumerate( [ 'x', 'y', 'z' ] ):
258        row.prop(params, "copy_rotaion_axes", index=i, toggle=True, text=axis)
259
260    r = layout.row()
261    r.prop(params, "tweak_extra_layers")
262    r.active = params.tweak_extra_layers
263
264    col = r.column(align=True)
265    row = col.row(align=True)
266
267    for i in range( 8 ): # Layers 0-7
268        row.prop(params, "tweak_layers", index=i, toggle=True, text="")
269
270    row = col.row(align=True)
271
272    for i in range( 16, 24 ): # Layers 16-23
273        row.prop(params, "tweak_layers", index=i, toggle=True, text="")
274
275    col = r.column(align=True)
276    row = col.row(align=True)
277
278    for i in range( 8, 16 ): # Layers 8-15
279        row.prop(params, "tweak_layers", index=i, toggle=True, text="")
280
281    row = col.row(align=True)
282
283    for i in range( 24, 32 ): # Layers 24-31
284        row.prop(params, "tweak_layers", index=i, toggle=True, text="")
285
286def create_sample(obj):
287    # generated by rigify.utils.write_metarig
288    bpy.ops.object.mode_set(mode='EDIT')
289    arm = obj.data
290
291    bones = {}
292
293    bone = arm.edit_bones.new('Bone')
294    bone.head[:] = 0.0000, 0.0000, 0.0000
295    bone.tail[:] = 0.0000, 0.0000, 0.3333
296    bone.roll = 0.0000
297    bone.use_connect = False
298    bones['Bone'] = bone.name
299
300    bone = arm.edit_bones.new('Bone.002')
301    bone.head[:] = 0.0000, 0.0000, 0.3333
302    bone.tail[:] = 0.0000, 0.0000, 0.6667
303    bone.roll = 0.0000
304    bone.use_connect = True
305    bone.parent = arm.edit_bones[bones['Bone']]
306    bones['Bone.002'] = bone.name
307
308    bone = arm.edit_bones.new('Bone.001')
309    bone.head[:] = 0.0000, 0.0000, 0.6667
310    bone.tail[:] = 0.0000, 0.0000, 1.0000
311    bone.roll = 0.0000
312    bone.use_connect = True
313    bone.parent = arm.edit_bones[bones['Bone.002']]
314    bones['Bone.001'] = bone.name
315
316    bpy.ops.object.mode_set(mode='OBJECT')
317    pbone = obj.pose.bones[bones['Bone']]
318    pbone.rigify_type = 'pitchipoy.simple_tentacle'
319    pbone.lock_location = (False, False, False)
320    pbone.lock_rotation = (False, False, False)
321    pbone.lock_rotation_w = False
322    pbone.lock_scale = (False, False, False)
323    pbone.rotation_mode = 'QUATERNION'
324    pbone = obj.pose.bones[bones['Bone.002']]
325    pbone.rigify_type = ''
326    pbone.lock_location = (False, False, False)
327    pbone.lock_rotation = (False, False, False)
328    pbone.lock_rotation_w = False
329    pbone.lock_scale = (False, False, False)
330    pbone.rotation_mode = 'QUATERNION'
331    pbone = obj.pose.bones[bones['Bone.001']]
332    pbone.rigify_type = ''
333    pbone.lock_location = (False, False, False)
334    pbone.lock_rotation = (False, False, False)
335    pbone.lock_rotation_w = False
336    pbone.lock_scale = (False, False, False)
337    pbone.rotation_mode = 'QUATERNION'
338
339    bpy.ops.object.mode_set(mode='EDIT')
340    for bone in arm.edit_bones:
341        bone.select = False
342        bone.select_head = False
343        bone.select_tail = False
344    for b in bones:
345        bone = arm.edit_bones[bones[b]]
346        bone.select = True
347        bone.select_head = True
348        bone.select_tail = True
349        arm.edit_bones.active = bone
350