export_unreal_psk_psa.py

	# to map our bones to vertex groups
	#print("///////////////////////")
	#print("set influence")
	if blender_bone.name in psk_file.VertexGroups:
		vertex_list = psk_file.VertexGroups[blender_bone.name]
		#print("vertex list:", len(vertex_list), " of >" ,blender_bone.name )
		for vertex_data in vertex_list:
			#print("set influence vettex")
			point_index = vertex_data[0]
			vertex_weight = vertex_data[1]
			influence = VRawBoneInfluence()
			influence.Weight = vertex_weight
			#influence.BoneIndex = my_id
			influence.BoneIndex = my_id
			influence.PointIndex = point_index
			#print(influence)
			#print ('Adding Bone Influence for [%s] = Point Index=%i, Weight=%f' % (blender_bone.name, point_index, vertex_weight))
			#print("adding influence")
			psk_file.AddInfluence(influence)
	
	#blender_bone.matrix_local
	#recursively dump child bones
	mainparent = parent_matrix
	#if len(blender_bone.children) > 0:
	for current_child_bone in blender_bone.children:
		parse_bone(current_child_bone, psk_file, psa_file, my_id, 0, mainparent, parent_root)

def parse_armature(blender_armature, psk_file, psa_file):
	print ("----- parsing armature -----")
	print ('blender_armature length: %i' % (len(blender_armature)))
	
	#magic 0 sized root bone for UT - this is where all armature dummy bones will attach
	#dont increment nbone here because we initialize it to 1 (hackity hackity hack)

	#count top level bones first. NOT EFFICIENT.
	child_count = 0
	for current_obj in blender_armature: 
		current_armature = current_obj.data
		bones = [x for x in current_armature.bones if not x.parent == None]
		child_count += len(bones)

	for current_obj in blender_armature:
		print ("Current Armature Name: " + current_obj.name)
		current_armature = current_obj.data
		#armature_id = make_armature_bone(current_obj, psk_file, psa_file)
		
		#we dont want children here - only the top level bones of the armature itself
		#we will recursively dump the child bones as we dump these bones
		"""
		bones = [x for x in current_armature.bones if not x.parent == None]
		#will ingore this part of the ocde
		"""
		for current_bone in current_armature.bones: #list the bone. #note this will list all the bones.
			if(current_bone.parent == None):
				parse_bone(current_bone, psk_file, psa_file, 0, 0, current_obj.matrix, None)
				break

# get blender objects by type		
def get_blender_objects(objects, intype):
	return [x for x in objects if x.type == intype]
			
#strips current extension (if any) from filename and replaces it with extension passed in
def make_filename_ext(filename, extension):
	new_filename = ''
	extension_index = filename.find('.')
	
	if extension_index == -1:
		new_filename = filename + extension
	else:
		new_filename = filename[0:extension_index] + extension
		
	return new_filename

# returns the quaternion Grassman product a*b
# this is the same as the rotation a(b(x)) 
# (ie. the same as B*A if A and B are matrices representing 
# the rotations described by quaternions a and b)
def grassman(a, b):	
	return mathutils.Quaternion((
		a.w*b.w - a.x*b.x - a.y*b.y - a.z*b.z,
		a.w*b.x + a.x*b.w + a.y*b.z - a.z*b.y,
		a.w*b.y - a.x*b.z + a.y*b.w + a.z*b.x,
		a.w*b.z + a.x*b.y - a.y*b.x + a.z*b.w))
		
