io_export_unreal_psk_psa.py

#  ***** 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 3 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, see <http://www.gnu.org/licenses/>.
#  All rights reserved.
#  ***** GPL LICENSE BLOCK *****

bl_addon_info = {
	"name": "Export Skeleletal Mesh/Animation Data",
	"author": "Darknet/Optimus_P-Fat/Active_Trash/Sinsoft",
	"version": (2,0),
	"blender": (2, 5, 3),
	"api": 31847,
	"location": "File > Export > Skeletal Mesh/Animation Data (.psk/.psa)",
	"description": "Export Unreal Engine (.psk)",
	"warning": "",
	"wiki_url": "http://wiki.blender.org/index.php/Extensions:2.5/Py/"\
		"Scripts/File_I-O/Unreal_psk_psa",
	"tracker_url": "https://projects.blender.org/tracker/index.php?"\
		"func=detail&aid=21366&group_id=153&atid=469",
	"category": "Import/Export"}

"""
-- Unreal Skeletal Mesh and Animation Export (.psk  and .psa) export script v0.0.1 --<br> 

- NOTES:
- This script Exports To Unreal's PSK and PSA file formats for Skeletal Meshes and Animations. <br>
- This script DOES NOT support vertex animation! These require completely different file formats. <br>

- v0.0.1
- Initial version

- v0.0.2
- This version adds support for more than one material index!

[ - Edit by: Darknet
- v0.0.3 - v0.0.12
- This will work on UT3 and it is a stable version that work with vehicle for testing. 
- Main Bone fix no dummy needed to be there.
- Just bone issues position, rotation, and offset for psk.
- The armature bone position, rotation, and the offset of the bone is fix. It was to deal with skeleton mesh export for psk.
- Animation is fix for position, offset, rotation bone support one rotation direction when armature build. 
- It will convert your mesh into triangular when exporting to psk file.
- Did not work with psa export yet.

- v0.0.13
- The animatoin will support different bone rotations when export the animation.

- v0.0.14
- Fixed Action set keys frames when there is no pose keys and it will ignore it.

- v0.0.15
- Fixed multiple objects when exporting to psk. Select one mesh to export to psk.
- ]

- v0.1.1
- Blender 2.50 svn (Support)

Credit to:
- export_cal3d.py (Position of the Bones Format)
- blender2md5.py (Animation Translation Format)
- export_obj.py (Blender 2.5/Pyhton 3.x Format)

- freenode #blendercoder -> user -> ideasman42

- Give Credit to those who work on this script.

- http://sinsoft.com
"""

import os
import time
import datetime
import bpy
import mathutils
import operator

from struct import pack, calcsize


# REFERENCE MATERIAL JUST IN CASE:
# 
# U = x / sqrt(x^2 + y^2 + z^2)
# V = y / sqrt(x^2 + y^2 + z^2)
#
# Triangles specifed counter clockwise for front face
#
#defines for sizeofs
SIZE_FQUAT = 16
SIZE_FVECTOR = 12
SIZE_VJOINTPOS = 44
SIZE_ANIMINFOBINARY = 168
SIZE_VCHUNKHEADER = 32
SIZE_VMATERIAL = 88
SIZE_VBONE = 120
SIZE_FNAMEDBONEBINARY = 120
SIZE_VRAWBONEINFLUENCE = 12
SIZE_VQUATANIMKEY = 32
SIZE_VVERTEX = 16
SIZE_VPOINT = 12
SIZE_VTRIANGLE = 12

########################################################################
# Generic Object->Integer mapping
# the object must be usable as a dictionary key
class ObjMap:
	def __init__(self):
		self.dict = {}
		self.next = 0
	def get(self, obj):
		if obj in self.dict:
			return self.dict[obj]
		else:
			id = self.next
			self.next = self.next + 1
			self.dict[obj] = id
			return id
		
	def items(self):
		getval = operator.itemgetter(0)
		getkey = operator.itemgetter(1)
		return map(getval, sorted(self.dict.items(), key=getkey))

