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
22
23from itertools import count
24
25from ...utils.naming import make_derived_name
26from ...utils.bones import align_bone_orientation
27from ...utils.widgets_basic import create_circle_widget
28from ...utils.widgets_special import create_neck_bend_widget, create_neck_tweak_widget
29from ...utils.switch_parent import SwitchParentBuilder
30from ...utils.misc import map_list
31
32from ...base_rig import stage
33
34from .spine_rigs import BaseHeadTailRig
35
36
37class Rig(BaseHeadTailRig):
38    """
39    Head rig with long neck support and connect option.
40    """
41
42    use_connect_reverse = False
43    min_chain_length = 1
44
45    def initialize(self):
46        super().initialize()
47
48        self.long_neck = len(self.bones.org) > 3
49        self.has_neck = len(self.bones.org) > 1
50
51    ####################################################
52    # BONES
53    #
54    # org[]:
55    #   ORG bones
56    # ctrl:
57    #   neck, head, neck_bend:
58    #     Main controls.
59    #   tweak[]:
60    #     Tweak control chain.
61    # mch:
62    #   rot_neck, rot_head:
63    #     Main control parents, implement FK follow.
64    #   stretch
65    #     Long neck stretch behavior.
66    #   ik[]
67    #     Long neck IK behavior.
68    #   chain[]
69    #     Tweak parents.
70    # deform[]:
71    #   DEF bones
72    #
73    ####################################################
74
75    ####################################################
76    # Main control bones
77
78    @stage.generate_bones
79    def make_control_chain(self):
80        orgs = self.bones.org
81        ctrl = self.bones.ctrl
82
83        if self.has_neck:
84            ctrl.neck = self.make_neck_control_bone(orgs[0], 'neck', orgs[-1])
85
86        ctrl.head = self.make_head_control_bone(orgs[-1], 'head')
87
88        if self.long_neck:
89            ctrl.neck_bend = self.make_neck_bend_control_bone(orgs[0], 'neck_bend', ctrl.neck)
90
91        self.default_prop_bone = ctrl.head
92
93    def make_neck_control_bone(self, org, name, org_head):
94        name = self.copy_bone(org, name, parent=False)
95
96        # Neck spans all neck bones (except head)
97        self.get_bone(name).tail = self.get_bone(org_head).head
98
99        return name
100
101    def make_neck_bend_control_bone(self, org, name, neck):
102        name = self.copy_bone(org, name, parent=False)
103        neck_bend_eb = self.get_bone(name)
104
105        # Neck pivot position
106        neck_bones = self.bones.org
107        if (len(neck_bones)-1) % 2:     # odd num of neck bones (head excluded)
108            center_bone = self.get_bone(neck_bones[int((len(neck_bones))/2) - 1])
109            neck_bend_eb.head = (center_bone.head + center_bone.tail)/2
110        else:
111            center_bone = self.get_bone(neck_bones[int((len(neck_bones)-1)/2) - 1])
112            neck_bend_eb.head = center_bone.tail
113
114        align_bone_orientation(self.obj, name, neck)
115        neck_bend_eb.length = self.get_bone(neck).length / 2
116
117        return name
118
119    def make_head_control_bone(self, org, name):
120        return self.copy_bone(org, name, parent=False)
121
122    @stage.parent_bones
123    def parent_control_chain(self):
124        ctrl = self.bones.ctrl
125        mch = self.bones.mch
126        if self.has_neck:
127            self.set_bone_parent(ctrl.neck, mch.rot_neck)
128        self.set_bone_parent(ctrl.head, mch.rot_head)
129        if self.long_neck:
130            self.set_bone_parent(ctrl.neck_bend, mch.stretch)
131
132    @stage.configure_bones
133    def configure_control_chain(self):
134        if self.has_neck:
135            self.configure_control_bone(0, self.bones.ctrl.neck, self.bones.org[0])
136        self.configure_control_bone(2, self.bones.ctrl.head, self.bones.org[-1])
137        if self.long_neck:
138            self.configure_control_bone(1, self.bones.ctrl.neck_bend, self.bones.org[0])
139
140    @stage.generate_widgets
141    def make_control_widgets(self):
142        ctrl = self.bones.ctrl
143        if self.has_neck:
144            self.make_neck_widget(ctrl.neck)
145        self.make_head_widget(ctrl.head)
146        if self.long_neck:
147            self.make_neck_bend_widget(ctrl.neck_bend)
148
149    def make_neck_widget(self, ctrl):
150        radius = 1/max(1, len(self.bones.mch.chain))
151
152        create_circle_widget(
153            self.obj, ctrl,
154            radius=radius,
155            head_tail=0.5,
156        )
157
158    def make_neck_bend_widget(self, ctrl):
159        radius = 1/max(1, len(self.bones.mch.chain))
160
161        create_neck_bend_widget(
162            self.obj, ctrl,
163            radius=radius/2,
164            head_tail=0.0,
165        )
166
167    def make_head_widget(self, ctrl):
168        # place wgt @ middle of head bone for long necks
169        if self.long_neck:
170            head_tail = 0.5
171        else:
172            head_tail = 1.0
173
174        create_circle_widget(
175            self.obj, ctrl,
176            radius              = 0.5,
177            head_tail           = head_tail,
178            with_line           = False,
179        )
180
181    ####################################################
182    # MCH bones associated with main controls
183
184    @stage.generate_bones
185    def make_mch_control_bones(self):
186        orgs = self.bones.org
187        mch = self.bones.mch
188
189        if self.has_neck:
190            mch.rot_neck = self.make_mch_follow_bone(orgs[0], 'neck', 0.5, copy_scale=True)
191            mch.stretch = self.make_mch_stretch_bone(orgs[0], 'STR-neck', orgs[-1])
192        mch.rot_head = self.make_mch_follow_bone(orgs[-1], 'head', 0.0, copy_scale=True)
193
194    def make_mch_stretch_bone(self, org, name, org_head):
195        name = self.copy_bone(org, make_derived_name(name, 'mch'), parent=False)
196        self.get_bone(name).tail = self.get_bone(org_head).head
197        return name
198
199    @stage.parent_bones
200    def parent_mch_control_bones(self):
201        if self.has_neck:
202            self.set_bone_parent(self.bones.mch.rot_neck, self.rig_parent_bone)
203            self.set_bone_parent(self.bones.mch.rot_head, self.bones.ctrl.neck)
204            self.set_bone_parent(self.bones.mch.stretch, self.bones.ctrl.neck)
205        else:
206            self.set_bone_parent(self.bones.mch.rot_head, self.rig_parent_bone)
207
208    @stage.rig_bones
209    def rig_mch_control_bones(self):
210        if self.has_neck:
211            self.rig_mch_stretch_bone(self.bones.mch.stretch, self.bones.ctrl.head)
212
213    def rig_mch_stretch_bone(self, mch, head):
214        self.make_constraint(mch, 'DAMPED_TRACK', head)
215        self.make_constraint(mch, 'STRETCH_TO', head)
216
217    ####################################################
218    # MCH IK chain for the long neck
219
220    @stage.generate_bones
221    def make_mch_ik_chain(self):
222        orgs = self.bones.org
223        if self.long_neck:
224            self.bones.mch.ik = map_list(self.make_mch_ik_bone, orgs[0:-1])
225
226    def make_mch_ik_bone(self, org):
227        return self.copy_bone(org, make_derived_name(org, 'mch', '_ik'), parent=False)
228
229    @stage.parent_bones
230    def parent_mch_ik_chain(self):
231        if self.long_neck:
232            ik = self.bones.mch.ik
233            self.set_bone_parent(ik[0], self.bones.ctrl.tweak[0])
234            self.parent_bone_chain(ik, use_connect=True)
235
236    @stage.rig_bones
237    def rig_mch_ik_chain(self):
238        if self.long_neck:
239            ik = self.bones.mch.ik
240            head = self.bones.ctrl.head
241            for args in zip(count(0), ik):
242                self.rig_mch_ik_bone(*args, len(ik), head)
243
244    def rig_mch_ik_bone(self, i, mch, ik_len, head):
245        if i == ik_len - 1:
246            self.make_constraint(mch, 'IK', head, chain_count=ik_len)
247
248        self.get_bone(mch).ik_stretch = 0.1
249
250    ####################################################
251    # MCH chain for the middle of the neck
252
253    @stage.generate_bones
254    def make_mch_chain(self):
255        orgs = self.bones.org
256        self.bones.mch.chain = map_list(self.make_mch_bone, orgs[1:-1])
257
258    def make_mch_bone(self, org):
259        return self.copy_bone(org, make_derived_name(org, 'mch'), parent=False, scale=1/4)
260
261    @stage.parent_bones
262    def align_mch_chain(self):
263        for mch in self.bones.mch.chain:
264            align_bone_orientation(self.obj, mch, self.bones.ctrl.neck)
265
266    @stage.parent_bones
267    def parent_mch_chain(self):
268        mch = self.bones.mch
269        for bone in mch.chain:
270            self.set_bone_parent(bone, mch.stretch, inherit_scale='NONE')
271
272    @stage.rig_bones
273    def rig_mch_chain(self):
274        chain = self.bones.mch.chain
275        if self.long_neck:
276            ik = self.bones.mch.ik
277            for args in zip(count(0), chain, ik[1:]):
278                self.rig_mch_bone_long(*args, len(chain))
279        else:
280            for args in zip(count(0), chain):
281                self.rig_mch_bone(*args, len(chain))
282
283    def rig_mch_bone_long(self, i, mch, ik, len_mch):
284        ctrl = self.bones.ctrl
285
286        self.make_constraint(mch, 'COPY_LOCATION', ik)
287
288        step = 2/(len_mch+1)
289        xval = (i+1)*step
290        influence = 2*xval - xval**2    #parabolic influence of pivot
291
292        self.make_constraint(
293            mch, 'COPY_LOCATION', ctrl.neck_bend,
294            influence=influence, use_offset=True, space='LOCAL'
295        )
296
297        self.make_constraint(mch, 'COPY_SCALE', ctrl.neck)
298
299    def rig_mch_bone(self, i, mch, len_mch):
300        ctrl = self.bones.ctrl
301
302        nfactor = float((i + 1) / (len_mch + 1))
303        self.make_constraint(
304            mch, 'COPY_ROTATION', ctrl.head,
305            influence=nfactor, space='LOCAL'
306        )
307
308    ####################################################
309    # Tweak bones
310
311    @stage.generate_bones
312    def make_tweak_chain(self):
313        orgs = self.bones.org
314        self.bones.ctrl.tweak = map_list(self.make_tweak_bone, count(0), orgs[0:-1])
315        if not self.has_neck:
316            self.check_connect_tweak(orgs[0])
317
318    @stage.parent_bones
319    def parent_tweak_chain(self):
320        ctrl = self.bones.ctrl
321        mch = self.bones.mch
322
323        if self.has_neck:
324            for args in zip(ctrl.tweak, [ctrl.neck, *mch.chain]):
325                self.set_bone_parent(*args)
326        elif self.connected_tweak:
327            self.set_bone_parent(self.connected_tweak, ctrl.head)
328
329    @stage.rig_bones
330    def generate_neck_tweak_widget(self):
331        # Generate the widget early to override connected parent
332        if self.long_neck:
333            bone = self.bones.ctrl.tweak[0]
334            create_neck_tweak_widget(self.obj, bone, size=1.0)
335
336    ####################################################
337    # ORG and DEF bones
338
339    @stage.generate_bones
340    def register_parent_bones(self):
341        rig = self.rigify_parent or self
342        builder = SwitchParentBuilder(self.generator)
343        builder.register_parent(
344            self, self.bones.org[-1], name='Head',
345            inject_into=rig, exclude_self=True, tags={'head'},
346        )
347
348    @stage.configure_bones
349    def configure_bbone_chain(self):
350        self.get_bone(self.bones.deform[-1]).bone.bbone_segments = 1
351
352    @stage.rig_bones
353    def rig_org_chain(self):
354        if self.has_neck:
355            tweaks = self.bones.ctrl.tweak + [self.bones.ctrl.head]
356        else:
357            tweaks = [self.connected_tweak or self.bones.ctrl.head]
358
359        for args in zip(count(0), self.bones.org, tweaks, tweaks[1:] + [None]):
360            self.rig_org_bone(*args)
361
362
363def create_sample(obj, *, parent=None):
364    # generated by rigify.utils.write_metarig
365    bpy.ops.object.mode_set(mode='EDIT')
366    arm = obj.data
367
368    bones = {}
369
370    bone = arm.edit_bones.new('neck')
371    bone.head[:] = 0.0000, 0.0114, 1.6582
372    bone.tail[:] = 0.0000, -0.0130, 1.7197
373    bone.roll = 0.0000
374    bone.use_connect = False
375    if parent:
376        bone.parent = arm.edit_bones[parent]
377    bones['neck'] = bone.name
378    bone = arm.edit_bones.new('neck.001')
379    bone.head[:] = 0.0000, -0.0130, 1.7197
380    bone.tail[:] = 0.0000, -0.0247, 1.7813
381    bone.roll = 0.0000
382    bone.use_connect = True
383    bone.parent = arm.edit_bones[bones['neck']]
384    bones['neck.001'] = bone.name
385    bone = arm.edit_bones.new('head')
386    bone.head[:] = 0.0000, -0.0247, 1.7813
387    bone.tail[:] = 0.0000, -0.0247, 1.9796
388    bone.roll = 0.0000
389    bone.use_connect = True
390    bone.parent = arm.edit_bones[bones['neck.001']]
391    bones['head'] = bone.name
392
393    bpy.ops.object.mode_set(mode='OBJECT')
394    pbone = obj.pose.bones[bones['neck']]
395    pbone.rigify_type = 'spines.super_head'
396    pbone.lock_location = (False, False, False)
397    pbone.lock_rotation = (False, False, False)
398    pbone.lock_rotation_w = False
399    pbone.lock_scale = (False, False, False)
400    pbone.rotation_mode = 'QUATERNION'
401    try:
402        pbone.rigify_parameters.connect_chain = bool(parent)
403    except AttributeError:
404        pass
405    try:
406        pbone.rigify_parameters.tweak_layers = [False, False, False, False, True, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False]
407    except AttributeError:
408        pass
409    pbone = obj.pose.bones[bones['neck.001']]
410    pbone.rigify_type = ''
411    pbone.lock_location = (False, False, False)
412    pbone.lock_rotation = (False, False, False)
413    pbone.lock_rotation_w = False
414    pbone.lock_scale = (False, False, False)
415    pbone.rotation_mode = 'QUATERNION'
416    pbone = obj.pose.bones[bones['head']]
417    pbone.rigify_type = ''
418    pbone.lock_location = (False, False, False)
419    pbone.lock_rotation = (False, False, False)
420    pbone.lock_rotation_w = False
421    pbone.lock_scale = (False, False, False)
422    pbone.rotation_mode = 'QUATERNION'
423
424    bpy.ops.object.mode_set(mode='EDIT')
425    for bone in arm.edit_bones:
426        bone.select = False
427        bone.select_head = False
428        bone.select_tail = False
429    for b in bones:
430        bone = arm.edit_bones[bones[b]]
431        bone.select = True
432        bone.select_head = True
433        bone.select_tail = True
434        arm.edit_bones.active = bone
435