#!BPY

# """
# Name: 'Camera Tracking (.ma)'
# Blender: 240
# Group: 'Export'
# Tip: 'Export camera tracking to Maya Ascii (for Shake, After Effects etc)'
# """
# --------------------------------------------------------------------------
# ***** 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., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
#
# ***** END GPL LICENCE BLOCK *****
# --------------------------------------------------------------------------

import Blender

native_dialog = 1
try:
	from Tkinter import *
	import tkFileDialog
except:
	native_dialog = 0
	from Blender.Window import FileSelector

import time
from Blender import Mathutils
from Blender.Mathutils import *
from Blender.Window import DrawProgressBar

destinations = ['ae','shake','maya']
destination = destinations[0] # change index to 1 or 2 for Shake or Maya respectively

# AE and Shake struggle with lots of tracking points so this caps the number of static objects exported
max_static = 500

# having to check every object for every frame can be slow in a big scene so only exporting certain objects can help
limit_export = 0 # set to 1 to limit export to cameras and empties
export_types = ['Camera','Empty'] # set this list to objects you want to use as trackers

def fs_callback(mafilename): # callback for the FileSelector
	if(mafilename != ''): exportTracking(mafilename)

def saveDialog():
	root = Tk()
	root.withdraw()
	try: root.tk.call('console','hide')
	except: pass
	fname = tkFileDialog.asksaveasfilename(title='Save Maya Ascii file As')
	root.destroy()
	return fname
	
def fixObjName(name):
	return name.replace(".","_")