########################################################################
# RG - UNREAL DATA STRUCTS - CONVERTED FROM C STRUCTS GIVEN ON UDN SITE 
# provided here: http://udn.epicgames.com/Two/BinaryFormatSpecifications.html
# updated UDK (Unreal Engine 3): http://udn.epicgames.com/Three/BinaryFormatSpecifications.html
class FQuat:
	def __init__(self): 
		self.X = 0.0
		self.Y = 0.0
		self.Z = 0.0
		self.W = 1.0
		
	def dump(self):
		data = pack('ffff', self.X, self.Y, self.Z, self.W)
		return data
		
	def __cmp__(self, other):
		return cmp(self.X, other.X) \
			or cmp(self.Y, other.Y) \
			or cmp(self.Z, other.Z) \
			or cmp(self.W, other.W)
		
	def __hash__(self):
		return hash(self.X) ^ hash(self.Y) ^ hash(self.Z) ^ hash(self.W)
		
	def __str__(self):
		return "[%f,%f,%f,%f](FQuat)" % (self.X, self.Y, self.Z, self.W)
		
class FVector(object):
	def __init__(self, X=0.0, Y=0.0, Z=0.0):
		self.X = X
		self.Y = Y
		self.Z = Z
		
	def dump(self):
		data = pack('fff', self.X, self.Y, self.Z)
		return data
		
	def __cmp__(self, other):
		return cmp(self.X, other.X) \
			or cmp(self.Y, other.Y) \
			or cmp(self.Z, other.Z)
		
	def _key(self):
		return (type(self).__name__, self.X, self.Y, self.Z)
		
	def __hash__(self):
		return hash(self._key())
		
	def __eq__(self, other):
		if not hasattr(other, '_key'):
			return False
		return self._key() == other._key() 
		
	def dot(self, other):
		return self.X * other.X + self.Y * other.Y + self.Z * other.Z
	
	def cross(self, other):
		return FVector(self.Y * other.Z - self.Z * other.Y,
				self.Z * other.X - self.X * other.Z,
				self.X * other.Y - self.Y * other.X)
				
	def sub(self, other):
		return FVector(self.X - other.X,
			self.Y - other.Y,
			self.Z - other.Z)

class VJointPos:
	def __init__(self):
		self.Orientation = FQuat()
		self.Position = FVector()
		self.Length = 0.0
		self.XSize = 0.0
		self.YSize = 0.0
		self.ZSize = 0.0
		
	def dump(self):
		data = self.Orientation.dump() + self.Position.dump() + pack('4f', self.Length, self.XSize, self.YSize, self.ZSize)
		return data
			
class AnimInfoBinary:
	def __init__(self):
		self.Name = "" # length=64
		self.Group = ""    # length=64
		self.TotalBones = 0
		self.RootInclude = 0
		self.KeyCompressionStyle = 0
		self.KeyQuotum = 0
		self.KeyPrediction = 0.0
		self.TrackTime = 0.0
		self.AnimRate = 0.0
		self.StartBone = 0
		self.FirstRawFrame = 0
		self.NumRawFrames = 0
		
	def dump(self):
		data = pack('64s64siiiifffiii', self.Name, self.Group, self.TotalBones, self.RootInclude, self.KeyCompressionStyle, self.KeyQuotum, self.KeyPrediction, self.TrackTime, self.AnimRate, self.StartBone, self.FirstRawFrame, self.NumRawFrames)
		return data

class VChunkHeader:
	def __init__(self, name, type_size):
		self.ChunkID = name # length=20
		self.TypeFlag = 1999801 # special value
		self.DataSize = type_size
		self.DataCount = 0
		
	def dump(self):
		data = pack('20siii', self.ChunkID, self.TypeFlag, self.DataSize, self.DataCount)
		return data
		
class VMaterial:
	def __init__(self):
		self.MaterialName = "" # length=64
		self.TextureIndex = 0
		self.PolyFlags = 0 # DWORD
		self.AuxMaterial = 0
		self.AuxFlags = 0 # DWORD
		self.LodBias = 0
		self.LodStyle = 0
		
	def dump(self):
		data = pack('64siLiLii', self.MaterialName, self.TextureIndex, self.PolyFlags, self.AuxMaterial, self.AuxFlags, self.LodBias, self.LodStyle)
		return data

class VBone:
	def __init__(self):
		self.Name = "" # length = 64
		self.Flags = 0 # DWORD
		self.NumChildren = 0
		self.ParentIndex = 0
		self.BonePos = VJointPos()
		
	def dump(self):
		data = pack('64sLii', self.Name, self.Flags, self.NumChildren, self.ParentIndex) + self.BonePos.dump()
		return data

