#!BPY
#HL2 SMD Importer

"""
Name: 'Half-Life 2 (.smd) Importer'
Blender: 242
Group: 'Import'
Tooltip: 'Import mesh/animation from Half-Life 2 File'
"""

__author__ = 'Joseph Gilbert'
__url__ = ('blender', 'elysiun','My script homepage, http://www.blender.org',)
__version__ = '1.0'
__bpydoc__ = """This is a Half-Life 2 Importer (.smd).
Usage:  Alt+P to run
Notes:  Designed to import either the skeletal mesh/rest frame file or an animation sequence
"""

# --------------------------------------------------------------------------
# HL2 SMD Importer
# --------------------------------------------------------------------------
# ***** 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
from Blender import Mathutils
from Blender import Armature
from Blender import Object
from Blender import Scene
from Blender import Mesh
import math

print "--------------- Script Starting ---------------"

FILE = 'C:/Documents and Settings/jgilbert/Desktop/H2exporter/SMD_mesh.smd'
MODEL_NAME = 'SMD_model'

#constants
Header_file = 'version 1'
Header_armature = 'nodes'
Header_animation = 'skeleton'
Header_mesh = 'triangles'
Block_end = 'end'

def parse_file(file):
	'''Parse the file'''
	myfile = open(FILE, 'r')
	try:
		file_data = {}
		parse_header(myfile)
		for x in range(3):
			data, type = check_next_block(myfile)
			file_data[type] = data
	finally:
		myfile.close()
	return file_data

def check_next_block(myfile):
	'''check which block comes next in the file and parse it'''
	type = None
	line = myfile.readline()
	if line.strip() == Header_armature:
		data = parse_armature(myfile)
		type = 'armature'
	elif line.strip() == Header_animation:
		data = parse_animation(myfile)
		type = 'animation'
	elif line.strip() == Header_mesh:
		data = parse_mesh(myfile)
		type = 'mesh'
	else:
		raise AttributeError, 'File format is not correct'
	return data, type
	
def parse_header(myfile):
	'''Parse the header line'''
	line = myfile.readline()
	if line.strip() != Header_file:
		raise AttributeError, 'This is not a valid .smd file'
	
def parse_armature(myfile):
	'''Parse armature data'''
	armature_data = {}
	line = myfile.readline()
	while line.strip() != Block_end:
		data = line.strip().split('\"')
		armature_data[int(data[0].strip())] = {'name': data[1], \
			'parent': int(data[2].strip())}
		line = myfile.readline()
	return armature_data
	
def parse_mesh(myfile):
	'''Parse the mesh data'''
	mesh_data = []
	triangle = {}
	line = myfile.readline()
	while line.strip() != Block_end:
		data = line.strip().split()
		if len(data) == 1:
			if triangle != {}:
				mesh_data.append(triangle.copy())
				triangle['texture'] = data[0]
				triangle['verts'] = []
			else:
				triangle['texture'] = data[0]
				triangle['verts'] = []
		else:
			vertex = {}
			vertex['bparent'] = data[0]
			vertex['loc'] = Mathutils.Vector(float(data[1]), \
			  float(data[2]), float(data[3]))
			vertex['norm'] = Mathutils.Vector(float(data[4]), \
			  float(data[5]), float(data[6]))
			vertex['tex'] = Mathutils.Vector(float(data[7]), \
			  float(data[8]))
			num_weights = int(data[9])
			vertex['weights'] = {}
			for x in range(num_weights):
				bone = int(data[10+(x*2)])
				weight = float(data[11+(x*2)])
				vertex['weights'][bone] = weight
			triangle['verts'].append(vertex)
		line = myfile.readline()
	mesh_data.append(triangle) #add last
	return mesh_data
	
def parse_animation(myfile):
	'''Parse the animation data'''
	animation_data = {}
	frame_data = {}
	frame_number = None
	line = myfile.readline()
	while line.strip() != Block_end:
		data = line.strip().split()
		if data[0] == 'time':
			if frame_data != {}:
				animation_data[frame_number] = frame_data
				frame_data = {}
				frame_number = int(data[1])
			else:
				frame_number = int(data[1])
		else:
			#reverse y and z
			loc = Mathutils.Vector(float(data[1]), float(data[2]), \
			  float(data[3]))
			rot = Mathutils.Euler(float(data[4]), float(data[5]), \
			  float(data[6]))
			frame_data[int(data[0])] = {'Loc':loc, 'Rot':rot}
		line = myfile.readline()
	animation_data[frame_number] = frame_data #add final
	return animation_data

def load_file_data(data):
	'''load the file data into blender'''
	create_resting_armature(file_data['armature'], file_data['animation'][0])
	create_mesh(file_data['mesh'])
	
