#
# Copyright 2006 Alex Fraser <alex@phatcore.com>
#
# 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 St, Fifth Floor, Boston, MA  02110-1301  USA
#
# TrackAndTransform.py:
# Makes an object look towards another, while inheriting the
# transform of a third. Useful for non-spherical eyes that
# need to track an object.
#
# Setting Up:
#     Select two objects, then run this script (Alt-P). The
#     first object will track the second. A number of controls
#     will be created: One for moving and rotating (looks like
#     a cross between a rotate and a translate manipulator),
#     and one for scaling (looks like a disjoint cube). The
#     third indicates which direction the tracking object is
#     pointing. Select this to change the tracking influence.
#
#     You may rename any of the widgets; just don't delete the TRK_
#     properties.
#
#     The existing transform and overall hierarchy are preserved.
#
# Tearing Down:
#     Select the tracking object or any of its related controls
#     and run this script (Alt-P).
#
#     The existing transform and overall hierarchy are preserved.
#
# Known Issues:
#     1) Only the negative Y-axis can be used to track.
#     2) The tracking object (arrow) must be on a visible layer.
#
# The only function we really need is TrackAndTransform. The
# rest just makes it more user-friendly!
#
 
import Blender

GeneMap = {}
def GetChildren(parent):
	"""Find the children of a given object."""
	# This seaches through all objects in the current scene! Results
	# are cached within this script instance to speed up successive
	# searches.
	if GeneMap.has_key(parent.name):
		return GeneMap[parent.name]
	children = []
	for ob in Blender.Object.Get():
		if ob.parent == parent:
			children.append(ob)
	GeneMap[parent.name] = children
	return children

def SearchUp(constituent, propName, propValue):
	"""Find an ancestor that contains the specified property."""
	if constituent == None:
		raise Exception, 'Error: ancestor %s=%s not found' % (propName, propValue)
	id = constituent.getProperty(propName)
	if id != None and id.data == propValue:
		return constituent
	else:
		return SearchUp(constituent.getParent(), propName, propValue)

def SearchDown(constituent, propName, propValue):
	"""Find a descendant that contains the specified property.
	
	Bredth-first search."""
	def _SearchDown(constituent, propName, propValue):
		for child in GetChildren(constituent):
			id = child.getProperty(propName)
			if id != None and id.data == propValue:
				return child
		for child in GetChildren(constituent):
			match = SearchDown(child, propName, propValue)
			if match != None:
				return match
		return None
	
	id = constituent.getProperty(propName)
	if id != None and id.data == propValue:
		return constituent
	match = _SearchDown(constituent, propName, propValue)
	if match != None:
		return match
	else:
		raise Exception, 'Error: descendant %s=%s not found' % (propName, propValue)

def TrackAndTransform(constituent):
	"""Orients the mesh to match the Tracker, while respecting the
	scale of the FinalSpace."""
	# Find constituents by searching for known properties.
	refSpace   = SearchUp(constituent, 'TRK_Role', 'ReferenceSpace')
	tracker    = SearchDown(refSpace, 'TRK_Role', 'Tracker')
	finalSpace = SearchDown(refSpace, 'TRK_Role', 'FinalSpace')
	object     = SearchDown(refSpace, 'TRK_Role', 'Object')
	
	# Transform the mesh to match the LOCAL rotation of the tracker.
	tmx = tracker.getMatrix('localspace').rotationPart().resize4x4()
	# Transform again, this time to match the FinalSpace object
	# (which will generally be used for scaling).
	fmx = finalSpace.getMatrix('worldspace')
	
	object.setMatrix(tmx *  fmx)