def parse_animation(blender_scene, blender_armatures, psa_file):
	#to do list:
	#need to list the action sets
	#need to check if there animation
	#need to check if animation is has one frame then exit it
	print ('\n----- parsing animation -----')
	#print(dir(blender_scene.render))
	render_data = blender_scene.render
	bHaveAction = True
	
	anim_rate = render_data.fps
	print("==== Blender Settings ====")
	print ('Scene: %s Start Frame: %i, End Frame: %i' % (blender_scene.name, blender_scene.start_frame, blender_scene.end_frame))
	print ('Frames Per Sec: %i' % anim_rate)
	print ("Default FPS: 24" )
	
	cur_frame_index = 0
	
	#print(dir(bpy.data.actions))
	#print(dir(bpy.context.scene.set))
	
	#list of armature objects
	for arm in blender_armatures:
		#check if there animation data from armature or something
		if not arm.animation_data:
			print("======================================")
			print("Check Animation Data: None")
			print("Armature has no animation, skipping...")
			print("======================================")
			break
			
		if not arm.animation_data.action:
			print("======================================")
			print("Check Action: None")
			print("Armature has no animation, skipping...")
			print("======================================")
			break
		act = arm.animation_data.action
		#print(dir(act))
		action_name = act.name
		
		if not len(act.fcurves):
			print("//===========================================================")
			print("// None bone pose set keys for this action set... skipping...")
			print("//===========================================================")
			bHaveAction = False
			
		#this deal with action export control
		if bHaveAction == True:
			print("")
			print("==== Action Set ====")
			print("Action Name:",action_name)
			#look for min and max frame that current set keys
			framemin, framemax = act.get_frame_range()
			#print("max frame:",framemax)
			start_frame = framemin
			end_frame = framemax
			scene_frames = range(start_frame, end_frame+1)
			frame_count = len(scene_frames)
			#===================================================
			anim = AnimInfoBinary()
			anim.Name = action_name
			anim.Group = "" #what is group?
			anim.NumRawFrames = frame_count
			anim.AnimRate = anim_rate
			anim.FirstRawFrame = cur_frame_index
			#===================================================
			count_previous_keys = len(psa_file.RawKeys.Data)
			print("Frame Key Set Count:",frame_count, "Total Frame:",frame_count)
			#print("init action bones...")
			unique_bone_indexes = {}
			# bone lookup table
			bones_lookup =  {}
		
			#build bone node for animation keys needed to be set
			for bone in arm.data.bones:
				bones_lookup[bone.name] = bone
			#print("bone name:",bone.name)
			frame_count = len(scene_frames)
			#print ('Frame Count: %i' % frame_count)
			pose_data = arm.pose
		
			#these must be ordered in the order the bones will show up in the PSA file!
			ordered_bones = {}
			ordered_bones = sorted([(psa_file.UseBone(x.name), x) for x in pose_data.bones], key=operator.itemgetter(0))
			
			#############################
			# ORDERED FRAME, BONE
			#for frame in scene_frames:
			
			for i in range(frame_count):
				frame = scene_frames[i]
				#LOUD
				#print ("==== outputting frame %i ===" % frame)
				
				if frame_count > i+1:
					next_frame = scene_frames[i+1]
					#print "This Frame: %i, Next Frame: %i" % (frame, next_frame)
				else:
					next_frame = -1
					#print "This Frame: %i, Next Frame: NONE" % frame
				
				#frame start from 1 as number one from blender
				blender_scene.set_frame(frame)
				
				cur_frame_index = cur_frame_index + 1
				for bone_data in ordered_bones:
					bone_index = bone_data[0]
					pose_bone = bone_data[1]
					#print("[=====POSE NAME:",pose_bone.name)
					
					#print("LENG >>.",len(bones_lookup))
					blender_bone = bones_lookup[pose_bone.name]
					
					#just need the total unique bones used, later for this AnimInfoBinary
					unique_bone_indexes[bone_index] = bone_index
					#LOUD
					#print ("-------------------", pose_bone.name)
					head = pose_bone.head
					
					posebonemat = mathutils.Matrix(pose_bone.matrix)
					parent_pose = pose_bone.parent
					if parent_pose != None:
						parentposemat = mathutils.Matrix(parent_pose.matrix)
						#blender 2.4X it been flip around with new 2.50 (mat1 * mat2) should now be (mat2 * mat1)
						posebonemat = parentposemat.invert() * posebonemat
					head = posebonemat.translation_part()
					quat = posebonemat.to_quat().normalize()
					vkey = VQuatAnimKey()
					vkey.Position.X = head.x
					vkey.Position.Y = head.y
					vkey.Position.Z = head.z
					
					if parent_pose != None:
						quat = make_fquat(quat)
					else:
						quat = make_fquat_default(quat)
					
					vkey.Orientation = quat
					#print("Head:",head)
					#print("Orientation",quat)
					
					#time from now till next frame = diff / framesPerSec
					if next_frame >= 0:
						diff = next_frame - frame
					else:
						diff = 1.0
					
					#print ("Diff = ", diff)
					vkey.Time = float(diff)/float(anim_rate)
					
					psa_file.AddRawKey(vkey)
					
			#done looping frames
			#done looping armatures
			#continue adding animInfoBinary counts here
		
			anim.TotalBones = len(unique_bone_indexes)
			print("Bones Count:",anim.TotalBones)
			anim.TrackTime = float(frame_count) / anim.AnimRate
			print("Time Track Frame:",anim.TrackTime)
			psa_file.AddAnimation(anim)
			print("==== Finish Action Build(s) ====")
		
