#!BPY

""" reg info:
Name: 'BRF Export (.brf)'
Blender: 235
Group: 'Export'
Tip: 'Export Skeletons, Meshes and Animations to brf-Format'
"""

__author__ = "Dirk Schulz"
__version__ = "0.1"

__bpydoc__ = """\
Exports from Blender to the brf-FileFormat (used by Mount&Blade)
"""

import Blender
from Blender.Mathutils import *
from Blender import Draw, BGL
from Blender.BGL import *
import struct, os, re
from math import *
import os

global toggle_scale_10, scale_10
global args

args = ''
toggle_scale_10 = 0
scale_10 = 0

scalemat = ScaleMatrix(0.1, 4)

class BoneType:
	def __init__(self, index, bone):
		self.index = index
		self.parent = None
		self.bone = bone
		self.pos = []
		self.name = ""
		self.matrix = Matrix()
		self.posemat = Matrix()
		
class BoneWeight:
	def __init__(self, index):
		self.index = index
		self.weights = []
		self.endweight = 0
		self.groups = []
		
	def addWeight(self, weight, group):
		self.weights.append(weight)
		self.groups.append(group)
		
	def normalizeWeights(self):
		totalweight = 0
		for weight in self.weights:
			totalweight += weight
			
		self.endweight = totalweight / len(self.weights)
		
class VertexUV:
	def __init__(self, index, uv, faceindex):
		self.vertindex = index
		self.uv = uv
		self.faceindex = faceindex

def writeInt(file, int):

	file.write(struct.pack('<i', int))
	
def writeFloat(file, f):

	file.write(struct.pack('<f', f))
	
def writeUInt(file, uint):

	file.write(struct.pack('<I', uint))
	
def writeChar(file, ch):

	file.write(struct.pack('<B', ch))
	
def equalVector2(vec1, vec2):

	if fabs(vec1[0] - vec2[0]) < 1e-3 and fabs(vec1[1] - vec2[1]) < 1e-3:
		return True
		
	return False
	
def printMatrix4x4(m):

	print '%f %f %f %f' % (m[0][0], m[0][1], m[0][2], m[0][3])
	print '%f %f %f %f' % (m[1][0], m[1][1], m[1][2], m[1][3])
	print '%f %f %f %f' % (m[2][0], m[2][1], m[2][2], m[2][3])
	print '%f %f %f %f' % (m[3][0], m[3][1], m[3][2], m[3][3])
	print '-----------------------'
	
def printMatrix3x3(m):

	print '%f %f %f' % (m[0][0], m[0][1], m[0][2])
	print '%f %f %f' % (m[1][0], m[1][1], m[1][2])
	print '%f %f %f' % (m[2][0], m[2][1], m[2][2])
	print '-----------------------'

class SkeletonExport:
	def __init__(self):
	
		self.bones = []
		self.ArmatureMat = Matrix()
		
	def callback(self, filename):
		self.export(filename)
		
	def addBone(self, bone, parent):
	
		newbone = BoneType(len(self.bones), bone)
		newbone.name = bone.name
		newbone.parent = parent
		newbone.matrix = Matrix(bone.matrix['BONESPACE'])
		self.bones.append(newbone)		
		
		if bone.children:
			for child in bone.children:
				self.addBone(child, newbone)

	def export(self, filename):

		global anim
		
		os.system('cls')

		file = open(filename, "wb")
		
		scn = Blender.Scene.GetCurrent()
		
		selobj = scn.objects.selected
		
		for obj in selobj:
		
			if obj.type == 'Armature':
			
				writeInt(file, len('skeleton'))
		
				file.write('skeleton')
				
				self.ArmatureMat = obj.getMatrix()
				
				if scale_10:
					self.ArmatureMat = scalemat * self.ArmatureMat
					
				printMatrix4x4(self.ArmatureMat)
				
				arm_obj = obj.getData()
							
				arm_bones = arm_obj.bones
				
				writeInt(file, 1)
				
				writeInt(file, len(arm_obj.name))
				
				file.write(arm_obj.name)
													
				for bone in arm_bones.values():
				
					if not bone.hasParent():
						self.addBone(bone, None)
						
				print 'found %i bones' % (len(self.bones))
						
				writeUInt(file, len(self.bones))						
						
				for bone in self.bones:
				
					if bone.parent:
						writeInt(file, bone.parent.index)
					else:
						writeInt(file, -1)
					
					writeUInt(file, len(bone.name))
					
					file.write('%s' % (bone.name))
					
					writeUInt(file, 0)
					
					bone.matrix = Matrix(bone.matrix).resize4x4() #* self.ArmatureMat
																	
					writeFloat(file, bone.matrix[0][0])
					writeFloat(file, bone.matrix[0][1])
					writeFloat(file, bone.matrix[0][2])
					writeFloat(file, bone.matrix[1][0])
					writeFloat(file, bone.matrix[1][1])
					writeFloat(file, bone.matrix[1][2])
					writeFloat(file, bone.matrix[2][0])
					writeFloat(file, bone.matrix[2][1])
					writeFloat(file, bone.matrix[2][2])		
																				
					#bone.pos = Vector(bone.bone.head['ARMATURESPACE'])
					bone.pos = Vector([0, 0, 0])
					
					outmat = bone.bone.matrix['ARMATURESPACE']
										
					if bone.parent:
						parentmat = Matrix(bone.parent.bone.matrix['ARMATURESPACE'])
						parentmat = parentmat.invert()
						outmat *= parentmat
						
					bone.pos *= outmat * self.ArmatureMat
																																						
					writeFloat(file, bone.pos[0])
					writeFloat(file, bone.pos[1])
					writeFloat(file, bone.pos[2])
			
				writeInt(file, 3)
			
				file.write('%s' % ('end'))
					
		file.close()
		
