# ##### BEGIN GPL LICENSE BLOCK #####
#
#  This program is free software; you can redistribute it and/or
#  modify it under the terms of the GNU General Public License
#  as published by the Free Software Foundation; either version 2
#  of the License, or (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software Foundation,
#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####

__bpydoc__ = """\
TimeOffset enables setting an animation frame offset on linked groups/objects
"""


bl_info = {
	"name": "TimeOffset",
	"author": "Gert De Roost",
	"version": (0, 5, 0),
	"blender": (2, 6, 5),
	"location": "View3D > Tools",
	"description": "Set frame offset on linked groups/objects",
	"warning": "",
	"wiki_url": "",
	"tracker_url": "",
	"category": "Object"}

if "bpy" in locals():
    import imp


import bpy
from random import random
from mathutils import *


activated = 0
end = 0
linkeddir = {}
groupdir = {}
savobjs = []
oldframe = []
root = None
rootname = ""
obchanged = 0
savbyname = 1
bynamechanged = 0
datasavdict = {}
oldobs = []
oldactob = None

bpy.types.Scene.Bottomlimit = bpy.props.IntProperty(name = "Lower", description = "Lower limit for randomizing offset values", default = 0)
bpy.types.Scene.Toplimit = bpy.props.IntProperty(name = "Upper", description = "Upper limit for randomizing offset values", default = 10)
bpy.types.Scene.Setvalue = bpy.props.IntProperty(name = "Value", description = "Value to set offset of selected objects to", default = 0)
bpy.types.Scene.Startvalue = bpy.props.IntProperty(name = "Start", description = "Value to start interval-offsets from", default = 0)
bpy.types.Scene.Setinterval = bpy.props.IntProperty(name = "Interval", description = "Set interval between object offsets", default = 1)
bpy.types.Scene.byName = bpy.props.BoolProperty(name = "by Name", description = "Copies found by name or by instance-similarity", default = 1)
bpy.types.Scene.Onlyrender = bpy.props.BoolProperty(name = "Only Render", description = "Only apply timeoffset when rendering", default = 0)



class Activate(bpy.types.Operator):
	bl_idname = "timeoffset.activate"
	bl_label = "Activate"
	bl_description = "Activate addon"
	bl_options = {"REGISTER", "UNDO"}
	
	def invoke(self, context, event):
		
		global activated, region, end, seloblist
		
		scn = bpy.context.scene
		region = bpy.context.region
		
		end = 0
		
		seloblist = bpy.context.selected_objects
		
		context.window_manager.modal_handler_add(self)
		if eval(str(bpy.app.build_revision)[2:7]) >= 53207:
			self._handle = bpy.types.SpaceView3D.draw_handler_add(update, (), 'WINDOW', 'POST_PIXEL')
		else:
			self._handle = context.region.callback_add(update, (), 'POST_PIXEL')
		bpy.app.handlers.scene_update_post.append(setoffsets)
		bpy.app.handlers.render_pre.append(removelinks)
		bpy.app.handlers.render_post.append(restorelinks)
		bpy.app.handlers.load_post.append(loaddata)
		
		checkobjs(1)
		loaddata()
		activated = 1
		
		return {'RUNNING_MODAL'}


	def modal(self, context, event):

		global activated

		if end:
			if eval(str(bpy.app.build_revision)[2:7]) >= 53207:
				bpy.types.SpaceView3D.draw_handler_remove(self._handle, "WINDOW")
			else:
				context.region.callback_remove(self._handle)
			bpy.app.handlers.scene_update_post.remove(setoffsets)
			bpy.app.handlers.render_pre.remove(removelinks)
			bpy.app.handlers.render_post.remove(restorelinks)
			bpy.app.handlers.load_post.remove(loaddata)
			activated = 0
			return {'FINISHED'}
		return {"PASS_THROUGH"}
		
			


class Randomize(bpy.types.Operator):
	bl_idname = "timeoffset.randomize"
	bl_label = "Randomize"
	bl_description = "Randomize timeoffset values"
	bl_options = {"REGISTER", "UNDO"}
	
	def invoke(self, context, event):
		
		global rootname
		
		scn = bpy.context.scene
		
		for i, ob in enumerate(linkeddir[root.name]):
			rootnm = root.name
			rootnm = rootnm.replace(".", "_")
			rootnm = rootnm.replace(" ", "_")
			exec('scn.' + rootnm + '_Frame' + str(i) + " = " + str(scn.Bottomlimit + random() * (scn.Toplimit - scn.Bottomlimit)))
					
		return {'FINISHED'}
		