exportmessage = "Export Finish"		
		
def fs_callback(filename, context, user_setting):
	#this deal with repeat export and the reset settings
	global bonedata, BBCount, nbone, exportmessage
	bonedata = []#clear array
	BBCount = 0
	nbone = 0
	
	start_time = time.clock()
	
	print ("========EXPORTING TO UNREAL SKELETAL MESH FORMATS========\r\n")
	print("Blender Version:", bpy.app.version_string)
	
	psk = PSKFile()
	psa = PSAFile()
	
	#sanity check - this should already have the extension, but just in case, we'll give it one if it doesn't
	psk_filename = make_filename_ext(filename, '.psk')
	
	#make the psa filename
	psa_filename = make_filename_ext(filename, '.psa')
	
	print ('PSK File: ' +  psk_filename)
	print ('PSA File: ' +  psa_filename)
	
	barmature = True
	bmesh = True
	blender_meshes = []
	blender_armature = []
	selectmesh = []
	selectarmature = []
	
	current_scene = context.scene
	cur_frame = current_scene.current_frame #store current frame before we start walking them during animation parse
	objects = current_scene.objects
	
	print("Checking object count...")
	for next_obj in objects:
		if next_obj.type == 'MESH':
			blender_meshes.append(next_obj)
			if (next_obj.selected):
				#print("mesh object select")
				selectmesh.append(next_obj)
		if next_obj.type == 'ARMATURE':
			blender_armature.append(next_obj)
			if (next_obj.selected):
				#print("armature object select")
				selectarmature.append(next_obj)
	
	print("Mesh Count:",len(blender_meshes)," Armature Count:",len(blender_armature))
	print("====================================")
	print("Checking Mesh Condtion(s):")
	if len(blender_meshes) == 1:
		print(" - One Mesh Scene")
	elif (len(blender_meshes) > 1) and (len(selectmesh) == 1):
		print(" - One Mesh [Select]")
	else:
		print(" - Too Many Meshes!")
		print(" - Select One Mesh Object!")
		bmesh = False
	print("====================================")
	print("Checking Armature Condtion(s):")
	if len(blender_armature) == 1:
		print(" - One Armature Scene")
	elif (len(blender_armature) > 1) and (len(selectarmature) == 1):
		print(" - One Armature [Select]")
	else:
		print(" - Too Armature Meshes!")
		print(" - Select One Armature Object Only!")
		barmature = False
	
	if 	(bmesh == False) or (barmature == False):
		exportmessage = "Export Fail! Check Log."
		print("=================================")
		print("= Export Fail!                  =")
		print("=================================")
	else:
		exportmessage = "Export Finish!"
		#need to build a temp bone index for mesh group vertex
		BoneIndexArmature(blender_armature)

		try:
			#######################
			# STEP 1: MESH DUMP
			# we build the vertexes, wedges, and faces in here, as well as a vertexgroup lookup table
			# for the armature parse
			print("//===============================")
			print("// STEP 1")
			print("//===============================")
			parse_meshes(blender_meshes, psk)
		except:
			context.scene.set_frame(cur_frame) #set frame back to original frame
			print ("Exception during Mesh Parse")
			raise
		
		try:
			#######################
			# STEP 2: ARMATURE DUMP
			# IMPORTANT: do this AFTER parsing meshes - we need to use the vertex group data from 
			# the mesh parse in here to generate bone influences
			print("//===============================")
			print("// STEP 2")
			print("//===============================")
			parse_armature(blender_armature, psk, psa) 
			
		except:
			context.scene.set_frame(cur_frame) #set frame back to original frame
			print ("Exception during Armature Parse")
			raise

		try:
			#######################
			# STEP 3: ANIMATION DUMP
			# IMPORTANT: do AFTER parsing bones - we need to do bone lookups in here during animation frames
			print("//===============================")
			print("// STEP 3")
			print("//===============================")
			parse_animation(current_scene, blender_armature, psa) 
			
		except:
			context.scene.set_frame(cur_frame) #set frame back to original frame
			print ("Exception during Animation Parse")
			raise

		# reset current frame
		
		context.scene.set_frame(cur_frame) #set frame back to original frame
		
		##########################
		# FILE WRITE
		print("//===========================================")
		print("// bExportPsk:",bpy.context.scene.unrealexportpsk," bExportPsa:",bpy.context.scene.unrealexportpsa)
		print("//===========================================")
		if bpy.context.scene.unrealexportpsk == True:
			print("Writing Skeleton Mesh Data...")
			#RG - dump psk file
			psk.PrintOut()
			file = open(psk_filename, "wb") 
			file.write(psk.dump())
			file.close() 
			print ("Successfully Exported File: " + psk_filename)
		if bpy.context.scene.unrealexportpsa == True:
			print("Writing Animaiton Data...")
			#RG - dump psa file
			if not psa.IsEmpty():
				psa.PrintOut()
				file = open(psa_filename, "wb") 
				file.write(psa.dump())
				file.close() 
				print ("Successfully Exported File: " + psa_filename)
			else:
				print ("No Animations (.psa file) to Export")

		print ('PSK/PSA Export Script finished in %.2f seconds' % (time.clock() - start_time))
		
		#MSG BOX EXPORT COMPLETE
		#...

		#DONE
		print ("PSK/PSA Export Complete")

