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