class SetValue(bpy.types.Operator):
	bl_idname = "timeoffset.setvalue"
	bl_label = "Set Value"
	bl_description = "Set selected objects to this offset value"
	bl_options = {"REGISTER", "UNDO"}
	
	def invoke(self, context, event):
		
		scn = bpy.context.scene
		
		def checkparent(obj):
			scn = bpy.context.scene
			if obj:
				for lst in linkeddir.values():
					if obj in lst:
						rootnm = root.name
						rootnm = rootnm.replace(".", "_")
						rootnm = rootnm.replace(" ", "_")
						exec('scn.' + rootnm + '_Frame' + str(lst.index(obj)) + " = scn.Setvalue")
					else:
						checkparent(obj.parent)
			else:
				return
				
		for obj in seloblist:
		 	checkparent(obj)
					
		return {'FINISHED'}
		
		


class SetInterval(bpy.types.Operator):
	bl_idname = "timeoffset.setinterval"
	bl_label = "Set Interval"
	bl_description = "Set selected objects to these interval values"
	bl_options = {"REGISTER", "UNDO"}
	
	def invoke(self, context, event):
		
		global value
		
		scn = bpy.context.scene
		
		def checkparent(obj):
			global value
			scn = bpy.context.scene
			if obj:
				for lst in linkeddir.values():
					if obj in lst:
						rootnm = root.name
						rootnm = rootnm.replace(".", "_")
						rootnm = rootnm.replace(" ", "_")
						exec('scn.' + rootnm + '_Frame' + str(lst.index(obj)) + " = value")
						value += scn.Setinterval
						return
					else:
						checkparent(obj.parent)
			else:
				return
				
		value = scn.Startvalue
		for obj in seloblist:
		 	checkparent(obj)
					
		return {'FINISHED'}
		
		


class Apply(bpy.types.Operator):
	bl_idname = "timeoffset.apply"
	bl_label = "Apply timeoffset"
	bl_description = "Apply all addon timeoffset settings"
	bl_options = {"REGISTER", "UNDO"}
	
	def invoke(self, context, event):
		
		global root, end
		
		scn = bpy.context.scene
		
		def applychildren(parent, refob):
			global datasavdict
			scn = bpy.context.scene
			exchil = list(parent.children)
			remob = refob
			for j, ob in enumerate(exchil):
				ref2ob = remob.children[j]
				if ob.animation_data:			
					ob.animation_data.action = ref2ob.animation_data.action.copy()
					if ob.data:
						mesh = bpy.data.meshes.new(ob.data.name)
						mesh = ob.data.copy()
						ob.data = mesh
						fshift = eval("scn." + rootnm + "_Frame" + str(i))
						for fc in ob.animation_data.action.fcurves:
							for k in fc.keyframe_points:
								k.co -= Vector((fshift, 0))
								k.handle_left -= Vector((fshift, 0))
								k.handle_right -= Vector((fshift, 0))					
				applychildren(ob, ref2ob)
			if parent.data:
				if parent.data.animation_data:
					if parent.data.animation_data.action:
						ref2ob = remob.data
						parent.data.animation_data.action = ref2ob.animation_data.action.copy()
						mesh = parent.data.copy()
						parent.data = mesh
						fshift = eval("scn." + rootnm + "_Frame" + str(i))
						for fc in parent.data.animation_data.action.fcurves:
							for k in fc.keyframe_points:
								k.co -= Vector((fshift, 0))
								k.handle_left -= Vector((fshift, 0))
								k.handle_right -= Vector((fshift, 0))					
					
		for lst in linkeddir.values():
			for i, empty in enumerate(lst):
				rootnm = scn.Root
				rtob = bpy.data.objects.get(rootnm)
				applychildren(empty, rtob)
				
		end = 1
		

		return {'FINISHED'}
		


class TimeOffset(bpy.types.Operator):
	bl_idname = "object.timeoffset"
	bl_label = "TimeOffset"
	bl_description = "Set frame offset on linked groups/objects"
	bl_options = {"REGISTER", "UNDO"}
	
	def invoke(self, context, event):
		
		global rootname
		
		scn = bpy.context.scene
		
		rootname = scn.Root
		do_timeoffset()
			
		return {'FINISHED'}




