[Edit Linked Library] Fixes based on code review from Campbell
[blender-addons-contrib.git] / animation_motion_trail.py
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
21
22 bl_info = {
23     "name": "Motion Trail",
24     "author": "Bart Crouch",
25     "version": (3, 1, 2),
26     "blender": (2, 65, 4),
27     "location": "View3D > Toolbar > Motion Trail tab",
28     "warning": "",
29     "description": "Display and edit motion trails in the 3d-view",
30     "wiki_url": "http://wiki.blender.org/index.php/Extensions:2.6/"
31                 "Py/Scripts/Animation/Motion_Trail",
32     "tracker_url": "http://projects.blender.org/tracker/index.php?"
33                     "func=detail&aid=26374",
34     "category": "Animation"}
35
36
37 import bgl
38 import blf
39 import bpy
40 from bpy_extras import view3d_utils
41 import math
42 import mathutils
43
44
45 # fake fcurve class, used if no fcurve is found for a path
46 class fake_fcurve():
47     def __init__(self, object, index, rotation=False, scale=False):
48         # location
49         if not rotation and not scale:
50             self.loc = object.location[index]
51         # scale
52         elif scale:
53             self.loc = object.scale[index]
54         # rotation
55         elif rotation == 'QUATERNION':
56             self.loc = object.rotation_quaternion[index]
57         elif rotation == 'AXIS_ANGLE':
58             self.loc = object.rotation_axis_angle[index]
59         else:
60             self.loc = object.rotation_euler[index]
61         self.keyframe_points = []
62
63     def evaluate(self, frame):
64         return(self.loc)
65
66     def range(self):
67         return([])
68
69
70 # get location curves of the given object
71 def get_curves(object, child=False):
72     if object.animation_data and object.animation_data.action:
73         action = object.animation_data.action
74         if child:
75             # posebone
76             curves = [fc for fc in action.fcurves if len(fc.data_path)>=14 \
77             and fc.data_path[-9:]=='.location' and \
78             child.name in fc.data_path.split("\"")]
79         else:
80             # normal object
81             curves = [fc for fc in action.fcurves if \
82             fc.data_path == 'location']
83     elif object.animation_data and object.animation_data.use_nla:
84         curves = []
85         strips = []
86         for track in object.animation_data.nla_tracks:
87             not_handled = [s for s in track.strips]
88             while not_handled:
89                 current_strip = not_handled.pop(-1)
90                 if current_strip.action:
91                     strips.append(current_strip)
92                 if current_strip.strips:
93                     # meta strip
94                     not_handled += [s for s in current_strip.strips]
95
96         for strip in strips:
97             if child:
98                 # posebone
99                 curves = [fc for fc in strip.action.fcurves if \
100                 len(fc.data_path)>=14 and fc.data_path[-9:]=='.location' \
101                 and child.name in fc.data_path.split("\"")]
102             else:
103                 # normal object
104                 curves = [fc for fc in strip.action.fcurves if \
105                 fc.data_path == 'location']
106             if curves:
107                 # use first strip with location fcurves
108                 break
109     else:
110         # should not happen?
111         curves = []
112
113     # ensure we have three curves per object
114     fcx = None
115     fcy = None
116     fcz = None
117     for fc in curves:
118         if fc.array_index == 0:
119             fcx = fc
120         elif fc.array_index == 1:
121             fcy = fc
122         elif fc.array_index == 2:
123             fcz = fc
124     if fcx == None:
125         fcx = fake_fcurve(object, 0)
126     if fcy == None:
127         fcy = fake_fcurve(object, 1)
128     if fcz == None:
129         fcz = fake_fcurve(object, 2)
130
131     return([fcx, fcy, fcz])
132
133
134 # turn screen coordinates (x,y) into world coordinates vector
135 def screen_to_world(context, x, y):
136     depth_vector = view3d_utils.region_2d_to_vector_3d(\
137         context.region, context.region_data, [x,y])
138     vector = view3d_utils.region_2d_to_location_3d(\
139         context.region, context.region_data, [x,y], depth_vector)
140
141     return(vector)
142
143
144 # turn 3d world coordinates vector into screen coordinate integers (x,y)
145 def world_to_screen(context, vector):
146     prj = context.region_data.perspective_matrix * \
147         mathutils.Vector((vector[0], vector[1], vector[2], 1.0))
148     width_half = context.region.width / 2.0
149     height_half = context.region.height / 2.0
150
151     x = int(width_half + width_half * (prj.x / prj.w))
152     y = int(height_half + height_half * (prj.y / prj.w))
153
154     # correction for corner cases in perspective mode
155     if prj.w < 0:
156         if x < 0:
157             x = context.region.width * 2
158         else:
159             x = context.region.width * -2
160         if y < 0:
161             y = context.region.height * 2
162         else:
163             y = context.region.height * -2
164
165     return(x, y)
166
167
168 # calculate location of display_ob in worldspace
169 def get_location(frame, display_ob, offset_ob, curves):
170     if offset_ob:
171         bpy.context.scene.frame_set(frame)
172         display_mat = getattr(display_ob, "matrix", False)
173         if not display_mat:
174             # posebones have "matrix", objects have "matrix_world"
175             display_mat = display_ob.matrix_world
176         if offset_ob:
177             loc = display_mat.to_translation() + \
178                 offset_ob.matrix_world.to_translation()
179         else:
180             loc = display_mat.to_translation()
181     else:
182         fcx, fcy, fcz = curves
183         locx = fcx.evaluate(frame)
184         locy = fcy.evaluate(frame)
185         locz = fcz.evaluate(frame)
186         loc = mathutils.Vector([locx, locy, locz])
187
188     return(loc)
189
190
191 # get position of keyframes and handles at the start of dragging
192 def get_original_animation_data(context, keyframes):
193     keyframes_ori = {}
194     handles_ori = {}
195
196     if context.active_object and context.active_object.mode == 'POSE':
197         armature_ob = context.active_object
198         objects = [[armature_ob, pb, armature_ob] for pb in \
199             context.selected_pose_bones]
200     else:
201         objects = [[ob, False, False] for ob in context.selected_objects]
202
203     for action_ob, child, offset_ob in objects:
204         if not action_ob.animation_data:
205             continue
206         curves = get_curves(action_ob, child)
207         if len(curves) == 0:
208             continue
209         fcx, fcy, fcz = curves
210         if child:
211             display_ob = child
212         else:
213             display_ob = action_ob
214
215         # get keyframe positions
216         frame_old = context.scene.frame_current
217         keyframes_ori[display_ob.name] = {}
218         for frame in keyframes[display_ob.name]:
219             loc = get_location(frame, display_ob, offset_ob, curves)
220             keyframes_ori[display_ob.name][frame] = [frame, loc]
221
222         # get handle positions
223         handles_ori[display_ob.name] = {}
224         for frame in keyframes[display_ob.name]:
225             handles_ori[display_ob.name][frame] = {}
226             left_x = [frame, fcx.evaluate(frame)]
227             right_x = [frame, fcx.evaluate(frame)]
228             for kf in fcx.keyframe_points:
229                 if kf.co[0] == frame:
230                     left_x = kf.handle_left[:]
231                     right_x = kf.handle_right[:]
232                     break
233             left_y = [frame, fcy.evaluate(frame)]
234             right_y = [frame, fcy.evaluate(frame)]
235             for kf in fcy.keyframe_points:
236                 if kf.co[0] == frame:
237                     left_y = kf.handle_left[:]
238                     right_y = kf.handle_right[:]
239                     break
240             left_z = [frame, fcz.evaluate(frame)]
241             right_z = [frame, fcz.evaluate(frame)]
242             for kf in fcz.keyframe_points:
243                 if kf.co[0] == frame:
244                     left_z = kf.handle_left[:]
245                     right_z = kf.handle_right[:]
246                     break
247             handles_ori[display_ob.name][frame]["left"] = [left_x, left_y,
248                 left_z]
249             handles_ori[display_ob.name][frame]["right"] = [right_x, right_y,
250                 right_z]
251
252         if context.scene.frame_current != frame_old:
253             context.scene.frame_set(frame_old)
254
255     return(keyframes_ori, handles_ori)
256
257
258 # callback function that calculates positions of all things that need be drawn
259 def calc_callback(self, context):
260     if context.active_object and context.active_object.mode == 'POSE':
261         armature_ob = context.active_object
262         objects = [[armature_ob, pb, armature_ob] for pb in \
263             context.selected_pose_bones]
264     else:
265         objects = [[ob, False, False] for ob in context.selected_objects]
266     if objects == self.displayed:
267         selection_change = False
268     else:
269         selection_change = True
270
271     if self.lock and not selection_change and \
272     context.region_data.perspective_matrix == self.perspective and not \
273     context.window_manager.motion_trail.force_update:
274         return
275
276     # dictionaries with key: objectname
277     self.paths = {} # value: list of lists with x, y, color
278     self.keyframes = {} # value: dict with frame as key and [x,y] as value
279     self.handles = {} # value: dict of dicts
280     self.timebeads = {} # value: dict with frame as key and [x,y] as value
281     self.click = {} # value: list of lists with frame, type, loc-vector
282     if selection_change:
283         # value: editbone inverted rotation matrix or None
284         self.edit_bones = {}
285     if selection_change or not self.lock or context.window_manager.\
286     motion_trail.force_update:
287         # contains locations of path, keyframes and timebeads
288         self.cached = {"path":{}, "keyframes":{}, "timebeads_timing":{},
289             "timebeads_speed":{}}
290     if self.cached["path"]:
291         use_cache = True
292     else:
293         use_cache = False
294     self.perspective = context.region_data.perspective_matrix.copy()
295     self.displayed = objects # store, so it can be checked next time
296     context.window_manager.motion_trail.force_update = False
297
298     global_undo = context.user_preferences.edit.use_global_undo
299     context.user_preferences.edit.use_global_undo = False
300
301     for action_ob, child, offset_ob in objects:
302         if selection_change:
303             if not child:
304                 self.edit_bones[action_ob.name] = None
305             else:
306                 bpy.ops.object.mode_set(mode='EDIT')
307                 editbones = action_ob.data.edit_bones
308                 mat = editbones[child.name].matrix.copy().to_3x3().inverted()
309                 bpy.ops.object.mode_set(mode='POSE')
310                 self.edit_bones[child.name] = mat
311
312         if not action_ob.animation_data:
313             continue
314         curves = get_curves(action_ob, child)
315         if len(curves) == 0:
316             continue
317
318         if context.window_manager.motion_trail.path_before == 0:
319             range_min = context.scene.frame_start
320         else:
321             range_min = max(context.scene.frame_start,
322                 context.scene.frame_current - \
323                 context.window_manager.motion_trail.path_before)
324         if context.window_manager.motion_trail.path_after == 0:
325             range_max = context.scene.frame_end
326         else:
327             range_max = min(context.scene.frame_end,
328                 context.scene.frame_current + \
329                 context.window_manager.motion_trail.path_after)
330         fcx, fcy, fcz = curves
331         if child:
332             display_ob = child
333         else:
334             display_ob = action_ob
335
336         # get location data of motion path
337         path = []
338         speeds = []
339         frame_old = context.scene.frame_current
340         step = 11 - context.window_manager.motion_trail.path_resolution
341
342         if not use_cache:
343             if display_ob.name not in self.cached["path"]:
344                 self.cached["path"][display_ob.name] = {}
345         if use_cache and range_min-1 in self.cached["path"][display_ob.name]:
346             prev_loc = self.cached["path"][display_ob.name][range_min-1]
347         else:
348             prev_loc = get_location(range_min-1, display_ob, offset_ob, curves)
349             self.cached["path"][display_ob.name][range_min-1] = prev_loc
350
351         for frame in range(range_min, range_max + 1, step):
352             if use_cache and frame in self.cached["path"][display_ob.name]:
353                 loc = self.cached["path"][display_ob.name][frame]
354             else:
355                 loc = get_location(frame, display_ob, offset_ob, curves)
356                 self.cached["path"][display_ob.name][frame] = loc
357             if not context.region or not context.space_data:
358                 continue
359             x, y = world_to_screen(context, loc)
360             if context.window_manager.motion_trail.path_style == 'simple':
361                 path.append([x, y, [0.0, 0.0, 0.0], frame, action_ob, child])
362             else:
363                 dloc = (loc - prev_loc).length
364                 path.append([x, y, dloc, frame, action_ob, child])
365                 speeds.append(dloc)
366                 prev_loc = loc
367
368         # calculate color of path
369         if context.window_manager.motion_trail.path_style == 'speed':
370             speeds.sort()
371             min_speed = speeds[0]
372             d_speed = speeds[-1] - min_speed
373             for i, [x, y, d_loc, frame, action_ob, child] in enumerate(path):
374                 relative_speed = (d_loc - min_speed) / d_speed # 0.0 to 1.0
375                 red = min(1.0, 2.0 * relative_speed)
376                 blue = min(1.0, 2.0 - (2.0 * relative_speed))
377                 path[i][2] = [red, 0.0, blue]
378         elif context.window_manager.motion_trail.path_style == 'acceleration':
379             accelerations = []
380             prev_speed = 0.0
381             for i, [x, y, d_loc, frame, action_ob, child] in enumerate(path):
382                 accel = d_loc - prev_speed
383                 accelerations.append(accel)
384                 path[i][2] = accel
385                 prev_speed = d_loc
386             accelerations.sort()
387             min_accel = accelerations[0]
388             max_accel = accelerations[-1]
389             for i, [x, y, accel, frame, action_ob, child] in enumerate(path):
390                 if accel < 0:
391                     relative_accel = accel / min_accel # values from 0.0 to 1.0
392                     green = 1.0 - relative_accel
393                     path[i][2] = [1.0, green, 0.0]
394                 elif accel > 0:
395                     relative_accel = accel / max_accel # values from 0.0 to 1.0
396                     red = 1.0 - relative_accel
397                     path[i][2] = [red, 1.0, 0.0]
398                 else:
399                     path[i][2] = [1.0, 1.0, 0.0]
400         self.paths[display_ob.name] = path
401
402         # get keyframes and handles
403         keyframes = {}
404         handle_difs = {}
405         kf_time = []
406         click = []
407         if not use_cache:
408             if display_ob.name not in self.cached["keyframes"]:
409                 self.cached["keyframes"][display_ob.name] = {}
410
411         for fc in curves:
412             for kf in fc.keyframe_points:
413                 # handles for location mode
414                 if context.window_manager.motion_trail.mode == 'location':
415                     if kf.co[0] not in handle_difs:
416                         handle_difs[kf.co[0]] = {"left":mathutils.Vector(),
417                             "right":mathutils.Vector(), "keyframe_loc":None}
418                     handle_difs[kf.co[0]]["left"][fc.array_index] = \
419                         (mathutils.Vector(kf.handle_left[:]) - \
420                         mathutils.Vector(kf.co[:])).normalized()[1]
421                     handle_difs[kf.co[0]]["right"][fc.array_index] = \
422                         (mathutils.Vector(kf.handle_right[:]) - \
423                         mathutils.Vector(kf.co[:])).normalized()[1]
424                 # keyframes
425                 if kf.co[0] in kf_time:
426                     continue
427                 kf_time.append(kf.co[0])
428                 co = kf.co[0]
429
430                 if use_cache and co in \
431                 self.cached["keyframes"][display_ob.name]:
432                     loc = self.cached["keyframes"][display_ob.name][co]
433                 else:
434                     loc = get_location(co, display_ob, offset_ob, curves)
435                     self.cached["keyframes"][display_ob.name][co] = loc
436                 if handle_difs:
437                     handle_difs[co]["keyframe_loc"] = loc
438
439                 x, y = world_to_screen(context, loc)
440                 keyframes[kf.co[0]] = [x, y]
441                 if context.window_manager.motion_trail.mode != 'speed':
442                     # can't select keyframes in speed mode
443                     click.append([kf.co[0], "keyframe",
444                         mathutils.Vector([x,y]), action_ob, child])
445         self.keyframes[display_ob.name] = keyframes
446
447         # handles are only shown in location-altering mode
448         if context.window_manager.motion_trail.mode == 'location' and \
449         context.window_manager.motion_trail.handle_display:
450             # calculate handle positions
451             handles = {}
452             for frame, vecs in handle_difs.items():
453                 if child:
454                     # bone space to world space
455                     mat = self.edit_bones[child.name].copy().inverted()
456                     vec_left = vecs["left"] * mat
457                     vec_right = vecs["right"] * mat
458                 else:
459                     vec_left = vecs["left"]
460                     vec_right = vecs["right"]
461                 if vecs["keyframe_loc"] != None:
462                     vec_keyframe = vecs["keyframe_loc"]
463                 else:
464                     vec_keyframe = get_location(frame, display_ob, offset_ob,
465                         curves)
466                 x_left, y_left = world_to_screen(context, vec_left*2 + \
467                     vec_keyframe)
468                 x_right, y_right = world_to_screen(context, vec_right*2 + \
469                     vec_keyframe)
470                 handles[frame] = {"left":[x_left, y_left],
471                     "right":[x_right, y_right]}
472                 click.append([frame, "handle_left",
473                     mathutils.Vector([x_left, y_left]), action_ob, child])
474                 click.append([frame, "handle_right",
475                     mathutils.Vector([x_right, y_right]), action_ob, child])
476             self.handles[display_ob.name] = handles
477
478         # calculate timebeads for timing mode
479         if context.window_manager.motion_trail.mode == 'timing':
480             timebeads = {}
481             n = context.window_manager.motion_trail.timebeads * (len(kf_time) \
482                 - 1)
483             dframe = (range_max - range_min) / (n + 1)
484             if not use_cache:
485                 if display_ob.name not in self.cached["timebeads_timing"]:
486                     self.cached["timebeads_timing"][display_ob.name] = {}
487
488             for i in range(1, n+1):
489                 frame = range_min + i * dframe
490                 if use_cache and frame in \
491                 self.cached["timebeads_timing"][display_ob.name]:
492                     loc = self.cached["timebeads_timing"][display_ob.name]\
493                         [frame]
494                 else:
495                     loc = get_location(frame, display_ob, offset_ob, curves)
496                     self.cached["timebeads_timing"][display_ob.name][frame] = \
497                         loc
498                 x, y = world_to_screen(context, loc)
499                 timebeads[frame] = [x, y]
500                 click.append([frame, "timebead", mathutils.Vector([x,y]),
501                     action_ob, child])
502             self.timebeads[display_ob.name] = timebeads
503
504         # calculate timebeads for speed mode
505         if context.window_manager.motion_trail.mode == 'speed':
506             angles = dict([[kf, {"left":[], "right":[]}] for kf in \
507                 self.keyframes[display_ob.name]])
508             for fc in curves:
509                 for i, kf in enumerate(fc.keyframe_points):
510                     if i != 0:
511                         angle = mathutils.Vector([-1, 0]).angle(mathutils.\
512                             Vector(kf.handle_left) - mathutils.Vector(kf.co),
513                             0)
514                         if angle != 0:
515                             angles[kf.co[0]]["left"].append(angle)
516                     if i != len(fc.keyframe_points) - 1:
517                         angle = mathutils.Vector([1, 0]).angle(mathutils.\
518                             Vector(kf.handle_right) - mathutils.Vector(kf.co),
519                             0)
520                         if angle != 0:
521                             angles[kf.co[0]]["right"].append(angle)
522             timebeads = {}
523             kf_time.sort()
524             if not use_cache:
525                 if display_ob.name not in self.cached["timebeads_speed"]:
526                     self.cached["timebeads_speed"][display_ob.name] = {}
527
528             for frame, sides in angles.items():
529                 if sides["left"]:
530                     perc = (sum(sides["left"]) / len(sides["left"])) / \
531                         (math.pi / 2)
532                     perc = max(0.4, min(1, perc * 5))
533                     previous = kf_time[kf_time.index(frame) - 1]
534                     bead_frame = frame - perc * ((frame - previous - 2) / 2)
535                     if use_cache and bead_frame in \
536                     self.cached["timebeads_speed"][display_ob.name]:
537                         loc = self.cached["timebeads_speed"][display_ob.name]\
538                             [bead_frame]
539                     else:
540                         loc = get_location(bead_frame, display_ob, offset_ob,
541                             curves)
542                         self.cached["timebeads_speed"][display_ob.name]\
543                             [bead_frame] = loc
544                     x, y = world_to_screen(context, loc)
545                     timebeads[bead_frame] = [x, y]
546                     click.append([bead_frame, "timebead", mathutils.\
547                         Vector([x,y]), action_ob, child])
548                 if sides["right"]:
549                     perc = (sum(sides["right"]) / len(sides["right"])) / \
550                         (math.pi / 2)
551                     perc = max(0.4, min(1, perc * 5))
552                     next = kf_time[kf_time.index(frame) + 1]
553                     bead_frame = frame + perc * ((next - frame - 2) / 2)
554                     if use_cache and bead_frame in \
555                     self.cached["timebeads_speed"][display_ob.name]:
556                         loc = self.cached["timebeads_speed"][display_ob.name]\
557                             [bead_frame]
558                     else:
559                         loc = get_location(bead_frame, display_ob, offset_ob,
560                             curves)
561                         self.cached["timebeads_speed"][display_ob.name]\
562                             [bead_frame] = loc
563                     x, y = world_to_screen(context, loc)
564                     timebeads[bead_frame] = [x, y]
565                     click.append([bead_frame, "timebead", mathutils.\
566                         Vector([x,y]), action_ob, child])
567             self.timebeads[display_ob.name] = timebeads
568
569         # add frame positions to click-list
570         if context.window_manager.motion_trail.frame_display:
571             path = self.paths[display_ob.name]
572             for x, y, color, frame, action_ob, child in path:
573                 click.append([frame, "frame", mathutils.Vector([x,y]),
574                     action_ob, child])
575
576         self.click[display_ob.name] = click
577
578         if context.scene.frame_current != frame_old:
579             context.scene.frame_set(frame_old)
580
581     context.user_preferences.edit.use_global_undo = global_undo
582
583
584 # draw in 3d-view
585 def draw_callback(self, context):
586     # polling
587     if (context.mode not in ('OBJECT', 'POSE') or
588         not context.window_manager.motion_trail.enabled):
589         return
590
591     # display limits
592     if context.window_manager.motion_trail.path_before != 0:
593         limit_min = context.scene.frame_current - \
594             context.window_manager.motion_trail.path_before
595     else:
596         limit_min = -1e6
597     if context.window_manager.motion_trail.path_after != 0:
598         limit_max = context.scene.frame_current + \
599             context.window_manager.motion_trail.path_after
600     else:
601         limit_max = 1e6
602
603     # draw motion path
604     bgl.glEnable(bgl.GL_BLEND)
605     bgl.glLineWidth(context.window_manager.motion_trail.path_width)
606     alpha = 1.0 - (context.window_manager.motion_trail.path_transparency / \
607         100.0)
608     if context.window_manager.motion_trail.path_style == 'simple':
609         bgl.glColor4f(0.0, 0.0, 0.0, alpha)
610         for objectname, path in self.paths.items():
611             bgl.glBegin(bgl.GL_LINE_STRIP)
612             for x, y, color, frame, action_ob, child in path:
613                 if frame < limit_min or frame > limit_max:
614                     continue
615                 bgl.glVertex2i(x, y)
616             bgl.glEnd()
617     else:
618         for objectname, path in self.paths.items():
619             for i, [x, y, color, frame, action_ob, child] in enumerate(path):
620                 if frame < limit_min or frame > limit_max:
621                     continue
622                 r, g, b = color
623                 if i != 0:
624                     prev_path = path[i-1]
625                     halfway = [(x + prev_path[0])/2, (y + prev_path[1])/2]
626                     bgl.glColor4f(r, g, b, alpha)
627                     bgl.glBegin(bgl.GL_LINE_STRIP)
628                     bgl.glVertex2i(int(halfway[0]), int(halfway[1]))
629                     bgl.glVertex2i(x, y)
630                     bgl.glEnd()
631                 if i != len(path) - 1:
632                     next_path = path[i+1]
633                     halfway = [(x + next_path[0])/2, (y + next_path[1])/2]
634                     bgl.glColor4f(r, g, b, alpha)
635                     bgl.glBegin(bgl.GL_LINE_STRIP)
636                     bgl.glVertex2i(x, y)
637                     bgl.glVertex2i(int(halfway[0]), int(halfway[1]))
638                     bgl.glEnd()
639
640     # draw frames
641     if context.window_manager.motion_trail.frame_display:
642         bgl.glColor4f(1.0, 1.0, 1.0, 1.0)
643         bgl.glPointSize(1)
644         bgl.glBegin(bgl.GL_POINTS)
645         for objectname, path in self.paths.items():
646             for x, y, color, frame, action_ob, child in path:
647                 if frame < limit_min or frame > limit_max:
648                     continue
649                 if self.active_frame and objectname == self.active_frame[0] \
650                 and abs(frame - self.active_frame[1]) < 1e-4:
651                     bgl.glEnd()
652                     bgl.glColor4f(1.0, 0.5, 0.0, 1.0)
653                     bgl.glPointSize(3)
654                     bgl.glBegin(bgl.GL_POINTS)
655                     bgl.glVertex2i(x,y)
656                     bgl.glEnd()
657                     bgl.glColor4f(1.0, 1.0, 1.0, 1.0)
658                     bgl.glPointSize(1)
659                     bgl.glBegin(bgl.GL_POINTS)
660                 else:
661                     bgl.glVertex2i(x,y)
662         bgl.glEnd()
663
664     # time beads are shown in speed and timing modes
665     if context.window_manager.motion_trail.mode in ('speed', 'timing'):
666         bgl.glColor4f(0.0, 1.0, 0.0, 1.0)
667         bgl.glPointSize(4)
668         bgl.glBegin(bgl.GL_POINTS)
669         for objectname, values in self.timebeads.items():
670             for frame, coords in values.items():
671                 if frame < limit_min or frame > limit_max:
672                     continue
673                 if self.active_timebead and \
674                 objectname == self.active_timebead[0] and \
675                 abs(frame - self.active_timebead[1]) < 1e-4:
676                     bgl.glEnd()
677                     bgl.glColor4f(1.0, 0.5, 0.0, 1.0)
678                     bgl.glBegin(bgl.GL_POINTS)
679                     bgl.glVertex2i(coords[0], coords[1])
680                     bgl.glEnd()
681                     bgl.glColor4f(0.0, 1.0, 0.0, 1.0)
682                     bgl.glBegin(bgl.GL_POINTS)
683                 else:
684                     bgl.glVertex2i(coords[0], coords[1])
685         bgl.glEnd()
686
687     # handles are only shown in location mode
688     if context.window_manager.motion_trail.mode == 'location':
689         # draw handle-lines
690         bgl.glColor4f(0.0, 0.0, 0.0, 1.0)
691         bgl.glLineWidth(1)
692         bgl.glBegin(bgl.GL_LINES)
693         for objectname, values in self.handles.items():
694             for frame, sides in values.items():
695                 if frame < limit_min or frame > limit_max:
696                     continue
697                 for side, coords in sides.items():
698                     if self.active_handle and \
699                     objectname == self.active_handle[0] and \
700                     side == self.active_handle[2] and \
701                     abs(frame - self.active_handle[1]) < 1e-4:
702                         bgl.glEnd()
703                         bgl.glColor4f(.75, 0.25, 0.0, 1.0)
704                         bgl.glBegin(bgl.GL_LINES)
705                         bgl.glVertex2i(self.keyframes[objectname][frame][0],
706                             self.keyframes[objectname][frame][1])
707                         bgl.glVertex2i(coords[0], coords[1])
708                         bgl.glEnd()
709                         bgl.glColor4f(0.0, 0.0, 0.0, 1.0)
710                         bgl.glBegin(bgl.GL_LINES)
711                     else:
712                         bgl.glVertex2i(self.keyframes[objectname][frame][0],
713                             self.keyframes[objectname][frame][1])
714                         bgl.glVertex2i(coords[0], coords[1])
715         bgl.glEnd()
716
717         # draw handles
718         bgl.glColor4f(1.0, 1.0, 0.0, 1.0)
719         bgl.glPointSize(4)
720         bgl.glBegin(bgl.GL_POINTS)
721         for objectname, values in self.handles.items():
722             for frame, sides in values.items():
723                 if frame < limit_min or frame > limit_max:
724                     continue
725                 for side, coords in sides.items():
726                     if self.active_handle and \
727                     objectname == self.active_handle[0] and \
728                     side == self.active_handle[2] and \
729                     abs(frame - self.active_handle[1]) < 1e-4:
730                         bgl.glEnd()
731                         bgl.glColor4f(1.0, 0.5, 0.0, 1.0)
732                         bgl.glBegin(bgl.GL_POINTS)
733                         bgl.glVertex2i(coords[0], coords[1])
734                         bgl.glEnd()
735                         bgl.glColor4f(1.0, 1.0, 0.0, 1.0)
736                         bgl.glBegin(bgl.GL_POINTS)
737                     else:
738                         bgl.glVertex2i(coords[0], coords[1])
739         bgl.glEnd()
740
741     # draw keyframes
742     bgl.glColor4f(1.0, 1.0, 0.0, 1.0)
743     bgl.glPointSize(6)
744     bgl.glBegin(bgl.GL_POINTS)
745     for objectname, values in self.keyframes.items():
746         for frame, coords in values.items():
747             if frame < limit_min or frame > limit_max:
748                 continue
749             if self.active_keyframe and \
750             objectname == self.active_keyframe[0] and \
751             abs(frame - self.active_keyframe[1]) < 1e-4:
752                 bgl.glEnd()
753                 bgl.glColor4f(1.0, 0.5, 0.0, 1.0)
754                 bgl.glBegin(bgl.GL_POINTS)
755                 bgl.glVertex2i(coords[0], coords[1])
756                 bgl.glEnd()
757                 bgl.glColor4f(1.0, 1.0, 0.0, 1.0)
758                 bgl.glBegin(bgl.GL_POINTS)
759             else:
760                 bgl.glVertex2i(coords[0], coords[1])
761     bgl.glEnd()
762
763     # draw keyframe-numbers
764     if context.window_manager.motion_trail.keyframe_numbers:
765         blf.size(0, 12, 72)
766         bgl.glColor4f(1.0, 1.0, 0.0, 1.0)
767         for objectname, values in self.keyframes.items():
768             for frame, coords in values.items():
769                 if frame < limit_min or frame > limit_max:
770                     continue
771                 blf.position(0, coords[0] + 3, coords[1] + 3, 0)
772                 text = str(frame).split(".")
773                 if len(text) == 1:
774                     text = text[0]
775                 elif len(text[1]) == 1 and text[1] == "0":
776                     text = text[0]
777                 else:
778                     text = text[0] + "." + text[1][0]
779                 if self.active_keyframe and \
780                 objectname == self.active_keyframe[0] and \
781                 abs(frame - self.active_keyframe[1]) < 1e-4:
782                     bgl.glColor4f(1.0, 0.5, 0.0, 1.0)
783                     blf.draw(0, text)
784                     bgl.glColor4f(1.0, 1.0, 0.0, 1.0)
785                 else:
786                     blf.draw(0, text)
787
788     # restore opengl defaults
789     bgl.glLineWidth(1)
790     bgl.glDisable(bgl.GL_BLEND)
791     bgl.glColor4f(0.0, 0.0, 0.0, 1.0)
792     bgl.glPointSize(1)
793
794
795 # change data based on mouse movement
796 def drag(context, event, drag_mouse_ori, active_keyframe, active_handle,
797 active_timebead, keyframes_ori, handles_ori, edit_bones):
798     # change 3d-location of keyframe
799     if context.window_manager.motion_trail.mode == 'location' and \
800     active_keyframe:
801         objectname, frame, frame_ori, action_ob, child = active_keyframe
802         if child:
803             mat = action_ob.matrix_world.copy().inverted() * \
804                 edit_bones[child.name].copy().to_4x4()
805         else:
806             mat = 1
807
808         mouse_ori_world = screen_to_world(context, drag_mouse_ori[0],
809             drag_mouse_ori[1]) * mat
810         vector = screen_to_world(context, event.mouse_region_x,
811             event.mouse_region_y) * mat
812         d = vector - mouse_ori_world
813
814         loc_ori_ws = keyframes_ori[objectname][frame][1]
815         loc_ori_bs = loc_ori_ws * mat
816         new_loc = loc_ori_bs + d
817         curves = get_curves(action_ob, child)
818
819         for i, curve in enumerate(curves):
820             for kf in curve.keyframe_points:
821                 if kf.co[0] == frame:
822                     kf.co[1] = new_loc[i]
823                     kf.handle_left[1] = handles_ori[objectname][frame]\
824                         ["left"][i][1] + d[i]
825                     kf.handle_right[1] = handles_ori[objectname][frame]\
826                         ["right"][i][1] + d[i]
827                     break
828
829     # change 3d-location of handle
830     elif context.window_manager.motion_trail.mode == 'location' and \
831     active_handle:
832         objectname, frame, side, action_ob, child = active_handle
833         if child:
834             mat = action_ob.matrix_world.copy().inverted() * \
835                 edit_bones[child.name].copy().to_4x4()
836         else:
837             mat = 1
838
839         mouse_ori_world = screen_to_world(context, drag_mouse_ori[0],
840             drag_mouse_ori[1]) * mat
841         vector = screen_to_world(context, event.mouse_region_x,
842             event.mouse_region_y) * mat
843         d = vector - mouse_ori_world
844         curves = get_curves(action_ob, child)
845
846         for i, curve in enumerate(curves):
847             for kf in curve.keyframe_points:
848                 if kf.co[0] == frame:
849                     if side == "left":
850                         # change handle type, if necessary
851                         if kf.handle_left_type in (
852                             'AUTO',
853                             'AUTO_CLAMPED',
854                             'ANIM_CLAMPED'):
855                             kf.handle_left_type = 'ALIGNED'
856                         elif kf.handle_left_type == 'VECTOR':
857                             kf.handle_left_type = 'FREE'
858                         # change handle position(s)
859                         kf.handle_left[1] = handles_ori[objectname][frame]\
860                             ["left"][i][1] + d[i]
861                         if kf.handle_left_type in (
862                             'ALIGNED',
863                             'ANIM_CLAMPED',
864                             'AUTO',
865                             'AUTO_CLAMPED'):
866                             dif = (abs(handles_ori[objectname][frame]["right"]\
867                                 [i][0] - kf.co[0]) / abs(kf.handle_left[0] - \
868                                 kf.co[0])) * d[i]
869                             kf.handle_right[1] = handles_ori[objectname]\
870                                 [frame]["right"][i][1] - dif
871                     elif side == "right":
872                         # change handle type, if necessary
873                         if kf.handle_right_type in (
874                             'AUTO',
875                             'AUTO_CLAMPED',
876                             'ANIM_CLAMPED'):
877                             kf.handle_left_type = 'ALIGNED'
878                             kf.handle_right_type = 'ALIGNED'
879                         elif kf.handle_right_type == 'VECTOR':
880                             kf.handle_left_type = 'FREE'
881                             kf.handle_right_type = 'FREE'
882                         # change handle position(s)
883                         kf.handle_right[1] = handles_ori[objectname][frame]\
884                             ["right"][i][1] + d[i]
885                         if kf.handle_right_type in (
886                             'ALIGNED',
887                             'ANIM_CLAMPED',
888                             'AUTO',
889                             'AUTO_CLAMPED'):
890                             dif = (abs(handles_ori[objectname][frame]["left"]\
891                                 [i][0] - kf.co[0]) / abs(kf.handle_right[0] - \
892                                 kf.co[0])) * d[i]
893                             kf.handle_left[1] = handles_ori[objectname]\
894                                 [frame]["left"][i][1] - dif
895                     break
896
897     # change position of all keyframes on timeline
898     elif context.window_manager.motion_trail.mode == 'timing' and \
899     active_timebead:
900         objectname, frame, frame_ori, action_ob, child = active_timebead
901         curves = get_curves(action_ob, child)
902         ranges = [val for c in curves for val in c.range()]
903         ranges.sort()
904         range_min = round(ranges[0])
905         range_max = round(ranges[-1])
906         range = range_max - range_min
907         dx_screen = -(mathutils.Vector([event.mouse_region_x,
908             event.mouse_region_y]) - drag_mouse_ori)[0]
909         dx_screen = dx_screen / context.region.width * range
910         new_frame = frame + dx_screen
911         shift_low = max(1e-4, (new_frame - range_min) / (frame - range_min))
912         shift_high = max(1e-4, (range_max - new_frame) / (range_max - frame))
913
914         new_mapping = {}
915         for i, curve in enumerate(curves):
916             for j, kf in enumerate(curve.keyframe_points):
917                 frame_map = kf.co[0]
918                 if frame_map < range_min + 1e-4 or \
919                 frame_map > range_max - 1e-4:
920                     continue
921                 frame_ori = False
922                 for f in keyframes_ori[objectname]:
923                     if abs(f - frame_map) < 1e-4:
924                         frame_ori = keyframes_ori[objectname][f][0]
925                         value_ori = keyframes_ori[objectname][f]
926                         break
927                 if not frame_ori:
928                     continue
929                 if frame_ori <= frame:
930                     frame_new = (frame_ori - range_min) * shift_low + \
931                         range_min
932                 else:
933                     frame_new = range_max - (range_max - frame_ori) * \
934                         shift_high
935                 frame_new = max(range_min + j, min(frame_new, range_max - \
936                     (len(curve.keyframe_points)-j)+1))
937                 d_frame = frame_new - frame_ori
938                 if frame_new not in new_mapping:
939                     new_mapping[frame_new] = value_ori
940                 kf.co[0] = frame_new
941                 kf.handle_left[0] = handles_ori[objectname][frame_ori]\
942                     ["left"][i][0] + d_frame
943                 kf.handle_right[0] = handles_ori[objectname][frame_ori]\
944                     ["right"][i][0] + d_frame
945         del keyframes_ori[objectname]
946         keyframes_ori[objectname] = {}
947         for new_frame, value in new_mapping.items():
948             keyframes_ori[objectname][new_frame] = value
949
950     # change position of active keyframe on the timeline
951     elif context.window_manager.motion_trail.mode == 'timing' and \
952     active_keyframe:
953         objectname, frame, frame_ori, action_ob, child = active_keyframe
954         if child:
955             mat = action_ob.matrix_world.copy().inverted() * \
956                 edit_bones[child.name].copy().to_4x4()
957         else:
958             mat = action_ob.matrix_world.copy().inverted()
959
960         mouse_ori_world = screen_to_world(context, drag_mouse_ori[0],
961             drag_mouse_ori[1]) * mat
962         vector = screen_to_world(context, event.mouse_region_x,
963             event.mouse_region_y) * mat
964         d = vector - mouse_ori_world
965
966         locs_ori = [[f_ori, coords] for f_mapped, [f_ori, coords] in \
967             keyframes_ori[objectname].items()]
968         locs_ori.sort()
969         direction = 1
970         range = False
971         for i, [f_ori, coords] in enumerate(locs_ori):
972             if abs(frame_ori - f_ori) < 1e-4:
973                 if i == 0:
974                     # first keyframe, nothing before it
975                     direction = -1
976                     range = [f_ori, locs_ori[i+1][0]]
977                 elif i == len(locs_ori) - 1:
978                     # last keyframe, nothing after it
979                     range = [locs_ori[i-1][0], f_ori]
980                 else:
981                     current = mathutils.Vector(coords)
982                     next = mathutils.Vector(locs_ori[i+1][1])
983                     previous = mathutils.Vector(locs_ori[i-1][1])
984                     angle_to_next = d.angle(next - current, 0)
985                     angle_to_previous = d.angle(previous-current, 0)
986                     if angle_to_previous < angle_to_next:
987                         # mouse movement is in direction of previous keyframe
988                         direction = -1
989                     range = [locs_ori[i-1][0], locs_ori[i+1][0]]
990                 break
991         direction *= -1 # feels more natural in 3d-view
992         if not range:
993             # keyframe not found, is impossible, but better safe than sorry
994             return(active_keyframe, active_timebead, keyframes_ori)
995         # calculate strength of movement
996         d_screen = mathutils.Vector([event.mouse_region_x,
997             event.mouse_region_y]) - drag_mouse_ori
998         if d_screen.length != 0:
999             d_screen = d_screen.length / (abs(d_screen[0])/d_screen.length*\
1000                 context.region.width + abs(d_screen[1])/d_screen.length*\
1001                 context.region.height)
1002             d_screen *= direction  # d_screen value ranges from -1.0 to 1.0
1003         else:
1004             d_screen = 0.0
1005         new_frame = d_screen * (range[1] - range[0]) + frame_ori
1006         max_frame = range[1]
1007         if max_frame == frame_ori:
1008             max_frame += 1
1009         min_frame = range[0]
1010         if min_frame == frame_ori:
1011             min_frame -= 1
1012         new_frame = min(max_frame - 1, max(min_frame + 1, new_frame))
1013         d_frame = new_frame - frame_ori
1014         curves = get_curves(action_ob, child)
1015
1016         for i, curve in enumerate(curves):
1017             for kf in curve.keyframe_points:
1018                 if abs(kf.co[0] - frame) < 1e-4:
1019                     kf.co[0] = new_frame
1020                     kf.handle_left[0] = handles_ori[objectname][frame_ori]\
1021                         ["left"][i][0] + d_frame
1022                     kf.handle_right[0] = handles_ori[objectname][frame_ori]\
1023                         ["right"][i][0] + d_frame
1024                     break
1025         active_keyframe = [objectname, new_frame, frame_ori, action_ob, child]
1026
1027     # change position of active timebead on the timeline, thus altering speed
1028     elif context.window_manager.motion_trail.mode == 'speed' and \
1029     active_timebead:
1030         objectname, frame, frame_ori, action_ob, child = active_timebead
1031         if child:
1032             mat = action_ob.matrix_world.copy().inverted() * \
1033                 edit_bones[child.name].copy().to_4x4()
1034         else:
1035             mat = 1
1036
1037         mouse_ori_world = screen_to_world(context, drag_mouse_ori[0],
1038             drag_mouse_ori[1]) * mat
1039         vector = screen_to_world(context, event.mouse_region_x,
1040             event.mouse_region_y) * mat
1041         d = vector - mouse_ori_world
1042
1043         # determine direction (to next or previous keyframe)
1044         curves = get_curves(action_ob, child)
1045         fcx, fcy, fcz = curves
1046         locx = fcx.evaluate(frame_ori)
1047         locy = fcy.evaluate(frame_ori)
1048         locz = fcz.evaluate(frame_ori)
1049         loc_ori = mathutils.Vector([locx, locy, locz]) # bonespace
1050         keyframes = [kf for kf in keyframes_ori[objectname]]
1051         keyframes.append(frame_ori)
1052         keyframes.sort()
1053         frame_index = keyframes.index(frame_ori)
1054         kf_prev = keyframes[frame_index - 1]
1055         kf_next = keyframes[frame_index + 1]
1056         vec_prev = (mathutils.Vector(keyframes_ori[objectname][kf_prev][1]) \
1057             * mat - loc_ori).normalized()
1058         vec_next = (mathutils.Vector(keyframes_ori[objectname][kf_next][1]) \
1059             * mat - loc_ori).normalized()
1060         d_normal = d.copy().normalized()
1061         dist_to_next = (d_normal - vec_next).length
1062         dist_to_prev = (d_normal - vec_prev).length
1063         if dist_to_prev < dist_to_next:
1064             direction = 1
1065         else:
1066             direction = -1
1067
1068         if (kf_next - frame_ori) < (frame_ori - kf_prev):
1069             kf_bead = kf_next
1070             side = "left"
1071         else:
1072             kf_bead = kf_prev
1073             side = "right"
1074         d_frame = d.length * direction * 2 # *2 to make it more sensitive
1075
1076         angles = []
1077         for i, curve in enumerate(curves):
1078             for kf in curve.keyframe_points:
1079                 if abs(kf.co[0] - kf_bead) < 1e-4:
1080                     if side == "left":
1081                         # left side
1082                         kf.handle_left[0] = min(handles_ori[objectname]\
1083                             [kf_bead]["left"][i][0] + d_frame, kf_bead - 1)
1084                         angle = mathutils.Vector([-1, 0]).angle(mathutils.\
1085                             Vector(kf.handle_left) - mathutils.Vector(kf.co),
1086                             0)
1087                         if angle != 0:
1088                             angles.append(angle)
1089                     else:
1090                         # right side
1091                         kf.handle_right[0] = max(handles_ori[objectname]\
1092                             [kf_bead]["right"][i][0] + d_frame, kf_bead + 1)
1093                         angle = mathutils.Vector([1, 0]).angle(mathutils.\
1094                             Vector(kf.handle_right) - mathutils.Vector(kf.co),
1095                             0)
1096                         if angle != 0:
1097                             angles.append(angle)
1098                     break
1099
1100         # update frame of active_timebead
1101         perc = (sum(angles) / len(angles)) / (math.pi / 2)
1102         perc = max(0.4, min(1, perc * 5))
1103         if side == "left":
1104             bead_frame = kf_bead - perc * ((kf_bead - kf_prev - 2) / 2)
1105         else:
1106             bead_frame = kf_bead + perc * ((kf_next - kf_bead - 2) / 2)
1107         active_timebead = [objectname, bead_frame, frame_ori, action_ob, child]
1108
1109     return(active_keyframe, active_timebead, keyframes_ori)
1110
1111
1112 # revert changes made by dragging
1113 def cancel_drag(context, active_keyframe, active_handle, active_timebead,
1114 keyframes_ori, handles_ori, edit_bones):
1115     # revert change in 3d-location of active keyframe and its handles
1116     if context.window_manager.motion_trail.mode == 'location' and \
1117     active_keyframe:
1118         objectname, frame, frame_ori, active_ob, child = active_keyframe
1119         curves = get_curves(active_ob, child)
1120         loc_ori = keyframes_ori[objectname][frame][1]
1121         if child:
1122             loc_ori = loc_ori * edit_bones[child.name] * \
1123                 active_ob.matrix_world.copy().inverted()
1124         for i, curve in enumerate(curves):
1125             for kf in curve.keyframe_points:
1126                 if kf.co[0] == frame:
1127                     kf.co[1] = loc_ori[i]
1128                     kf.handle_left[1] = handles_ori[objectname][frame]\
1129                         ["left"][i][1]
1130                     kf.handle_right[1] = handles_ori[objectname][frame]\
1131                         ["right"][i][1]
1132                     break
1133
1134     # revert change in 3d-location of active handle
1135     elif context.window_manager.motion_trail.mode == 'location' and \
1136     active_handle:
1137         objectname, frame, side, active_ob, child = active_handle
1138         curves = get_curves(active_ob, child)
1139         for i, curve in enumerate(curves):
1140             for kf in curve.keyframe_points:
1141                 if kf.co[0] == frame:
1142                     kf.handle_left[1] = handles_ori[objectname][frame]\
1143                         ["left"][i][1]
1144                     kf.handle_right[1] = handles_ori[objectname][frame]\
1145                         ["right"][i][1]
1146                     break
1147
1148     # revert position of all keyframes and handles on timeline
1149     elif context.window_manager.motion_trail.mode == 'timing' and \
1150     active_timebead:
1151         objectname, frame, frame_ori, active_ob, child = active_timebead
1152         curves = get_curves(active_ob, child)
1153         for i, curve in enumerate(curves):
1154             for kf in curve.keyframe_points:
1155                 for kf_ori, [frame_ori, loc] in keyframes_ori[objectname].\
1156                 items():
1157                     if abs(kf.co[0] - kf_ori) < 1e-4:
1158                         kf.co[0] = frame_ori
1159                         kf.handle_left[0] = handles_ori[objectname]\
1160                             [frame_ori]["left"][i][0]
1161                         kf.handle_right[0] = handles_ori[objectname]\
1162                             [frame_ori]["right"][i][0]
1163                         break
1164
1165     # revert position of active keyframe and its handles on the timeline
1166     elif context.window_manager.motion_trail.mode == 'timing' and \
1167     active_keyframe:
1168         objectname, frame, frame_ori, active_ob, child = active_keyframe
1169         curves = get_curves(active_ob, child)
1170         for i, curve in enumerate(curves):
1171             for kf in curve.keyframe_points:
1172                 if abs(kf.co[0] - frame) < 1e-4:
1173                     kf.co[0] = keyframes_ori[objectname][frame_ori][0]
1174                     kf.handle_left[0] = handles_ori[objectname][frame_ori]\
1175                         ["left"][i][0]
1176                     kf.handle_right[0] = handles_ori[objectname][frame_ori]\
1177                         ["right"][i][0]
1178                     break
1179         active_keyframe = [objectname, frame_ori, frame_ori, active_ob, child]
1180
1181     # revert position of handles on the timeline
1182     elif context.window_manager.motion_trail.mode == 'speed' and \
1183     active_timebead:
1184         objectname, frame, frame_ori, active_ob, child = active_timebead
1185         curves = get_curves(active_ob, child)
1186         keyframes = [kf for kf in keyframes_ori[objectname]]
1187         keyframes.append(frame_ori)
1188         keyframes.sort()
1189         frame_index = keyframes.index(frame_ori)
1190         kf_prev = keyframes[frame_index - 1]
1191         kf_next = keyframes[frame_index + 1]
1192         if (kf_next - frame_ori) < (frame_ori - kf_prev):
1193             kf_frame = kf_next
1194         else:
1195             kf_frame = kf_prev
1196         for i, curve in enumerate(curves):
1197             for kf in curve.keyframe_points:
1198                 if kf.co[0] == kf_frame:
1199                     kf.handle_left[0] = handles_ori[objectname][kf_frame]\
1200                         ["left"][i][0]
1201                     kf.handle_right[0] = handles_ori[objectname][kf_frame]\
1202                         ["right"][i][0]
1203                     break
1204         active_timebead = [objectname, frame_ori, frame_ori, active_ob, child]
1205
1206     return(active_keyframe, active_timebead)
1207
1208
1209 # return the handle type of the active selection
1210 def get_handle_type(active_keyframe, active_handle):
1211     if active_keyframe:
1212         objectname, frame, side, action_ob, child = active_keyframe
1213         side = "both"
1214     elif active_handle:
1215         objectname, frame, side, action_ob, child = active_handle
1216     else:
1217         # no active handle(s)
1218         return(False)
1219
1220     # properties used when changing handle type
1221     bpy.context.window_manager.motion_trail.handle_type_frame = frame
1222     bpy.context.window_manager.motion_trail.handle_type_side = side
1223     bpy.context.window_manager.motion_trail.handle_type_action_ob = \
1224         action_ob.name
1225     if child:
1226         bpy.context.window_manager.motion_trail.handle_type_child = child.name
1227     else:
1228         bpy.context.window_manager.motion_trail.handle_type_child = ""
1229
1230     curves = get_curves(action_ob, child=child)
1231     for c in curves:
1232         for kf in c.keyframe_points:
1233             if kf.co[0] == frame:
1234                 if side in ("left", "both"):
1235                     return(kf.handle_left_type)
1236                 else:
1237                     return(kf.handle_right_type)
1238
1239     return("AUTO")
1240
1241
1242 # turn the given frame into a keyframe
1243 def insert_keyframe(self, context, frame):
1244     objectname, frame, frame, action_ob, child = frame
1245     curves = get_curves(action_ob, child)
1246     for c in curves:
1247         y = c.evaluate(frame)
1248         if c.keyframe_points:
1249             c.keyframe_points.insert(frame, y)
1250
1251     bpy.context.window_manager.motion_trail.force_update = True
1252     calc_callback(self, context)
1253
1254
1255 # change the handle type of the active selection
1256 def set_handle_type(self, context):
1257     if not context.window_manager.motion_trail.handle_type_enabled:
1258         return
1259     if context.window_manager.motion_trail.handle_type_old == \
1260     context.window_manager.motion_trail.handle_type:
1261         # function called because of selection change, not change in type
1262         return
1263     context.window_manager.motion_trail.handle_type_old = \
1264         context.window_manager.motion_trail.handle_type
1265
1266     frame = bpy.context.window_manager.motion_trail.handle_type_frame
1267     side = bpy.context.window_manager.motion_trail.handle_type_side
1268     action_ob = bpy.context.window_manager.motion_trail.handle_type_action_ob
1269     action_ob = bpy.data.objects[action_ob]
1270     child = bpy.context.window_manager.motion_trail.handle_type_child
1271     if child:
1272         child = action_ob.pose.bones[child]
1273     new_type = context.window_manager.motion_trail.handle_type
1274
1275     curves = get_curves(action_ob, child=child)
1276     for c in curves:
1277         for kf in c.keyframe_points:
1278             if kf.co[0] == frame:
1279                 # align if necessary
1280                 if side in ("right", "both") and new_type in (
1281                     "AUTO", "AUTO_CLAMPED", "ALIGNED"):
1282                     # change right handle
1283                     normal = (kf.co - kf.handle_left).normalized()
1284                     size = (kf.handle_right[0] - kf.co[0]) / normal[0]
1285                     normal = normal*size + kf.co
1286                     kf.handle_right[1] = normal[1]
1287                 elif side == "left" and new_type in (
1288                     "AUTO", "AUTO_CLAMPED", "ALIGNED"):
1289                     # change left handle
1290                     normal = (kf.co - kf.handle_right).normalized()
1291                     size = (kf.handle_left[0] - kf.co[0]) / normal[0]
1292                     normal = normal*size + kf.co
1293                     kf.handle_left[1] = normal[1]
1294                 # change type
1295                 if side in ("left", "both"):
1296                     kf.handle_left_type = new_type
1297                 if side in ("right", "both"):
1298                     kf.handle_right_type = new_type
1299
1300     context.window_manager.motion_trail.force_update = True
1301
1302
1303 class MotionTrailOperator(bpy.types.Operator):
1304     """Edit motion trails in 3d-view"""
1305     bl_idname = "view3d.motion_trail"
1306     bl_label = "Motion Trail"
1307
1308     _handle_calc = None
1309     _handle_draw = None
1310
1311     @staticmethod
1312     def handle_add(self, context):
1313         MotionTrailOperator._handle_calc = bpy.types.SpaceView3D.draw_handler_add(
1314             calc_callback, (self, context), 'WINDOW', 'POST_VIEW')
1315         MotionTrailOperator._handle_draw = bpy.types.SpaceView3D.draw_handler_add(
1316             draw_callback, (self, context), 'WINDOW', 'POST_PIXEL')
1317
1318     @staticmethod
1319     def handle_remove():
1320         if MotionTrailOperator._handle_calc is not None:
1321             bpy.types.SpaceView3D.draw_handler_remove(MotionTrailOperator._handle_calc, 'WINDOW')
1322         if MotionTrailOperator._handle_draw is not None:
1323             bpy.types.SpaceView3D.draw_handler_remove(MotionTrailOperator._handle_draw, 'WINDOW')
1324         MotionTrailOperator._handle_calc = None
1325         MotionTrailOperator._handle_draw = None
1326
1327
1328     def modal(self, context, event):
1329
1330         #XXX Required, or custom transform.translate will break!
1331         #XXX If one disables and re-enables motion trail, modal op will still be running,
1332         #XXX default translate op will unintentionally get called, followed by custom translate.
1333         if not context.window_manager.motion_trail.enabled:
1334             MotionTrailOperator.handle_remove()
1335             context.area.tag_redraw()
1336             return {'FINISHED'}
1337
1338         if not context.area or not context.region or event.type == 'NONE':
1339             context.area.tag_redraw()
1340             return {'PASS_THROUGH'}
1341
1342         select = context.user_preferences.inputs.select_mouse
1343         if (not context.active_object or
1344             context.active_object.mode not in ('OBJECT', 'POSE')):
1345             if self.drag:
1346                 self.drag = False
1347                 self.lock = True
1348                 context.window_manager.motion_trail.force_update = True
1349             # default hotkeys should still work
1350             if event.type == self.transform_key and event.value == 'PRESS':
1351                 if bpy.ops.transform.translate.poll():
1352                     bpy.ops.transform.translate('INVOKE_DEFAULT')
1353             elif event.type == select + 'MOUSE' and event.value == 'PRESS' \
1354             and not self.drag and not event.shift and not event.alt \
1355             and not event.ctrl:
1356                 if bpy.ops.view3d.select.poll():
1357                     bpy.ops.view3d.select('INVOKE_DEFAULT')
1358             elif event.type == 'LEFTMOUSE' and event.value == 'PRESS' and not\
1359             event.alt and not event.ctrl and not event.shift:
1360                 if eval("bpy.ops."+self.left_action+".poll()"):
1361                     eval("bpy.ops."+self.left_action+"('INVOKE_DEFAULT')")
1362             return {'PASS_THROUGH'}
1363         # check if event was generated within 3d-window, dragging is exception
1364         if not self.drag:
1365             if not (0 < event.mouse_region_x < context.region.width) or \
1366             not (0 < event.mouse_region_y < context.region.height):
1367                 return {'PASS_THROUGH'}
1368
1369         if (event.type == self.transform_key and event.value == 'PRESS' and
1370            (self.active_keyframe or
1371             self.active_handle or
1372             self.active_timebead or
1373             self.active_frame)):
1374             # override default translate()
1375             if not self.drag:
1376                 # start drag
1377                 if self.active_frame:
1378                     insert_keyframe(self, context, self.active_frame)
1379                     self.active_keyframe = self.active_frame
1380                     self.active_frame = False
1381                 self.keyframes_ori, self.handles_ori = \
1382                     get_original_animation_data(context, self.keyframes)
1383                 self.drag_mouse_ori = mathutils.Vector([event.mouse_region_x,
1384                     event.mouse_region_y])
1385                 self.drag = True
1386                 self.lock = False
1387             else:
1388                 # stop drag
1389                 self.drag = False
1390                 self.lock = True
1391                 context.window_manager.motion_trail.force_update = True
1392         elif event.type == self.transform_key and event.value == 'PRESS':
1393             # call default translate()
1394             if bpy.ops.transform.translate.poll():
1395                 bpy.ops.transform.translate('INVOKE_DEFAULT')
1396         elif (event.type == 'ESC' and self.drag and event.value == 'PRESS') \
1397         or (event.type == 'RIGHTMOUSE' and self.drag and event.value == \
1398         'PRESS'):
1399             # cancel drag
1400             self.drag = False
1401             self.lock = True
1402             context.window_manager.motion_trail.force_update = True
1403             self.active_keyframe, self.active_timebead = cancel_drag(context,
1404                 self.active_keyframe, self.active_handle,
1405                 self.active_timebead, self.keyframes_ori, self.handles_ori,
1406                 self.edit_bones)
1407         elif event.type == 'MOUSEMOVE' and self.drag:
1408             # drag
1409             self.active_keyframe, self.active_timebead, self.keyframes_ori = \
1410                 drag(context, event, self.drag_mouse_ori,
1411                 self.active_keyframe, self.active_handle,
1412                 self.active_timebead, self.keyframes_ori, self.handles_ori,
1413                 self.edit_bones)
1414         elif event.type == select + 'MOUSE' and event.value == 'PRESS' and \
1415         not self.drag and not event.shift and not event.alt and not \
1416         event.ctrl:
1417             # select
1418             treshold = 10
1419             clicked = mathutils.Vector([event.mouse_region_x,
1420                 event.mouse_region_y])
1421             self.active_keyframe = False
1422             self.active_handle = False
1423             self.active_timebead = False
1424             self.active_frame = False
1425             context.window_manager.motion_trail.force_update = True
1426             context.window_manager.motion_trail.handle_type_enabled = True
1427             found = False
1428
1429             if context.window_manager.motion_trail.path_before == 0:
1430                 frame_min = context.scene.frame_start
1431             else:
1432                 frame_min = max(context.scene.frame_start,
1433                     context.scene.frame_current - \
1434                     context.window_manager.motion_trail.path_before)
1435             if context.window_manager.motion_trail.path_after == 0:
1436                 frame_max = context.scene.frame_end
1437             else:
1438                 frame_max = min(context.scene.frame_end,
1439                     context.scene.frame_current + \
1440                     context.window_manager.motion_trail.path_after)
1441
1442             for objectname, values in self.click.items():
1443                 if found:
1444                     break
1445                 for frame, type, coord, action_ob, child in values:
1446                     if frame < frame_min or frame > frame_max:
1447                         continue
1448                     if (coord - clicked).length <= treshold:
1449                         found = True
1450                         if type == "keyframe":
1451                             self.active_keyframe = [objectname, frame, frame,
1452                                 action_ob, child]
1453                         elif type == "handle_left":
1454                             self.active_handle = [objectname, frame, "left",
1455                                 action_ob, child]
1456                         elif type == "handle_right":
1457                             self.active_handle = [objectname, frame, "right",
1458                                 action_ob, child]
1459                         elif type == "timebead":
1460                             self.active_timebead = [objectname, frame, frame,
1461                                 action_ob, child]
1462                         elif type == "frame":
1463                             self.active_frame = [objectname, frame, frame,
1464                                 action_ob, child]
1465                         break
1466             if not found:
1467                 context.window_manager.motion_trail.handle_type_enabled = False
1468                 # no motion trail selections, so pass on to normal select()
1469                 if bpy.ops.view3d.select.poll():
1470                     bpy.ops.view3d.select('INVOKE_DEFAULT')
1471             else:
1472                 handle_type = get_handle_type(self.active_keyframe,
1473                     self.active_handle)
1474                 if handle_type:
1475                     context.window_manager.motion_trail.handle_type_old = \
1476                         handle_type
1477                     context.window_manager.motion_trail.handle_type = \
1478                         handle_type
1479                 else:
1480                     context.window_manager.motion_trail.handle_type_enabled = \
1481                         False
1482         elif event.type == 'LEFTMOUSE' and event.value == 'PRESS' and \
1483         self.drag:
1484             # stop drag
1485             self.drag = False
1486             self.lock = True
1487             context.window_manager.motion_trail.force_update = True
1488         elif event.type == 'LEFTMOUSE' and event.value == 'PRESS' and not\
1489         event.alt and not event.ctrl and not event.shift:
1490             if eval("bpy.ops."+self.left_action+".poll()"):
1491                 eval("bpy.ops."+self.left_action+"('INVOKE_DEFAULT')")
1492
1493         if context.area: # not available if other window-type is fullscreen
1494             context.area.tag_redraw()
1495
1496         return {'PASS_THROUGH'}
1497
1498     def invoke(self, context, event):
1499         if context.area.type == 'VIEW_3D':
1500             # get clashing keymap items
1501             select = context.user_preferences.inputs.select_mouse
1502             kms = [bpy.context.window_manager.keyconfigs.active.\
1503                 keymaps['3D View'], bpy.context.window_manager.keyconfigs.\
1504                 active.keymaps['Object Mode']]
1505             kmis = []
1506             self.left_action = None
1507             self.right_action = None
1508             for km in kms:
1509                 for kmi in km.keymap_items:
1510                     if kmi.idname == "transform.translate" and \
1511                     kmi.map_type == 'KEYBOARD' and not \
1512                     kmi.properties.texture_space:
1513                         kmis.append(kmi)
1514                         self.transform_key = kmi.type
1515                     elif (kmi.type == 'ACTIONMOUSE' and select == 'RIGHT') \
1516                     and not kmi.alt and not kmi.any and not kmi.ctrl \
1517                     and not kmi.shift:
1518                         kmis.append(kmi)
1519                         self.left_action = kmi.idname
1520                     elif kmi.type == 'SELECTMOUSE' and not kmi.alt and not \
1521                     kmi.any and not kmi.ctrl and not kmi.shift:
1522                         kmis.append(kmi)
1523                         if select == 'RIGHT':
1524                             self.right_action = kmi.idname
1525                         else:
1526                             self.left_action = kmi.idname
1527                     elif kmi.type == 'LEFTMOUSE' and not kmi.alt and not \
1528                     kmi.any and not kmi.ctrl and not kmi.shift:
1529                         kmis.append(kmi)
1530                         self.left_action = kmi.idname
1531
1532             if not context.window_manager.motion_trail.enabled:
1533                 # enable
1534                 self.active_keyframe = False
1535                 self.active_handle = False
1536                 self.active_timebead = False
1537                 self.active_frame = False
1538                 self.click = {}
1539                 self.drag = False
1540                 self.lock = True
1541                 self.perspective = context.region_data.perspective_matrix
1542                 self.displayed = []
1543                 context.window_manager.motion_trail.force_update = True
1544                 context.window_manager.motion_trail.handle_type_enabled = False
1545                 self.cached = {"path":{}, "keyframes":{},
1546                     "timebeads_timing":{}, "timebeads_speed":{}}
1547
1548                 for kmi in kmis:
1549                     kmi.active = False
1550
1551                 MotionTrailOperator.handle_add(self, context)
1552                 context.window_manager.motion_trail.enabled = True
1553
1554                 if context.area:
1555                     context.area.tag_redraw()
1556
1557                 context.window_manager.modal_handler_add(self)
1558                 return {'RUNNING_MODAL'}
1559
1560             else:
1561                 # disable
1562                 for kmi in kmis:
1563                     kmi.active = True
1564                 MotionTrailOperator.handle_remove()
1565                 context.window_manager.motion_trail.enabled = False
1566
1567                 if context.area:
1568                     context.area.tag_redraw()
1569
1570                 return {'FINISHED'}
1571
1572         else:
1573             self.report({'WARNING'}, "View3D not found, cannot run operator")
1574             return {'CANCELLED'}
1575
1576
1577 class MotionTrailPanel(bpy.types.Panel):
1578     bl_space_type = 'VIEW_3D'
1579     bl_region_type = 'TOOLS'
1580     bl_label = "Motion Trail"
1581     bl_options = {'DEFAULT_CLOSED'}
1582
1583     @classmethod
1584     def poll(cls, context):
1585         if context.active_object is None:
1586             return False
1587         return context.active_object.mode in ('OBJECT', 'POSE')
1588
1589     def draw(self, context):
1590         col = self.layout.column()
1591         if not context.window_manager.motion_trail.enabled:
1592             col.operator("view3d.motion_trail", text="Enable motion trail")
1593         else:
1594             col.operator("view3d.motion_trail", text="Disable motion trail")
1595
1596         box = self.layout.box()
1597         box.prop(context.window_manager.motion_trail, "mode")
1598         #box.prop(context.window_manager.motion_trail, "calculate")
1599         if context.window_manager.motion_trail.mode == 'timing':
1600             box.prop(context.window_manager.motion_trail, "timebeads")
1601
1602         box = self.layout.box()
1603         col = box.column()
1604         row = col.row()
1605
1606         if context.window_manager.motion_trail.path_display:
1607             row.prop(context.window_manager.motion_trail, "path_display",
1608                 icon="DOWNARROW_HLT", text="", emboss=False)
1609         else:
1610             row.prop(context.window_manager.motion_trail, "path_display",
1611                 icon="RIGHTARROW", text="", emboss=False)
1612
1613         row.label("Path options")
1614
1615         if context.window_manager.motion_trail.path_display:
1616             col.prop(context.window_manager.motion_trail, "path_style",
1617                 text="Style")
1618             grouped = col.column(align=True)
1619             grouped.prop(context.window_manager.motion_trail, "path_width",
1620                 text="Width")
1621             grouped.prop(context.window_manager.motion_trail,
1622                 "path_transparency", text="Transparency")
1623             grouped.prop(context.window_manager.motion_trail,
1624                 "path_resolution")
1625             row = grouped.row(align=True)
1626             row.prop(context.window_manager.motion_trail, "path_before")
1627             row.prop(context.window_manager.motion_trail, "path_after")
1628             col = col.column(align=True)
1629             col.prop(context.window_manager.motion_trail, "keyframe_numbers")
1630             col.prop(context.window_manager.motion_trail, "frame_display")
1631
1632         if context.window_manager.motion_trail.mode == 'location':
1633             box = self.layout.box()
1634             col = box.column(align=True)
1635             col.prop(context.window_manager.motion_trail, "handle_display",
1636                 text="Handles")
1637             if context.window_manager.motion_trail.handle_display:
1638                 row = col.row()
1639                 row.enabled = context.window_manager.motion_trail.\
1640                     handle_type_enabled
1641                 row.prop(context.window_manager.motion_trail, "handle_type")
1642
1643
1644 class MotionTrailProps(bpy.types.PropertyGroup):
1645     def internal_update(self, context):
1646         context.window_manager.motion_trail.force_update = True
1647         if context.area:
1648             context.area.tag_redraw()
1649
1650     # internal use
1651     enabled = bpy.props.BoolProperty(default=False)
1652
1653     force_update = bpy.props.BoolProperty(name="internal use",
1654         description="Force calc_callback to fully execute",
1655         default=False)
1656
1657     handle_type_enabled = bpy.props.BoolProperty(default=False)
1658     handle_type_frame = bpy.props.FloatProperty()
1659     handle_type_side = bpy.props.StringProperty()
1660     handle_type_action_ob = bpy.props.StringProperty()
1661     handle_type_child = bpy.props.StringProperty()
1662     handle_type_old = bpy.props.EnumProperty(items=(
1663         ("AUTO", "", ""),
1664         ("AUTO_CLAMPED", "", ""),
1665         ("VECTOR", "", ""),
1666         ("ALIGNED", "", ""),
1667         ("FREE", "", "")),
1668         default='AUTO')
1669
1670     # visible in user interface
1671     calculate = bpy.props.EnumProperty(name="Calculate", items=(
1672         ("fast", "Fast", "Recommended setting, change if the "\
1673          "motion path is positioned incorrectly"),
1674         ("full", "Full", "Takes parenting and modifiers into account, "\
1675          "but can be very slow on complicated scenes")),
1676         description="Calculation method for determining locations",
1677         default='full',
1678         update=internal_update)
1679
1680     frame_display = bpy.props.BoolProperty(name="Frames",
1681         description="Display frames, \n test",
1682         default=True,
1683         update=internal_update)
1684
1685     handle_display = bpy.props.BoolProperty(name="Display",
1686         description="Display handles",
1687         default=True,
1688         update=internal_update)
1689
1690     handle_type = bpy.props.EnumProperty(name="Type", items=(
1691         ("AUTO", "Automatic", ""),
1692         ("AUTO_CLAMPED", "Auto Clamped", ""),
1693         ("VECTOR", "Vector", ""),
1694         ("ALIGNED", "Aligned", ""),
1695         ("FREE", "Free", "")),
1696         description="Set handle type for the selected handle",
1697         default='AUTO',
1698         update=set_handle_type)
1699
1700     keyframe_numbers = bpy.props.BoolProperty(name="Keyframe numbers",
1701         description="Display keyframe numbers",
1702         default=False,
1703         update=internal_update)
1704
1705     mode = bpy.props.EnumProperty(name="Mode", items=(
1706         ("location", "Location", "Change path that is followed"),
1707         ("speed", "Speed", "Change speed between keyframes"),
1708         ("timing", "Timing", "Change position of keyframes on timeline")),
1709         description="Enable editing of certain properties in the 3d-view",
1710         default='location',
1711         update=internal_update)
1712
1713     path_after = bpy.props.IntProperty(name="After",
1714         description="Number of frames to show after the current frame, "\
1715             "0 = display all",
1716         default=50,
1717         min=0,
1718         update=internal_update)
1719
1720     path_before = bpy.props.IntProperty(name="Before",
1721         description="Number of frames to show before the current frame, "\
1722             "0 = display all",
1723         default=50,
1724         min=0,
1725         update=internal_update)
1726
1727     path_display = bpy.props.BoolProperty(name="Path options",
1728         description="Display path options",
1729         default=True)
1730
1731     path_resolution = bpy.props.IntProperty(name="Resolution",
1732         description="10 is smoothest, but could be "\
1733         "slow when adjusting keyframes, handles or timebeads",
1734         default=10,
1735         min=1,
1736         max=10,
1737         update=internal_update)
1738
1739     path_style = bpy.props.EnumProperty(name="Path style", items=(
1740         ("acceleration", "Acceleration", "Gradient based on relative acceleration"),
1741         ("simple", "Simple", "Black line"),
1742         ("speed", "Speed", "Gradient based on relative speed")),
1743         description="Information conveyed by path color",
1744         default='simple',
1745         update=internal_update)
1746
1747     path_transparency = bpy.props.IntProperty(name="Path transparency",
1748         description="Determines visibility of path",
1749         default=0,
1750         min=0,
1751         max=100,
1752         subtype='PERCENTAGE',
1753         update=internal_update)
1754
1755     path_width = bpy.props.IntProperty(name="Path width",
1756         description="Width in pixels",
1757         default=1,
1758         min=1,
1759         soft_max=5,
1760         update=internal_update)
1761
1762     timebeads = bpy.props.IntProperty(name="Time beads",
1763         description="Number of time beads to display per segment",
1764         default=5,
1765         min=1,
1766         soft_max = 10,
1767         update=internal_update)
1768
1769
1770 def register():
1771     bpy.utils.register_module(__name__)
1772     bpy.types.WindowManager.motion_trail = bpy.props.PointerProperty(
1773         type=MotionTrailProps)
1774
1775
1776 def unregister():
1777     MotionTrailOperator.handle_remove()
1778     bpy.utils.unregister_module(__name__)
1779     del bpy.types.WindowManager.motion_trail
1780
1781
1782 if __name__ == "__main__":
1783     register()