#same as above - whatever - this is how Epic does it...        
class FNamedBoneBinary:
	def __init__(self):
		self.Name = "" # length = 64
		self.Flags = 0 # DWORD
		self.NumChildren = 0
		self.ParentIndex = 0
		self.BonePos = VJointPos()
		
		self.IsRealBone = 0  # this is set to 1 when the bone is actually a bone in the mesh and not a dummy
		
	def dump(self):
		data = pack('64sLii', self.Name, self.Flags, self.NumChildren, self.ParentIndex) + self.BonePos.dump()
		return data
	
class VRawBoneInfluence:
	def __init__(self):
		self.Weight = 0.0
		self.PointIndex = 0
		self.BoneIndex = 0
		
	def dump(self):
		data = pack('fii', self.Weight, self.PointIndex, self.BoneIndex)
		return data
		
class VQuatAnimKey:
	def __init__(self):
		self.Position = FVector()
		self.Orientation = FQuat()
		self.Time = 0.0
		
	def dump(self):
		data = self.Position.dump() + self.Orientation.dump() + pack('f', self.Time)
		return data
		
class VVertex(object):
	def __init__(self):
		self.PointIndex = 0 # WORD
		self.U = 0.0
		self.V = 0.0
		self.MatIndex = 0 #BYTE
		self.Reserved = 0 #BYTE
		
	def dump(self):
		data = pack('HHffBBH', self.PointIndex, 0, self.U, self.V, self.MatIndex, self.Reserved, 0)
		return data
		
	def __cmp__(self, other):
		return cmp(self.PointIndex, other.PointIndex) \
			or cmp(self.U, other.U) \
			or cmp(self.V, other.V) \
			or cmp(self.MatIndex, other.MatIndex) \
			or cmp(self.Reserved, other.Reserved)
	
	def _key(self):
		return (type(self).__name__,self.PointIndex, self.U, self.V,self.MatIndex,self.Reserved)
		
	def __hash__(self):
		return hash(self._key())
		
	def __eq__(self, other):
		if not hasattr(other, '_key'):
			return False
		return self._key() == other._key()
		
class VPoint(object):
	def __init__(self):
		self.Point = FVector()
		
	def dump(self):
		return self.Point.dump()
		
	def __cmp__(self, other):
		return cmp(self.Point, other.Point)
	
	def _key(self):
		return (type(self).__name__, self.Point)
	
	def __hash__(self):
		return hash(self._key())
		
	def __eq__(self, other):
		if not hasattr(other, '_key'):
			return False
		return self._key() == other._key() 
		
class VTriangle:
	def __init__(self):
		self.WedgeIndex0 = 0 # WORD
		self.WedgeIndex1 = 0 # WORD
		self.WedgeIndex2 = 0 # WORD
		self.MatIndex = 0 # BYTE
		self.AuxMatIndex = 0 # BYTE
		self.SmoothingGroups = 0 # DWORD
		
	def dump(self):
		data = pack('HHHBBL', self.WedgeIndex0, self.WedgeIndex1, self.WedgeIndex2, self.MatIndex, self.AuxMatIndex, self.SmoothingGroups)
		return data

# END UNREAL DATA STRUCTS
########################################################################

########################################################################
#RG - helper class to handle the normal way the UT files are stored 
#as sections consisting of a header and then a list of data structures
class FileSection:
	def __init__(self, name, type_size):
		self.Header = VChunkHeader(name, type_size)
		self.Data = [] # list of datatypes
		
	def dump(self):
		data = self.Header.dump()
		for i in range(len(self.Data)):
			data = data + self.Data[i].dump()
		return data
		
	def UpdateHeader(self):
		self.Header.DataCount = len(self.Data)
		