def panel_func(self, context):
	
	global oldframe, groupdir, savbyname, bynamechanged
	
	scn = bpy.context.scene
	
	self.layout.label(text="TimeOffset:")
	if not(activated):
		self.layout.operator("timeoffset.activate", text="Activate")
	else:
		self.layout.prop(scn, "Root")
		self.layout.operator("object.timeoffset", text="Set root object")
		self.layout.prop(scn, "byName")
		if scn.byName != savbyname:
			bynamechanged = 1
			bpy.context.region.tag_redraw()
			savbyname = scn.byName
		if root:
			row = self.layout.row()
			row.operator("timeoffset.setvalue", text="Selected to:")
			row.prop(scn, "Setvalue")
			self.layout.operator("timeoffset.setinterval", text="Selected to:")
			row = self.layout.row()
			row.prop(scn, "Startvalue")
			row.prop(scn, "Setinterval")
			self.layout.operator("timeoffset.randomize", text="Randomize")
			row = self.layout.row()
			row.prop(scn, "Bottomlimit")
			row.prop(scn, "Toplimit")
			self.layout.prop(scn, "Onlyrender")
			self.layout.operator("timeoffset.apply", text="Apply")
			for i, ob in enumerate(linkeddir[root.name]):
				self.layout.label(text=ob.name)
				rootnm = root.name
				rootnm = rootnm.replace(".", "_")
				rootnm = rootnm.replace(" ", "_")
				self.layout.prop(scn, rootnm + "_Frame" + str(i))
				frame = eval("scn." + rootnm + "_Frame" + str(i))
				if oldframe[i] != frame:
					savedata()

def register():
	bpy.utils.register_module(__name__)
	bpy.types.VIEW3D_PT_tools_meshedit.append(panel_func)
	bpy.types.VIEW3D_PT_tools_objectmode.append(panel_func)


def unregister():
	bpy.utils.unregister_module(__name__)
	bpy.types.VIEW3D_PT_tools_meshedit.remove(panel_func)
	bpy.types.VIEW3D_PT_tools_objectmode.append(panel_func)


if __name__ == "__main__":
	register()






def do_timeoffset():

	global linkeddir, root, rootname, oldframe
	
	scn = bpy.context.scene
	
	scn.Root = rootname
	root = bpy.data.objects.get(rootname)
	linkeddir[root.name] = []
	sellist = []
	for ob in scn.objects:
		if ob.select:
			sellist.append(ob)
	savact = bpy.context.active_object
	bpy.ops.object.select_all(action="DESELECT")
	root.select = 1
	scn.objects.active = root
	bpy.ops.object.select_linked(extend=False, type='OBDATA')
	root.select = 0
	lst = []
	for ob in scn.objects:
		if ob == root:
			continue
		if scn.byName:
			if ob.name.count(root.name):
				lst.append(ob)
		else:
			cnt = 0
			for ch1 in root.children:
				for ch2 in ob.children:
					if ch1.data == ch2.data:
						cnt += 1
			if cnt != len(root.children) or len(ob.children) == 0:
				continue
			lst.append(ob)
	oldframe = []
	lst.reverse()
	for i, ob in enumerate(lst):
		rootnm = root.name
		rootnm = rootnm.replace(".", "_")
		rootnm = rootnm.replace(" ", "_")
		linkeddir[root.name].append(ob)
		exec('bpy.types.Scene.' + rootnm + '_Frame' + str(i) + ' = bpy.props.IntProperty(name = "Frame offset", description = "Frame offset", default = 0)')
		oldframe.append(0)
	bpy.ops.object.select_all(action="DESELECT")
	for ob in sellist:
		ob.select = 1
	scn.objects.active = savact
	
	savedata()




def setoffsets(dummy=0):
	
	global oldobs, oldactob, seloblist
	
	scn = bpy.context.scene
	
	actob = bpy.context.active_object
	obs = bpy.context.selected_objects
	if obs != oldobs:
		for ob in obs:
			if not(ob in seloblist):
				seloblist.append(ob)
		for ob in seloblist:
			if not(ob in obs):
				seloblist.pop(seloblist.index(ob))
		oldobs = obs
		oldactob = actob
	elif actob != oldactob:
		seloblist.pop(seloblist.index(actob))
		seloblist.append(actob)
		oldactob = actob
	
	if scn.Onlyrender and dummy != 2:
		return
	
	checkobjs()
	rtval = 0
	keys = list(linkeddir.keys())
	
	def setchildren(parent, obrec):
		scn = bpy.context.scene
		remrec = obrec
		exchil = list(parent.children) + [parent.data]
		for j, ob in enumerate(exchil):
			if ob == exchil[-1]:
				obrec = remrec.data
			else:
				obrec = remrec.children[j]
			if ob == None:
				continue
			if ob.animation_data == None:
				if ob != parent.data:
					setchildren(ob, obrec)
				continue
			if ob.animation_data.action == None:
				continue
			for fc in ob.animation_data.action.fcurves:
				for fcrt in obrec.animation_data.action.fcurves:
					if fc.data_path == fcrt.data_path and fc.array_index == fcrt.array_index:
						frm = eval("scn.frame_current + scn." + rt.name + "_Frame" + str(i))
						rtval = fcrt.evaluate(frm)
						try:
							exec("ob." + fc.data_path + "[" + str(fc.array_index) + "] = " + str(rtval))
						except:	
							exec("ob." + fc.data_path + " = " + str(rtval))
			if ob != parent.data:
				setchildren(ob, obrec)

	for name in keys:
		rt = bpy.data.objects.get(name)
		lst = list(linkeddir[name])
		for i, empty in enumerate(lst):
			obrec = rt
			setchildren(empty, obrec)
					
					