def Create3dBezierCurve(name, data):
	"""Creates a Bezier curve object.
	
	data:
		The top level is a list of curves.
		Each curve is a dictionary of Bezier data.
		Each dictionary contains the fields 'FlagU', 'FlagV', 'Triples'.
		Each Triple is a list of three points.
		Each point is a list floats.
	
	To create such data, see the SerialiseCurve script."""
	
	def TripleToBezTriple(t):
		"""Convert a list of points into a Bezier Triple"""
		return Blender.BezTriple.New(t[0] + t[1] + t[2])

	cd = Blender.Curve.New(name)
	cd.setFlag(data['Flag'])
	for curve in data['Curves']:
		t1 = curve['Triples'].pop(0)
		cn = cd.appendNurb(TripleToBezTriple(t1))
		for t in curve['Triples']:
			cn.append(TripleToBezTriple(t))
		cn.setFlagU(curve['FlagU'])
		cn.setFlagV(curve['FlagV'])
	cd.update()
	ob = Blender.Object.New('Curve', name)
	ob.link(cd)
	return ob

def CreateTransRotManip(name):
	"""Creates an object shaped like translation and scale
	manipulators."""
	# curveData created using the SerialiseCurve script.
	curveData = {
		"Flag": 1, "Curves": [
			{
				"FlagU": 0, "FlagV": 0, "Triples": [
					[ [ 0, -1, 0, ], [ 0, -1, 0, ], [ 0, -1, 0, ], ],
					[ [ 0, 1, 0, ], [ 0, 1, 0, ], [ 0, 1, 0, ], ],
				],
			},
			{
				"FlagU": 0, "FlagV": 0, "Triples": [
					[ [ 1, 0, 0, ], [ 1, 0, 0, ], [ 1, 0, 0, ], ],
					[ [ -1, 0, 0, ], [ -1, 0, 0, ], [ -1, 0, 0, ], ],
				],
			},
			{
				"FlagU": 0, "FlagV": 0, "Triples": [
					[ [ 0, 0, 1, ], [ 0, 0, 1, ], [ 0, 0, 1, ], ],
					[ [ 0, 0, -1, ], [ 0, 0, -1, ], [ 0, 0, -1, ], ],
				],
			},
			{
				"FlagU": 1, "FlagV": 0, "Triples": [
					[ [ -1, 0, -0.552125, ], [ -1, 0, 0, ], [ -1, 0, 0.552125, ], ],
					[ [ -0.552125, 0, 1, ], [ 0, 0, 1, ], [ 0.552125, 0, 1, ], ],
					[ [ 1, 0, 0.552125, ], [ 1, 0, 0, ], [ 1, 0, -0.552125, ], ],
					[ [ 0.552125, 0, -1, ], [ 0, 0, -1, ], [ -0.552125, 0, -1, ], ],
				],
			},
			{
				"FlagU": 0, "FlagV": 0, "Triples": [
					[ [ -0.944079, -0.525128, 0, ], [ -0.707047, -0.707047, 0, ], [ -0.470015, -0.888966, 0, ], ],
					[ [ -0.298795, -1.00001, 0, ], [ 0, -1, 0, ], [ 0.298802, -0.999992, 0, ], ],
					[ [ 0.526062, -0.888002, 0, ], [ 0.707047, -0.707004, 0, ], [ 0.888031, -0.526005, 0, ], ],
				],
			},
			{
				"FlagU": 0, "FlagV": 0, "Triples": [
					[ [ 0, -0.525128, -0.944079, ], [ 0, -0.707047, -0.707047, ], [ 0, -0.888966, -0.470015, ], ],
					[ [ 0, -1.00001, -0.298783, ], [ 0, -1, 0, ], [ 0, -0.999992, 0.298802, ], ],
					[ [ 0, -0.888002, 0.526063, ], [ 0, -0.707004, 0.707047, ], [ 0, -0.526005, 0.888031, ], ],
				],
			},
		],
	}
	return Create3dBezierCurve(name, curveData)