class AnimationExport ( SkeletonExport ) :
	def callback(self, filename):
		self.export(filename)

	def export(self, filename):
	
		os.system('cls')

		file = open(filename, "wb")
		
		scn = Blender.Scene.GetCurrent()
		
		selobj = scn.objects.selected
		
		for obj in selobj:
		
			if obj.type == 'Armature':
			
				writeUInt(file, len('skeleton_anim'))
			
				file.write('skeleton_anim')
				
				self.ArmatureMat = obj.getMatrix()				
				
				writeUInt(file, 1)
				
				pose = obj.getPose()
				action = obj.getAction()
				
				writeUInt(file, len(action.name))
				
				file.write(action.name)
				
				arm_obj = obj.getData()
							
				arm_bones = arm_obj.bones
																				
				for bone in arm_bones.values():
				
					if not bone.hasParent():
						self.addBone(bone, None)
						
				writeUInt(file, len(self.bones))
				
				playtime = 0
				
				ipos = action.getAllChannelIpos()
				
				for boneobjs in self.bones:
					bone = boneobjs.bone
					if bone.name in ipos:
						boneIpo = ipos[bone.name]
						ipoCurve = boneIpo.getCurve( 'LocX' )
						if ipoCurve:
							for knot in ipoCurve.bezierPoints:
								if knot.pt[ 0 ] > playtime:
									playtime = knot.pt[ 0 ]
									
				playtime = int(round(playtime, 0))
				if playtime < 1:
					playtime = 1
					
				action.setActive(obj)
				
				matFixAnim = RotationMatrix( 90.0, 4, 'z' )				
									
				for bones in self.bones:
				
					realbone = bones.bone
													
					writeUInt(file, playtime)
				
					for time in range(1, playtime+1):
					
						obj.evaluatePose( time)
						
						writeUInt(file, time-1)
						
						posebone = pose.bones[ realbone.name]												
											
						posemat = Matrix(posebone.poseMatrix)
																							
						if bones.parent:
							posemat = posemat * Matrix(pose.bones[ bones.parent.name].poseMatrix).invert() 

						quat = posemat.toQuat()				
										
						posequat = Quaternion(quat[1], quat[2], quat[3], quat[0])
																						
						writeFloat(file, posequat[0])
						writeFloat(file, posequat[1])
						writeFloat(file, posequat[2])
						writeFloat(file, posequat[3])				
					
				writeUInt(file, playtime)
				
				for time in range(1, playtime+1):
				
					writeUInt(file, time-1)
					
					obj.evaluatePose( time)					
					
					for bone in self.bones:
					
						if not bone.parent:
						
							posebone = pose.bones[ bone.bone.name]
														
							pos = posebone.loc #* posebone.poseMatrix
													
							writeFloat(file, pos[0])
							writeFloat(file, pos[2])
							writeFloat(file, pos[1])				
												
	
				writeInt(file, 3)
				file.write('%s' % ('end'))
					
		file.close()
		