def exportTracking(mafilename):
	sscale = 1 # scene scale
	if destination == 'ae': sscale = 100
	tscale = 1 # tracking point scale
	
	oframe = Blender.Get('curframe')
	Blender.Set('curframe', Blender.Get('staframe'))
	start_frame = Blender.Get('staframe')
	end_frame = Blender.Get('endframe')
	num_frames = end_frame+1 - start_frame
			
	scene = Blender.Scene.GetCurrent()
	#obj = scene.getCurrentCamera()
	camera = Blender.Camera.Get()[0]
	camobj = scene.objects.camera
	
	context = scene.getRenderingContext()
	yRes = context.imageSizeY()
	xRes = context.imageSizeX()
	aspX = context.aspectX
	aspY = context.aspectY
	aspect = xRes*aspX / float(yRes*aspY)

	t = {}
	t['header'] = []
	t['header'].append('//Maya ASCII scene\n//Name: camera.ma\n//Last modified: %s\n' %time.strftime('%X %d/%m/%y'))
	if context.fps < 25: fps = 'film'
	elif context.fps < 30: fps = 'pal'
	else: fps = 'ntsc'
	t['header'].append('currentUnit -l centimeter -a degree -t %s;\n' %fps)
	
	if destination == 'ae':
		t['header'].append('select -ne :defaultResolution;\n')
		t['header'].append('setAttr ".w" %s;\n' %(xRes*aspX/aspY))
		t['header'].append('setAttr ".h" %s;\n' %yRes)
		t['header'].append('setAttr ".pa" 1.0;\n')
		
	matrix = camobj.matrix
	mx = RotationMatrix(-90, 4, "x")
	matrix = mx*matrix
				
	x = matrix.translationPart().x*sscale
	y = matrix.translationPart().y*sscale
	z = matrix.translationPart().z*sscale
	
	rx = matrix.toEuler().x
	ry = matrix.toEuler().y
	rz = matrix.toEuler().z
	
	sx = matrix.scalePart().x
	sy = matrix.scalePart().y
	sz = matrix.scalePart().z

	t['header'].append('createNode transform -n "%s";\n' %camobj.name)
	t['header'].append('\tsetAttr ".t" -type "double3" %s %s %s ;\n' %(x, z, -1*y))
	t['header'].append('\tsetAttr ".r" -type "double3" %s %s %s ;\n' %(rx, rz, ry))
	t['header'].append('\tsetAttr ".s" -type "double3" %s %s %s;\n' %(sx, sz, sy))
	
	t['header'].append('createNode camera -n "%s_null" -p "%s";\n' %(camobj.name,camobj.name))
	t['header'].append('\tsetAttr -k off ".v";\n')
	t['header'].append('\tsetAttr ".rnd" no;\n')
	
	t['header'].append('\tsetAttr ".w" %s;\n' %xRes)
	t['header'].append('\tsetAttr ".h" %s;\n' %yRes)
	
	# if aspect > 1: x is constant, otherwise y constant
	# constant comes from fov equations: fov = 2 * atan(aspect * 16.0 / camera.lens), aperturey = tan(fov/2)*camera.lens*2/25.4
	# tracking seems to work better using 1.25, change as needed
	#c = 16.0*2/25.4
	c = 1.25
	aperturex = min(c,c*aspect)
	aperturey = min(c,c/aspect)
		
	t['header'].append('\tsetAttr ".cap" -type "double2" %s %s ;\n' %(aperturex,aperturey))	
	t['header'].append('\tsetAttr ".dar" %s;\n' %aspect)
	t['header'].append('\tsetAttr ".ldar" no;\n')
	
	t['header'].append('\tsetAttr ".lsr" 1;\n') # lens squeeze ratio
	t['header'].append('\tsetAttr ".ff" 0;\n')
	t['header'].append('\tsetAttr ".fl" %s;\n' %camera.lens) # focal length
	t['header'].append('\tsetAttr ".ncp" 0.01;\n')
	t['header'].append('\tsetAttr ".ow" 30;\n')
	t['header'].append('\tsetAttr ".imn" -type "string" "camera1";\n')
	t['header'].append('\tsetAttr ".den" -type "string" "camera1_depth";\n')
	t['header'].append('\tsetAttr ".man" -type "string" "camera1_mask";\n')
	
	t['nodes'] = {}
	t['names'] = {}
	attrlist = ['_translateX','_translateY','_translateZ','_rotateX','_rotateY','_rotateZ','_scaleX','_scaleY','_scaleZ']
	for obj in scene.objects:
		if filter(lambda x:x in obj.layers, Blender.Window.ViewLayers() )==[]: continue
		#if limit_export and obj.type not in export_types: continue
		name = fixObjName(obj.name)
		t['names'][obj.name] = name
		t['nodes'][name] = []
		for p in attrlist:
			t[name+"_null"+p] = {'function':[], 'data':{}}
		
	DrawProgressBar (0, "starting...")
	# this creates major memory usage as it makes a key for every object every frame in memory
	# instead iterate over each object and then every frame and clean up data as you go
	fnum = 0
	pvals = {}
	frames = Blender.Get('endframe')+1-Blender.Get('staframe')
	tt12 = 0
		
	for frame in xrange(Blender.Get('staframe'),Blender.Get('endframe')+1):
		Blender.Set('curframe', frame)
		pc = (frame + Blender.Get('staframe') - 1)/float(frames)
		DrawProgressBar (pc, "frame %s" %str(frame))
		for obj in scene.objects:
			if fnum>0 and limit_export and obj.type not in export_types: continue
			if not t['names'].has_key(obj.name): continue # quicker than re-checking layers
			name = t['names'][obj.name]
					
			matrix = obj.matrix			
			matrix = mx*matrix			
			
			mt = matrix.translationPart()
			x = mt.x*sscale
			y = mt.y*sscale
			z = mt.z*sscale
			t[name+'_null_translateX']['data'][frame] = x			
			t[name+'_null_translateY']['data'][frame] = z
			t[name+'_null_translateZ']['data'][frame] = -1*y
			
			mr = matrix.toEuler()
			rx = mr.x
			ry = mr.y
			rz = mr.z
			t[name+'_null_rotateX']['data'][frame] = rx
			t[name+'_null_rotateY']['data'][frame] = rz
			t[name+'_null_rotateZ']['data'][frame] = ry
			
			ms = matrix.scalePart()
			sx = ms.x
			sy = ms.y
			sz = ms.z
			t[name+'_null_scaleX']['data'][frame] = sx
			t[name+'_null_scaleY']['data'][frame] = sz
			t[name+'_null_scaleZ']['data'][frame] = sy
						
			if fnum==0:
				if not pvals.has_key(name): pvals[name] = dict((p,[None,None]) for p in attrlist)
			del_last_frame = 0
			
			pvaldata = pvals[name]
			for p in attrlist:
				vv = t[name+"_null"+p]['data'][frame]
				if vv==pvaldata[p][0] and vv==pvaldata[p][1]: del_last_frame += 1
				pvaldata[p][1] = pvaldata[p][0]
				pvaldata[p][0] = vv			
				
			if del_last_frame == 9:
				for p in attrlist:
					del t[name+"_null"+p]['data'][frame-1]
					if fnum==2: del t[name+"_null"+p]['data'][frame-2]
					if fnum==frames-1: del t[name+"_null"+p]['data'][frame]
								
			if fnum==0:
				t['nodes'][name].append('createNode transform -n "%s";\n' %(name))
				t['nodes'][name].append('\tsetAttr ".t" -type "double3" %s %s %s ;\n' %(x, z, -1*y))
				t['nodes'][name].append('\tsetAttr ".r" -type "double3" %s %s %s ;\n' %(rx, rz, ry))
				t['nodes'][name].append('\tsetAttr ".s" -type "double3" %s %s %s;\n' %(sx, sz, sy))
				
				t['nodes'][name].append('createNode locator -n "%s_null" -p "%s";\n' %(name,name))
				t['nodes'][name].append('\tsetAttr -k off ".v";\n')
				
		fnum += 1
			
	t['footer'] = []
	static_objs = {}
	for obj in scene.objects:
		if filter(lambda x:x in obj.layers, Blender.Window.ViewLayers() )==[]: continue
		name = fixObjName(obj.name)
		if len(t[name+"_null_translateX"]['data'])==0: static_objs[name] = 0; continue
		if limit_export and obj.type not in export_types: static_objs[name] = 0; continue
		for p in attrlist:
			
			curve_type = "animCurveTL"
			if p in ['_rotateX','_rotateY','_rotateZ']: curve_type = "animCurveTA"
			elif p in ['_scaleX','_scaleY','_scaleZ']: curve_type = "animCurveTU"
			t[name+"_null"+p]['function'].append('createNode %s -n "%s_null%s";\n' %(curve_type,name,p))
			t[name+"_null"+p]['function'].append('\tsetAttr ".tan" 10;\n')
			t[name+"_null"+p]['function'].append('\tsetAttr ".wgt" no;\n')
			t[name+"_null"+p]['function'].append('\tsetAttr -s %s ".ktv[%s:%s]" ' %(len(t[name+"_null"+p]['data']), 0, len(t[name+"_null"+p]['data'])-1))
	
		t['footer'].append('connectAttr "%s_null_translateX.o" "%s.tx";\n' %(obj.name,obj.name))
		t['footer'].append('connectAttr "%s_null_translateY.o" "%s.ty";\n' %(obj.name,obj.name))
		t['footer'].append('connectAttr "%s_null_translateZ.o" "%s.tz";\n' %(obj.name,obj.name))
		t['footer'].append('connectAttr "%s_null_rotateX.o" "%s.rx";\n' %(obj.name,obj.name))
		t['footer'].append('connectAttr "%s_null_rotateY.o" "%s.ry";\n' %(obj.name,obj.name))
		t['footer'].append('connectAttr "%s_null_rotateZ.o" "%s.rz";\n' %(obj.name,obj.name))
		t['footer'].append('connectAttr "%s_null_scaleX.o" "%s.sx";\n' %(obj.name,obj.name))
		t['footer'].append('connectAttr "%s_null_scaleY.o" "%s.sy";\n' %(obj.name,obj.name))
		t['footer'].append('connectAttr "%s_null_scaleZ.o" "%s.sz";\n\n' %(obj.name,obj.name))
	
	t['footer'].append('\n// End of camera.ma\n')
			
	if not mafilename.endswith('.ma'):
		mafilename += '.ma'
	mafile = open(mafilename, 'w')
	for line in t['header']:
		mafile.write(line)
	num_static = 0
	for name,lines in t['nodes'].items():
		if static_objs.has_key(name) and destination!='maya':
			num_static += 1
			if(num_static>max_static): continue
		mafile.write(''.join(lines))
	for obj in scene.objects:
		if filter(lambda x:x in obj.layers, Blender.Window.ViewLayers() )==[]: continue
		if limit_export and obj.type not in export_types: continue
		name = fixObjName(obj.name)
		if len(t[name+"_null_translateX"]['data'])==0: continue
		#if obj.type not in ['Camera','Empty']: continue
		for p in attrlist:
			for line in t[name+"_null"+p]['function']:
				mafile.write(line)
			for frame,val in t[name+"_null"+p]['data'].items():
				mafile.write('%s %s ' %(frame, val))
			mafile.write(';\n')
			del t[name+"_null"+p]['function']
			del t[name+"_null"+p]['data']
	for line in t['footer']:
		mafile.write(line)
	
	# free up all data
	del t
	del pvals
	Blender.Set('curframe', oframe) # reset frame back to where it was

if native_dialog:
	mafilename = saveDialog()
	if(mafilename != ''): exportTracking(mafilename)
else:
	FileSelector(fs_callback, "Save Maya Ascii file As")