def write_data(path, context, user_setting):
	print("//============================")
	print("// running psk/psa export...")
	print("//============================")
	fs_callback(path, context, user_setting)
	pass

from bpy.props import *

exporttypedata = []

# [index,text field,0] #or something like that
exporttypedata.append(("0","PSK","Export PSK"))
exporttypedata.append(("1","PSA","Export PSA"))
exporttypedata.append(("2","ALL","Export ALL"))

IntProperty= bpy.types.Scene.IntProperty

IntProperty(attr="unrealfpsrate", name="fps rate",
    description="Set the frame per second (fps) for unreal.",
    default=24,min=1,max=100)
	
bpy.types.Scene.EnumProperty( attr="unrealexport_settings",
    name="Export:",
    description="Select a export settings (psk/psa/all)...",
    items = exporttypedata, default = '0')
		
bpy.types.Scene.BoolProperty( attr="unrealtriangulatebool",
    name="Triangulate Mesh",
    description="Convert Quad to Tri Mesh Boolean...",
    default=False)
	
bpy.types.Scene.BoolProperty( attr="unrealexportpsk",
    name="bool export psa",
    description="bool for exporting this psk format",
    default=False)
	
bpy.types.Scene.BoolProperty( attr="unrealexportpsa",
    name="bool export psa",
    description="bool for exporting this psa format",
    default=False)

class ExportUDKAnimData(bpy.types.Operator):
	global exportmessage
	'''Export Skeleton Mesh / Animation Data file(s)'''
	bl_idname = "export.udk_anim_data" # this is important since its how bpy.ops.export.udk_anim_data is constructed
	bl_label = "Export PSK/PSA"
	__doc__ = "One mesh and one armature else select one mesh or armature to be exported."

	# List of operator properties, the attributes will be assigned
	# to the class instance from the operator settings before calling.

	# TODO, add props
	path = StringProperty(name="File Path", description="File path used for exporting the PSA file", maxlen= 1024, default= "")
	use_setting = BoolProperty(name="No Options Yet", description="No Options Yet", default= True)
	filename = StringProperty(name="filename", description="", maxlen= 1024, default= "")
	directory = StringProperty(name="directory", description="", maxlen= 1024, default= "")
	pskexportbool = BoolProperty(name="Export PSK", description="Export Skeletal Mesh", default= True)
	psaexportbool = BoolProperty(name="Export PSA", description="Export Action Set (Animation Data)", default= True)
	#fpssettomg = IntProperty(name="FPS", attr="fpsunrealexport",description="", default=24, min=1, max=120, soft_min=1, soft_max=120, step=1, options={'ANIMATABLE'})

	def poll(self, context):
		return context.active_object != None

	def execute(self, context):
		#check if  skeleton mesh is needed to be exported
		if (self.properties.pskexportbool):
			bpy.context.scene.unrealexportpsk = True
		else:
			bpy.context.scene.unrealexportpsk = False
		#check if  animation data is needed to be exported
		if (self.properties.psaexportbool):
			bpy.context.scene.unrealexportpsa = True
		else:
			bpy.context.scene.unrealexportpsa = False
			
		#print(">>>>>",dir(self))
		#self.properties.pskexportbool
		#self.layout.prop(context,context.scene.render,"fps")
		write_data(self.properties.path, context, self.properties.use_setting)
		
		self.report({'WARNING', 'INFO'}, exportmessage)
		return {'FINISHED'}
		
	def invoke(self, context, event):
		wm = context.manager
		wm.add_fileselect(self)
		return {'RUNNING_MODAL'}	
		