def create_resting_armature(armature_data, frame_0_data):
	'''creates the armature in it's rest pose'''
	myscene = Scene.GetCurrent()
	myobj = Object.New('Armature', MODEL_NAME)
	myobj.setDrawMode(64) #xray bones
	myarm = Armature.Armature(MODEL_NAME)
	myobj.link(myarm)
	myscene.link(myobj)
	myarm.makeEditable()
	#create bones
	for bone in armature_data.keys():
		mybone = Armature.Editbone()
		mybone.name = armature_data[bone]['name']
		myarm.bones[mybone.name] = mybone
	#now fix the parenting
	for bone in armature_data.keys():
		ebone = myarm.bones[armature_data[bone]['name']]
		parent = armature_data[bone]['parent']
		if parent != -1:
			parent_name = armature_data[parent]['name']
			ebone.parent = myarm.bones[parent_name]
	bone_mapping = {}
	for key, value in armature_data.items():
		bone_mapping[value['name']] = {'key':key, 'matrix': None, 'head': None, 'tail': None}
	#now fix the position data
	#calculate all matrices
	for bone in armature_data.keys():
		parent = armature_data[bone]['parent']
		if parent == -1:
			ebone = myarm.bones[armature_data[bone]['name']]
			bone_mapping[ebone.name]['matrix'] = generate_local_bone_matrix( \
			  frame_0_data[bone]['Loc'], frame_0_data[bone]['Rot'])
			calculate_matrices(myarm, ebone, frame_0_data, bone_mapping)
	#update them in the ebones
	for bone_name, data in bone_mapping.items():
		myarm.bones[bone_name].matrix = data['matrix']	
	bone_mapping_w_headtail = bone_mapping.copy()
	for bone_name, data in bone_mapping.items():
		bone_mapping_w_headtail[bone_name]['head'] = myarm.bones[bone_name].head
		bone_mapping_w_headtail[bone_name]['tail'] = myarm.bones[bone_name].tail
	for bone_name, data in bone_mapping.items():
		parent = myarm.bones[bone_name].parent
		if parent:
			myarm.bones[bone_name].head = bone_mapping_w_headtail[parent.name]['tail']
	myarm.update()	
	
def generate_local_bone_matrix(loc, rot):
	'''generate a localspace matrix from file data'''
	new_rot = convert_Euler(rot)
	new_loc = Mathutils.Vector(loc).resize4D()
	transform = new_rot.toMatrix().resize4x4()
	transform[3] = new_loc
	return transform
	
def convert_Euler(euler):
	'''convert a euler to the appropriate type'''
	x = (180/math.pi)*euler[0]
	y = (180/math.pi)*euler[1]
	z = (180/math.pi)*euler[2]
	return Mathutils.Euler(x,y,z)	
		
def calculate_matrices(myarm, ebone, frame_0_data, bone_mapping):
	children = []
	if ebone.parent:
		bone_id = bone_mapping[ebone.name]['key']
		local_matrix = generate_local_bone_matrix(frame_0_data[bone_id]['Loc'] , \
		  rot = frame_0_data[bone_id]['Rot'])
		parent_matrix = bone_mapping[ebone.parent.name]['matrix']
		bone_mapping[ebone.name]['matrix'] = local_matrix * parent_matrix
	for bone in myarm.bones.values():
		if bone.parent:
			if bone.parent.name == ebone.name:
				children.append(bone)
	for child in children:
		calculate_matrices(myarm, child, frame_0_data, bone_mapping)
		
def create_mesh(mesh_data):
	'''create the mesh from the imported data'''
	myscene = Scene.GetCurrent()
	myobj = Object.New('Mesh', MODEL_NAME)
	myobj.setDrawType(2) #draw wire mesh
	mymesh= Mesh.New(MODEL_NAME)
	myobj.link(mymesh)
	myscene.link(myobj)
	verts = []
	faces = []
	for tris in mesh_data:
		for vert in tris['verts']:
			verts.append(vert['loc'])
	#get unique vertices
	unique_verts = []
	for v in verts:
		if v not in unique_verts:
			unique_verts.append(v)
	mymesh.verts.extend(unique_verts)
	#build vert data map
	unique_vert_mapping = {}
	for x in range(len(unique_verts)):
		unique_vert_mapping[x] = {'loc':unique_verts[x]}
	#this is a little slow - we need unique vert data by vert index in the mesh
	unique_vert_map = unique_vert_mapping.copy()
	for key, value in unique_vert_mapping.items():
		found = 0
		for tris in mesh_data:
			for vert in tris['verts']:
				if value['loc'] == vert['loc']:
					unique_vert_map[key]['norm'] = vert['norm']
					unique_vert_map[key]['tex'] = vert['tex']
					unique_vert_map[key]['weights'] = vert['weights']
					unique_vert_map[key]['bparent'] = vert['bparent']
					unique_vert_map[key]['texture'] = tris['texture']
					found = 1
					break
			if found:
				break
	unique_vert_mapping = unique_vert_map.copy()
	del unique_vert_map
	#set normals per vertex
	for x in range(len(mymesh.verts)):
		mymesh.verts[x].no = unique_vert_mapping[x]['norm']
	#build faces
	for tris in mesh_data:
		face = []
		for vert in tris['verts']:
			index = unique_verts.index(vert['loc'])
			face.append(index)
		faces.append(face[:])
	mymesh.faces.extend(faces)
	
if __name__=='__main__':
	file_data = parse_file(FILE)
	load_file_data(file_data)