def CreateScaleManip(name):
	"""Creates an object shaped like a scale manipulator."""
	# curveData created using the SerialiseCurve script.
	curveData = {
		"Flag": 1, "Curves": [
			{
				"FlagU": 1, "FlagV": 0, "Triples": [
					[ [ 0.2, 0.2, 1, ], [ 0.2, 0.2, 1, ], [ 0.2, 0.2, 1, ], ],
					[ [ 0.2, -0.2, 1, ], [ 0.2, -0.2, 1, ], [ 0.2, -0.2, 1, ], ],
					[ [ -0.2, -0.2, 1, ], [ -0.2, -0.2, 1, ], [ -0.2, -0.2, 1, ], ],
					[ [ -0.2, 0.2, 1, ], [ -0.2, 0.2, 1, ], [ -0.2, 0.2, 1, ], ],
				],
			},
			{
				"FlagU": 1, "FlagV": 0, "Triples": [
					[ [ 0.2, 0.2, -1, ], [ 0.2, 0.2, -1, ], [ 0.2, 0.2, -1, ], ],
					[ [ 0.2, -0.2, -1, ], [ 0.2, -0.2, -1, ], [ 0.2, -0.2, -1, ], ],
					[ [ -0.2, -0.2, -1, ], [ -0.2, -0.2, -1, ], [ -0.2, -0.2, -1, ], ],
					[ [ -0.2, 0.2, -1, ], [ -0.2, 0.2, -1, ], [ -0.2, 0.2, -1, ], ],
				],
			},
			{
				"FlagU": 1, "FlagV": 0, "Triples": [
					[ [ 1, 0.2, -0.2, ], [ 1, 0.2, -0.2, ], [ 1, 0.2, -0.2, ], ],
					[ [ 1, -0.2, -0.2, ], [ 1, -0.2, -0.2, ], [ 1, -0.2, -0.2, ], ],
					[ [ 1, -0.2, 0.2, ], [ 1, -0.2, 0.2, ], [ 1, -0.2, 0.2, ], ],
					[ [ 1, 0.2, 0.2, ], [ 1, 0.2, 0.2, ], [ 1, 0.2, 0.2, ], ],
				],
			},
			{
				"FlagU": 1, "FlagV": 0, "Triples": [
					[ [ -1, 0.2, -0.2, ], [ -1, 0.2, -0.2, ], [ -1, 0.2, -0.2, ], ],
					[ [ -1, -0.2, -0.2, ], [ -1, -0.2, -0.2, ], [ -1, -0.2, -0.2, ], ],
					[ [ -1, -0.2, 0.2, ], [ -1, -0.2, 0.2, ], [ -1, -0.2, 0.2, ], ],
					[ [ -1, 0.2, 0.2, ], [ -1, 0.2, 0.2, ], [ -1, 0.2, 0.2, ], ],
				],
			},
			{
				"FlagU": 1, "FlagV": 0, "Triples": [
					[ [ -0.2, 1, -0.2, ], [ -0.2, 1, -0.2, ], [ -0.2, 1, -0.2, ], ],
					[ [ 0.2, 1, -0.2, ], [ 0.2, 1, -0.2, ], [ 0.2, 1, -0.2, ], ],
					[ [ 0.2, 1, 0.2, ], [ 0.2, 1, 0.2, ], [ 0.2, 1, 0.2, ], ],
					[ [ -0.2, 1, 0.2, ], [ -0.2, 1, 0.2, ], [ -0.2, 1, 0.2, ], ],
				],
			},
			{
				"FlagU": 1, "FlagV": 0, "Triples": [
					[ [ -0.2, -1, -0.2, ], [ -0.2, -1, -0.2, ], [ -0.2, -1, -0.2, ], ],
					[ [ 0.2, -1, -0.2, ], [ 0.2, -1, -0.2, ], [ 0.2, -1, -0.2, ], ],
					[ [ 0.2, -1, 0.2, ], [ 0.2, -1, 0.2, ], [ 0.2, -1, 0.2, ], ],
					[ [ -0.2, -1, 0.2, ], [ -0.2, -1, 0.2, ], [ -0.2, -1, 0.2, ], ],
				],
			},
		],
	}
	return Create3dBezierCurve(name, curveData)