def removelinks(dummy):
	
	scn = bpy.context.scene

	setoffsets(2)

	def remchildren(parent):
		global datasavdict
		exchil = list(parent.children)
		for ob in exchil:
			if ob.animation_data:
				datasavdict[ob] = ob.animation_data.action.copy()
				ob.animation_data.action = None
			remchildren(ob)
		if parent.data:
			if parent.data.animation_data:
				if parent.data.animation_data.action:
					datasavdict[parent.data] = parent.data.animation_data.action.copy()
					parent.data.animation_data.action = None
				
	for lst in linkeddir.values():
		for empty in lst:
			remchildren(empty)
		
		
def restorelinks(dummy):
	
	global datasavdict
	
	def restchildren(parent, obrec):
		remrec = obrec
		for i, ob in enumerate(parent.children):
			obrec = remrec.children[i]
			if ob.animation_data:
				ob.animation_data.action = datasavdict[ob]
			restchildren(ob, obrec)
		if parent.data:
			if parent.data.animation_data:
				parent.data.animation_data.action = datasavdict[parent.data]
				
	for name in linkeddir.keys():
		rt = bpy.data.objects.get(name)
		lst = linkeddir[name]
		for empty in lst:
			obrec = rt
			restchildren(empty, obrec)
			
	datasavdict = {}
	
	
	
def checkobjs(stage=0):
	
	global savobjs, groupdir, obchanged
	
	scn = bpy.context.scene
	
	lst = []
	for ob in scn.objects:
		lst.append(ob.name)

	if lst[:] != savobjs:
		itemlist = []
		for ob in scn.objects:
			if ob.type == "EMPTY":
				itemlist.append((ob.name, ob.name, "Choose root empty"))
		if rootname:
			bpy.types.Scene.Root = bpy.props.EnumProperty(
					items = itemlist,
					default = rootname,
					name = "Root:", 
					description = "Object to get animation settings from")
			obchanged = 1
			region.tag_redraw()
		else:
			bpy.types.Scene.Root = bpy.props.EnumProperty(
					items = itemlist,
					name = "Root:", 
					description = "Object to get animation settings from")
			
	savobjs = lst[:]
	
	
	
def update():
	
	global obchanged, bynamechanged, rootname
	
	if rootname and obchanged:
		obchanged = 0
		do_timeoffset()
		
	if bynamechanged:
		bynamechanged = 0
		scn = bpy.context.scene
		rootname = scn.Root
		do_timeoffset()
		





def savedata():
	
	scn = bpy.context.scene
	
	try:
		txt = bpy.data.texts["timeoffset_save"]
		txt.clear()
	except:
		#create new text
		txt = bpy.data.texts.new("timeoffset_save")

	for key in linkeddir.keys():		
		rootname = key
		oblist = []
		framelist = []
		for i, ob in enumerate(linkeddir[rootname]):
			oblist.append(ob.name)
			rootnm = rootname
			rootnm = rootnm.replace(".", "_")
			rootnm = rootnm.replace(" ", "_")
			frm = eval('scn.' + rootnm + '_Frame' + str(i))
			framelist.append(frm)
    
		#write rootname to text object
		txt.write("'"+rootname+"'")
		txt.write("\n")
		
		#write oblist to text object
		txt.write("[")
		for i in range(len(oblist)):
			if i != 0:
				txt.write(", ")
			txt.write("'"+ str(oblist[i])+"'")
		txt.write("]\n")
		
		#write framelist to text object
		txt.write("[")
		for i in range(len(framelist)):
			if i != 0:
				txt.write(", ")
			txt.write(str(framelist[i]))
		txt.write("]\n")


def loaddata():
	
	global linkeddir, rootname, oldframe
	
	scn = bpy.context.scene
	
	linkeddir = {}
	txt = bpy.data.texts.get("timeoffset_save")
	if txt == None:
		return
	#load the saved data
	evallist = txt.lines
	oldframe = []
	for i in range(0, len(evallist) - 1, 3):
		rootname = eval(evallist[i].body)
		linkeddir[rootname] = []
		root = bpy.data.objects.get(rootname)
		oblist = eval(evallist[i + 1].body)
		framelist = eval(evallist[i + 2].body)
		for j, name in enumerate(oblist):
			ob = bpy.data.objects.get(name)
			linkeddir[rootname].append(ob)
			exec('bpy.types.Scene.' + rootname + '_Frame' + str(j) + ' = bpy.props.IntProperty(name = "Frame offset", description = "Frame offset", default = framelist[' + str(j) + '])')
			oldframe.append(0)
	checkobjs(1)
	if rootname:
		do_timeoffset()