class PSKFile:
	def __init__(self):
		self.GeneralHeader = VChunkHeader("ACTRHEAD", 0)
		self.Points = FileSection("PNTS0000", SIZE_VPOINT)        #VPoint
		self.Wedges = FileSection("VTXW0000", SIZE_VVERTEX)        #VVertex
		self.Faces = FileSection("FACE0000", SIZE_VTRIANGLE)        #VTriangle
		self.Materials = FileSection("MATT0000", SIZE_VMATERIAL)    #VMaterial
		self.Bones = FileSection("REFSKELT", SIZE_VBONE)        #VBone
		self.Influences = FileSection("RAWWEIGHTS", SIZE_VRAWBONEINFLUENCE)    #VRawBoneInfluence
		
		#RG - this mapping is not dumped, but is used internally to store the new point indices 
		# for vertex groups calculated during the mesh dump, so they can be used again
		# to dump bone influences during the armature dump
		#
		# the key in this dictionary is the VertexGroup/Bone Name, and the value
		# is a list of tuples containing the new point index and the weight, in that order
		#
		# Layout:
		# { groupname : [ (index, weight), ... ], ... }
		#
		# example: 
		# { 'MyVertexGroup' : [ (0, 1.0), (5, 1.0), (3, 0.5) ] , 'OtherGroup' : [(2, 1.0)] }
		
		self.VertexGroups = {} 
		
	def AddPoint(self, p):
		#print ('AddPoint')
		self.Points.Data.append(p)
		
	def AddWedge(self, w):
		#print ('AddWedge')
		self.Wedges.Data.append(w)
	
	def AddFace(self, f):
		#print ('AddFace')
		self.Faces.Data.append(f)
		
	def AddMaterial(self, m):
		#print ('AddMaterial')
		self.Materials.Data.append(m)
		
	def AddBone(self, b):
		#print ('AddBone [%s]: Position: (x=%f, y=%f, z=%f) Rotation=(%f,%f,%f,%f)'  % (b.Name, b.BonePos.Position.X, b.BonePos.Position.Y, b.BonePos.Position.Z, b.BonePos.Orientation.X,b.BonePos.Orientation.Y,b.BonePos.Orientation.Z,b.BonePos.Orientation.W))
		self.Bones.Data.append(b)
		
	def AddInfluence(self, i):
		#print ('AddInfluence')
		self.Influences.Data.append(i)
		
	def UpdateHeaders(self):
		self.Points.UpdateHeader()
		self.Wedges.UpdateHeader()
		self.Faces.UpdateHeader()
		self.Materials.UpdateHeader()
		self.Bones.UpdateHeader()
		self.Influences.UpdateHeader()
		
	def dump(self):
		self.UpdateHeaders()
		data = self.GeneralHeader.dump() + self.Points.dump() + self.Wedges.dump() + self.Faces.dump() + self.Materials.dump() + self.Bones.dump() + self.Influences.dump()
		return data
		
	def GetMatByIndex(self, mat_index):
		if mat_index >= 0 and len(self.Materials.Data) > mat_index:
			return self.Materials.Data[mat_index]
		else:
			m = VMaterial()
			m.MaterialName = "Mat%i" % mat_index
			self.AddMaterial(m)
			return m
		
	def PrintOut(self):
		print ("--- PSK FILE EXPORTED ---")
		print ('point count: %i' % len(self.Points.Data))
		print ('wedge count: %i' % len(self.Wedges.Data))
		print ('face count: %i' % len(self.Faces.Data))
		print ('material count: %i' % len(self.Materials.Data))
		print ('bone count: %i' % len(self.Bones.Data))
		print ('inlfuence count: %i' % len(self.Influences.Data))
		print ('-------------------------')

# PSA FILE NOTES FROM UDN:
#
#    The raw key array holds all the keys for all the bones in all the specified sequences, 
#    organized as follows:
#    For each AnimInfoBinary's sequence there are [Number of bones] times [Number of frames keys] 
#    in the VQuatAnimKeys, laid out as tracks of [numframes] keys for each bone in the order of 
#    the bones as defined in the array of FnamedBoneBinary in the PSA. 
#
#    Once the data from the PSK (now digested into native skeletal mesh) and PSA (digested into 
#    a native animation object containing one or more sequences) are associated together at runtime, 
#    bones are linked up by name. Any bone in a skeleton (from the PSK) that finds no partner in 
#    the animation sequence (from the PSA) will assume its reference pose stance ( as defined in 
#    the offsets & rotations that are in the VBones making up the reference skeleton from the PSK)

class PSAFile:
	def __init__(self):
		self.GeneralHeader = VChunkHeader("ANIMHEAD", 0)
		self.Bones = FileSection("BONENAMES", SIZE_FNAMEDBONEBINARY)    #FNamedBoneBinary
		self.Animations = FileSection("ANIMINFO", SIZE_ANIMINFOBINARY)    #AnimInfoBinary
		self.RawKeys = FileSection("ANIMKEYS", SIZE_VQUATANIMKEY)    #VQuatAnimKey
		
		# this will take the format of key=Bone Name, value = (BoneIndex, Bone Object)
		# THIS IS NOT DUMPED
		self.BoneLookup = {} 
		
	def dump(self