def CreateArrow(name):
	"""Creates an object shaped like a 3D arrow. The arrow is 1 unit
	long and composed of Bezier curves."""
	# curveData created using the SerialiseCurve script.
	curveData = {
		"Flag": 1, "Curves": [
			{
				"FlagU": 0, "FlagV": 0, "Triples": [
					[ [ 0, -1, 0, ], [ 0, -1, 0, ], [ 0, -1, 0, ], ],
					[ [ 0, 0, 0, ], [ 0, 0, 0, ], [ 0, 0, 0, ], ],
				],
			},
			{
				"FlagU": 1, "FlagV": 0, "Triples": [
					[ [ 0.1, -0.9, 0.0552126, ], [ 0.1, -0.9, 0, ], [ 0.1, -0.9, -0.0552125, ], ],
					[ [ 0.0552125, -0.9, -0.1, ], [ 0, -0.9, -0.1, ], [ -0.0552125, -0.9, -0.1, ], ],
					[ [ -0.1, -0.9, -0.0552125, ], [ -0.1, -0.9, 0, ], [ -0.1, -0.9, 0.0552126, ], ],
					[ [ -0.0552125, -0.9, 0.1, ], [ 0, -0.9, 0.1, ], [ 0.0552125, -0.9, 0.1, ], ],
				],
			},
			{
				"FlagU": 1, "FlagV": 0, "Triples": [
					[ [ 0.2, -0.8, 0.110425, ], [ 0.2, -0.8, 0, ], [ 0.2, -0.8, -0.110425, ], ],
					[ [ 0.110425, -0.8, -0.2, ], [ 0, -0.8, -0.2, ], [ -0.110425, -0.8, -0.2, ], ],
					[ [ -0.2, -0.8, -0.110425, ], [ -0.2, -0.8, 0, ], [ -0.2, -0.8, 0.110425, ], ],
					[ [ -0.110425, -0.8, 0.2, ], [ 0, -0.8, 0.2, ], [ 0.110425, -0.8, 0.2, ], ],
				],
			},
		],
	}
	return Create3dBezierCurve(name, curveData)

def SetUpTracker(object, target):
	"""Sets up the hierarchy required for an object to track another,
	and then inherit the transform of a third.
	
	object: The object to set tracking for. The existing transformation and
	        hierarchy are preserved.
	target: The object being tracked."""
	
	# Create some empties! The matrices of these nodes will be
	# used to transform the nominated mesh.
	scene = Blender.Scene.getCurrent()
	refSpace = CreateTransRotManip(object.name + '.TRK_LocRot')
	scene.link(refSpace)
	refSpace.layers = scene.layers
	tracker = CreateArrow(object.name + '.TRK_Track')
	scene.link(tracker)
	tracker.layers = scene.layers
	finalSpace = CreateScaleManip(object.name + '.TRK_Scale')
	scene.link(finalSpace)
	finalSpace.layers = scene.layers
	
	# Add properties to identify each object within the hierarchy.
	object.addProperty('TRK_Role', 'Object')
	refSpace.addProperty('TRK_Role', 'ReferenceSpace')
	tracker.addProperty('TRK_Role', 'Tracker')
	finalSpace.addProperty('TRK_Role', 'FinalSpace')
	
	# Set up object hierarchy. The nominated mesh is made a child of
	# a new empty, so its transform is transferred to the empty.
	size = object.getSize('worldspace')
	object.setSize(1, 1, 1)
	mx = object.getMatrix('worldspace')
	parent = object.getParent()
	if parent != None:
		parent.makeParent([refSpace])
	refSpace.makeParent([object, tracker, finalSpace])
	refSpace.setMatrix(mx)
	finalSpace.setSize(size)
	
	# Set up tracking.
	sets = Blender.Constraint.Settings
	ttCon = tracker.constraints.append(Blender.Constraint.Type['TRACKTO'])
	ttCon.influence = 1
	ttCon.name = 'focus'
	ttCon[sets['TARGET']] = target
	ttCon[sets['TRACK']] = sets['TRACKNEGY']
	ttCon[sets['UP']] = sets['UPZ']
	
	# Lock channels that shouldn't be messed with.
	LOCK_T_X = 0x01 << 0
	LOCK_T_Y = 0x01 << 1
	LOCK_T_Z = 0x01 << 2
	LOCK_T = LOCK_T_X & LOCK_T_Y & LOCK_T_Z
	LOCK_R_X = 0x01 << 3
	LOCK_R_Y = 0x01 << 4
	LOCK_R_Z = 0x01 << 5
	LOCK_R = LOCK_R_X & LOCK_R_Y & LOCK_R_Z
	LOCK_S_X = 0x01 << 6
	LOCK_S_Y = 0x01 << 7
	LOCK_S_Z = 0x01 << 8
	LOCK_S = LOCK_S_X & LOCK_S_Y & LOCK_S_Z
	
	finalSpace.protectFlags = LOCK_T
	
	# Bind this script to empties. We want to know of any changes!
	textName = __script__.name
	object.addScriptLink(textName, 'Redraw')
	refSpace.addScriptLink(textName, 'Redraw')
	tracker.addScriptLink(textName, 'Redraw')
	finalSpace.addScriptLink(textName, 'Redraw')
	
	Blender.Scene.GetCurrent().update()
	Blender.Redraw()
	
	Blender.Draw.PupMenu(("Tracker setup complete." +
		" To change the tracking influence, select the arrow (%s)" +
		" and look at the Object buttons.")
		% tracker.name)