class MeshExport ( SkeletonExport ) :

	def __init__(self):
	
		self.weights = []
		self.faces = []
		self.meshes = []
		self.objects = []
		self.bones = []
		self.individualverts = []		
		
	def callback(self, filename):
		self.export(filename)		
		
	def export(self, filename):
	
		os.system('cls')

		file = open(filename, "wb")
		
		if not file:
			print 'something happened'
		
		writeUInt(file, len('mesh'))
		
		file.write('mesh')
		
		scn = Blender.Scene.GetCurrent()
		
		selobj = scn.objects.selected
		
		nummeshes = 0
		
		armafound = False
		
		for obj in selobj:
		
			if obj.type == 'Armature':
			
				armafound = True
							
				arm_obj = obj.getData()
							
				arm_bones = arm_obj.bones
																				
				for bone in arm_bones.values():
				
					if not bone.hasParent():
						self.addBone(bone, None)		
		
			if obj.type == 'Mesh':	
			
			
				nummeshes += 1
				self.meshes.append(obj.getData(False, True))
				self.objects.append(obj)
				
		if nummeshes == 0:
			Blender.Draw.PupMenu('Error%t|No Mesh selected, exit!')
			file.flush()
			file.close()
			return		
				
		print '%u' % (nummeshes)
				
		writeUInt(file, nummeshes)
		
		for i, mesh in enumerate(self.meshes):
		
			storeverts = mesh.verts[:]
			self.faces = mesh.faces					
			
			mesh.quadToTriangle()
			mesh.transform(self.objects[i].getMatrix(), True)
			if scale_10:
				mesh.transform(scalemat, True)
			mesh.recalcNormals()
					
			writeUInt(file, len(mesh.name))
			
			file.write(mesh.name)	
			
			writeInt(file, 0)
			
			if mesh.materials[0] != None:
				writeUInt(file, len(mesh.materials[0].name))
				file.write(mesh.materials[0].name)
			else:
				writeUInt(file, len('None'))
			
				file.write('None')
						
			writeUInt(file, len(mesh.verts))
			
			for vert in mesh.verts:
				
				writeFloat(file, vert.co.x)
				writeFloat(file, vert.co.y)
				writeFloat(file, vert.co.z)
				
			if armafound == True:
				
				groupnames = mesh.getVertGroupNames()
				
				writeUInt(file, len(groupnames))

				boneweights = []
				
				for group in groupnames:
								
					vertweightlist = mesh.getVertsFromGroup(group, 1)
							
					for vertweightpair in vertweightlist:
						if vertweightpair[1] > 1e-6:
							found = False
							for boweight in boneweights:
								if boweight.index == vertweightpair[0]:
									boweight.addWeight(vertweightpair[1], group)
									found = True
									break
							if found == False:
								newindex = BoneWeight(vertweightpair[0])
								newindex.addWeight(vertweightpair[1], group)
								boneweights.append(newindex)	
							
				for boweight in boneweights:
					boweight.normalizeWeights()
						
				for group in groupnames:
				
					for i, bone in enumerate(self.bones):
						if bone.name == group:
							writeUInt(file, i)
							break
					
					#print 'bone %u' % (i)
						
					realvertindices = []
					realvertweights = []
							
					for boweight in boneweights:
						for j, bogroup in enumerate(boweight.groups):
							if bogroup == group:
								realvertindices.append(boweight.index)
								realvertweights.append(boweight.weights[j])
								#print 'index: %u weight %f' % (boweight.index, boweight.weights[j])
										
					
					writeUInt(file, len(realvertweights))
					
					for j in range(0, len(realvertweights)):
									
						writeUInt(file, realvertindices[j])
						writeFloat(file, realvertweights[j])
			else:
				writeUInt(file, 0)
					
			for face in self.faces:
				
				for i in range(0, len(face.verts)):
					newvertex = VertexUV(face.verts[i].index, face.uv[i], face.index)
					found = False
					for vert in self.individualverts:
						if equalVector2(newvertex.uv, vert.uv) and newvertex.vertindex == vert.vertindex:
							found = True
							break
					if found == False:
						self.individualverts.append(newvertex)
						

			print 'indyverts: %u' % (len(self.individualverts))
			
			writeInt(file, 0)
			
			writeUInt(file, len(self.individualverts))
									
			for vert in self.individualverts:
				
				writeUInt(file, vert.vertindex)
				
				if mesh.vertexColors:
					writeChar(file, self.faces[vert.faceindex].col.a)
					writeChar(file, self.faces[vert.faceindex].col.b)
					writeChar(file, self.faces[vert.faceindex].col.g)
					writeChar(file, self.faces[vert.faceindex].col.r)								

				else:
					writeChar(file, 255)
					writeChar(file, 255)
					writeChar(file, 255)
					writeChar(file, 255)
					
				writeFloat(file, mesh.verts[vert.vertindex].no.x)
				writeFloat(file, mesh.verts[vert.vertindex].no.y)
				writeFloat(file, mesh.verts[vert.vertindex].no.z)
				
				if mesh.faceUV:
					writeFloat(file, vert.uv[0])
					writeFloat(file, vert.uv[1])
					writeFloat(file, 0.0)
					writeFloat(file, 0.0)
				else:

					writeFloat(file, 0.0)
					writeFloat(file, 0.0)
					writeFloat(file, 0.0)
					writeFloat(file, 0.0)
					
			writeUInt(file, len(self.faces))
					
			for face in self.faces:
			
				for i, vert in enumerate(self.individualverts):
					if vert.vertindex == face.v[0].index and equalVector2(face.uv[0], vert.uv):
						writeUInt(file, i)
						break
				for i, vert in enumerate(self.individualverts):							
					if vert.vertindex == face.v[1].index and equalVector2(face.uv[1], vert.uv):
						writeUInt(file, i)
						break
				for i, vert in enumerate(self.individualverts):							
					if vert.vertindex == face.v[2].index and equalVector2(face.uv[2], vert.uv):
						writeUInt(file, i)
						break
							
					
			
			#	writeUInt(file, face.v[0].index)
			#	writeUInt(file, face.v[1].index)
			#	writeUInt(file, face.v[2].index)	

			mesh.verts = storeverts
				
			
		writeInt(file, 3)
			
		file.write('%s' % ('end'))
			
		file.close()
			