def menu_func(self, context):
	bpy.context.scene.unrealexportpsk = True
	bpy.context.scene.unrealexportpsa = True
	default_path = bpy.data.filename.replace(".blend", ".psk")
	self.layout.operator ("export.udk_anim_data", text="Skeleton Mesh / Animation Data (.psk/.psa)").path = default_path

class VIEW3D_PT_unrealtools_objectmode(bpy.types.Panel):
	bl_space_type = "VIEW_3D"
	bl_region_type = "TOOLS"
	bl_label = "Unreal Tools"
	
	def poll(self, context):
		return context.active_object
    
	def draw(self, context):
		layout = self.layout
		#layout.label(text="Unreal Tools")
		rd = context.scene
		#drop box
		layout.prop(rd, "unrealexport_settings",expand=True)
		#layout.prop(rd, "unrealexport_settings")
		#button
		layout.operator("object.UnrealExport")
		#FPS #it use the real data from your scene
		layout.prop(rd.render, "fps")
		
		#row = layout.row()
		#row.label(text="Action Set(s)(not build)")
		#for action in  bpy.data.actions:
			#print(dir( action))
			#print(action.get_frame_range())
			#row = layout.row()
			#row.prop(action, "name")
			
			#print(dir(action.groups[0]))
			#for g in action.groups:#those are bones
				#print("group...")
				#print(dir(g))
				#print("////////////")
				#print((g.name))
				#print("////////////")
			
			#row.label(text="Active:" + action.select)
		btrimesh = False
		
class OBJECT_OT_UnrealExport(bpy.types.Operator):
	global exportmessage
	bl_idname = "OBJECT_OT_UnrealExport"
	bl_label = "Unreal Export"
	__doc__ = "Select export setting for .psk/.psa or both."
	
	def invoke(self, context, event):
		path = StringProperty(name="File Path", description="File path used for exporting the PSA file", maxlen= 1024, default= "")
		print("Init Export Script:")
		if(int(bpy.context.scene.unrealexport_settings) == 0):
			bpy.context.scene.unrealexportpsk = True
			bpy.context.scene.unrealexportpsa = False
			print("Exporting PSK...")
		if(int(bpy.context.scene.unrealexport_settings) == 1):
			bpy.context.scene.unrealexportpsk = False
			bpy.context.scene.unrealexportpsa = True
			print("Exporting PSA...")
		if(int(bpy.context.scene.unrealexport_settings) == 2):
			bpy.context.scene.unrealexportpsk = True
			bpy.context.scene.unrealexportpsa = True
			print("Exporting ALL...")
		default_path = bpy.data.filename.replace(".blend", ".psk")
		fs_callback(default_path, bpy.context, False)
		self.report({'WARNING', 'INFO'}, exportmessage)
		return{'FINISHED'}	
	
def register():
	global MENUPANELBOOL
	if MENUPANELBOOL:
		bpy.types.register(OBJECT_OT_UnrealExport)
		bpy.types.register(VIEW3D_PT_unrealtools_objectmode)
	bpy.types.register(ExportUDKAnimData)
	bpy.types.INFO_MT_file_export.append(menu_func)
    
def unregister():
	global MENUPANELBOOL
	if MENUPANELBOOL:
		bpy.types.unregister(OBJECT_OT_UnrealExport)
		bpy.types.unregister(VIEW3D_PT_unrealtools_objectmode)
	bpy.types.unregister(ExportUDKAnimData)
	bpy.types.INFO_MT_file_export.remove(menu_func)

if __name__ == "__main__":
    register()