def TearDownTracker(constituent):
	"""Removes the tracking constraint. All related tracking controls are
	deleted. The overall hierarchy and transformation are preserved."""
	# Find constituents by searching for known properties.
	refSpace   = SearchUp(constituent, 'TRK_Role', 'ReferenceSpace')
	tracker    = SearchDown(refSpace, 'TRK_Role', 'Tracker')
	finalSpace = SearchDown(refSpace, 'TRK_Role', 'FinalSpace')
	object     = SearchDown(refSpace, 'TRK_Role', 'Object')
	
	object.clearScriptLinks([__script__.name])
	object.removeProperty(object.getProperty('TRK_Role'))
	
	# Remove widgets from hierarchy, shifting the transform to the object.
	mx = finalSpace.getMatrix('worldspace')
	parent = refSpace.getParent()
	if parent != None:
		parent.makeParent([object])
	else:
		object.clrParent()
	object.setMatrix(mx)
	
	for scene in Blender.Scene.Get():
		scene.unlink(finalSpace)
		scene.unlink(tracker)
		scene.unlink(refSpace)
		scene.update()
	
	Blender.Redraw()

if __name__ == "__main__":
	try:
		if Blender.bylink:
			# Running from a callback. Perform transform.
			TrackAndTransform(Blender.link)
		else:
			option = Blender.Draw.PupMenu(
				"Track And Transform%t" +
				"|Set up tracker" +
				"|Remove tracker"
			)
			if option == 1:
				# Set up tracker.
				sel = Blender.Object.GetSelected()
				if len(sel) != 2:
					raise Exception, "Select two objects. The first will track the second."
				object = sel[1]
				target = sel[0]
				try:
					object.getProperty('TRK_Role')
				except:
					pass
				else:
					raise Exception, "Object seems to contain old tracker properties. Delete them first."
				SetUpTracker(object, target)
			elif option == 2:
				# Tear down tracker.
				sel = Blender.Object.GetSelected()
				if len(sel) != 1:
					raise Exception, "Select one object from a tracker hierarchy."
				try:
					sel[0].getProperty('TRK_Role')
				except:
					raise Exception, "Object is not controlled by this script."
				TearDownTracker(sel[0])
	except Exception, e:
		print e
		Blender.Draw.PupMenu(str(e))
	#finally:
	#	pass