def event(evt, val):
	if evt == Draw.ESCKEY:
		Draw.Exit()
		return
		
def button_event(evt):

	global toggle_scale_10, scale_10
	
	if evt == 0:
		export = MeshExport()
		filename = Blender.sys.makename(ext = ".brf")
		Blender.Window.FileSelector(export.callback, "Export brf", filename);	
	elif evt == 1:
		export = SkeletonExport()
		filename = Blender.sys.makename(ext = ".brf")
		Blender.Window.FileSelector(export.callback, "Export brf", filename);	
	elif evt == 2:
		export = AnimationExport()
		filename = Blender.sys.makename(ext = ".brf")
		Blender.Window.FileSelector(export.callback, "Export brf", filename);	
	elif evt == 3:
		Draw.Exit()
	elif evt == 4:
		toggle_scale_10 = 1 - toggle_scale_10
		scale_10 = toggle_scale_10
		Draw.Redraw(1)
	else:
		return
		
	#Draw.Redraw(1)
		
def draw():

	glClearColor(0.55, 0.6, 0.6, 1)
	glClear(BGL.GL_COLOR_BUFFER_BIT)
	
	glColor3f(0.2, 0.3, 0.3)
	rect(10, 180, 300, 180)
	
	glColor3f(0.8,.8,0.6)
	glRasterPos2i(20, 150)	
	Draw.Text("dsm-Exporter ", 'large')
	glRasterPos2i(20, 130)
	Draw.Text("for rigged mesh, skeleton has to be selected", 'small')
	
	exportmesh_button = Draw.Button("Export Mesh", 0, 220, 155, 75, 30, "Export Mesh")
	exportskel_button = Draw.Button("Export Skeleton", 1, 220, 120, 75, 30, "Export Skeleton")
	exportanim_button = Draw.Button("Export Animation", 2, 220, 85, 75, 30, "Export Animation")
	exit_button = Draw.Button("Exit", 3, 220, 25, 75, 30, "Exit")
	Draw.Toggle("Imported smd from OpenBrf", 4, 10, 90, 200, 30, toggle_scale_10, "scales by 0.1")
	
	glRasterPos2i(20, 70)
	if toggle_scale_10 :
		scale_msg = "object will be scaled by 0.1"
	else :
		scale_msg = "no scaling will be applied"
	Draw.Text(scale_msg)	
		
Draw.Register(draw, event, button_event)

def rect(x,y,width,height):
		glBegin(GL_LINE_LOOP)
		glVertex2i(x,y)
		glVertex2i(x+width,y)
		glVertex2i(x+width,y-height)
		glVertex2i(x,y-height)
		glEnd()