# LightWave Object File exporter for Blender 2.69+
# Copyright (C) 2014  Morshidul Chowdhury

# 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/>.

bl_info = {
    "name": "Export LightWave Objects",
    "author": "Morshidul Chowdhury (iPLEOMAX)",
    "version": (1, 1),
    "blender": (2, 6, 9),
    "location": "File > Export > LightWave Object (.lwo)",
    "description": "Exports selected objects to LightWave (.lwo) format",
    "warning": "",
    "category": "Import-Export"
}

# http://pleo.code5gaming.com/
# http://blenderartists.org/forum/member.php?220524-iPLEOMAX

# This is a complete remake of Anthony D'Agostino's LightWave plugin.

# Features in V1.1 (current):
# - Added better material export data
# - Added support for multi-vertex weigths
# - Added support for multi-shape keys
# - Added support for large meshes.
# - Surfaces are exported as 'smooth' by default, you need to set flat ones by yourself.
# - Fixed other bugs that would cause failed imports.
#
# Features in V1.0 (old):
# - Creates an extra '.lwo' option in export menu.
# - Links all image paths to the lightwave file.
# - Supports multi-surfaces (although there are some incompatibility issues)
# - Supports multi-meshes.
# - Supports multi-uv layers.
# - Supports multi-vertex colors.

import bpy
import io
import struct
import time
import operator
from functools import reduce

LWO_VCOLOR_MATERIAL = "LWO.VColorMaterial"
LWO_DEFAULT_MATERIAL = "LWO.DefaultMaterial"
LWO_ACTIVE_UV = "LWO.ActiveUV"
LWO_ACTIVE_VC = "LWO.ActiveVCol"
LWO_ACTIVE_VG = "LWO.ActiveVGroup"
LWO_ACTIVE_EM = "LWO.ActiveMorph"

# Function from Scorpius's Plugin (Terminates the string and keeps the length even)
def lwopad(str):
	if len(str) % 2: str += '\0' #odd
	else: str += '\0\0' #even
	return str

# Function from Scorpius's Plugin (Decides whether to occupy a 4byte or 2byte space in file)
def generate_vx(index):
	if index < 0xFF00:
		value = struct.pack(">H", index)                 # 2-byte index
	else:
		value = struct.pack(">L", index | 0xFF000000)    # 4-byte index
	return value
	
def get_materials(meshes):
	print("\t- Collecting materials..")
	materials = []
	for mesh in meshes:
		if mesh.materials:
			for mat in mesh.materials:
				if mat:
					if not mat.name in materials:
						materials.append(mat.name)
		elif mesh.vertex_colors:
			if not LWO_VCOLOR_MATERIAL in materials:
				materials.append(LWO_VCOLOR_MATERIAL)
		else:
			if not LWO_DEFAULT_MATERIAL in materials:
				materials.append(LWO_DEFAULT_MATERIAL)
	return materials
	
def get_tags(materials):
	print("\t- Generating tags..")
	data = io.BytesIO()
	for mat in materials:
		data.write(bytes(lwopad(mat), 'UTF-8'))
	return data.getvalue()

def get_layr(mesh, index, pivot):
	print("\t- Generating layer for mesh '" + mesh.name + "'")
	data = io.BytesIO()
	data.write(struct.pack(">h", index))
	data.write(struct.pack(">h", 0))
	data.write(struct.pack(">fff", pivot.x, pivot.z, pivot.y))
	data.write(bytes(lwopad(mesh.name), 'UTF-8'))
	return data.getvalue()
	
def get_pnts(mesh):
	print("\t- Generating vertices for mesh '" + mesh.name + "'")
	data = io.BytesIO()
	for v in mesh.vertices:
		x, y, z = v.co
		data.write(struct.pack(">fff", x, z, y))
	return data.getvalue()
	
def get_bbox(mesh):
	print("\t- Generating bounding box for mesh '" + mesh.name + "'")
	data = io.BytesIO()
	if mesh.vertices:
		nv = [v.co for v in mesh.vertices]
		xx = [ co[0] for co in nv ]
		yy = [ co[1] for co in nv ]
		zz = [ co[2] for co in nv ]
	else:
		xx = yy = zz = [0.0,]
	
	data.write(struct.pack(">6f", min(xx), min(zz), min(yy), max(xx), max(zz), max(yy)))
	return data.getvalue()
	
def get_vmad_uv(mesh, uvlayer, activelayername):
	print("\t\tUV Layer: '" + uvlayer.name + "'")
	data = io.BytesIO()
	data.write(b"TXUV")
	data.write(struct.pack(">H", 2))
	
	if uvlayer.name == activelayername:
		uvname = lwopad(LWO_ACTIVE_UV)
		data.write(bytes(uvname, 'UTF-8'))
	else:
		uvname = lwopad(uvlayer.name)
		data.write(bytes(uvname, 'UTF-8'))
		
	for face in mesh.polygons:
		for vert, loop in zip(face.vertices, face.loop_indices):
			uv = uvlayer.data[loop].uv
			#data.write(struct.pack(">H", vert))
			#data.write(struct.pack(">H", face.index))
			data.write(generate_vx(vert))
			data.write(generate_vx(face.index))
			data.write(struct.pack(">ff", uv.x, uv.y))
	return data.getvalue()

def get_vmad_vc(mesh, vclayer, activevcname):
	print("\t\tVCOL Layer: '" + vclayer.name + "'")
	data = io.BytesIO()
	data.write(b"RGB ")
	data.write(struct.pack(">H", 3))
	
	if vclayer.name == activevcname:
		vcname = lwopad(LWO_ACTIVE_VC)
		data.write(bytes(vcname, 'UTF-8'))
	else:
		vcname = lwopad(vclayer.name)
		data.write(bytes(vcname, 'UTF-8'))
		
	for face in mesh.polygons:
		for vert, loop in zip(face.vertices, face.loop_indices):
			color =  vclayer.data[loop].color
			#data.write(struct.pack(">H", vert))
			#data.write(struct.pack(">H", face.index))
			data.write(generate_vx(vert))
			data.write(generate_vx(face.index))
			data.write(struct.pack(">fff", color.r, color.g, color.b))
	return data.getvalue()

def get_vmap_vg(mesh, vglayer, activevgname):
	print("\t\tVertex Group Layer: '" + vglayer.name + "'")
	data = io.BytesIO()
	data.write(b"WGHT")
	data.write(struct.pack(">H", 1))
	
	if vglayer.name == activevgname:
		vcname = lwopad(LWO_ACTIVE_VG)
		data.write(bytes(vcname, 'UTF-8'))
	else:
		vgname = lwopad(vglayer.name)
		data.write(bytes(vgname, 'UTF-8'))
	
	for v in mesh.vertices:
		w = 0.0
		try:
			w =  vglayer.weight(v.index)
		except:
			pass
		data.write(generate_vx(v.index))
		data.write(struct.pack(">f", w))
		
	return data.getvalue()

def get_vmap_em(mesh, emlayer, activeemname):
	print("\t\tShape Keys Layer: '" + emlayer.name + "'")
	data = io.BytesIO()
	data.write(b"MORF")
	data.write(struct.pack(">H", 3))
	
	if emlayer.name == LWO_ACTIVE_EM:
		emname = lwopad(LWO_ACTIVE_EM)
		data.write(bytes(emname, 'UTF-8'))
	else:
		emname = lwopad(emlayer.name)
		data.write(bytes(emname, 'UTF-8'))
	
	for v in mesh.vertices:
		x, y, z = emlayer.data[v.index].co - v.co
		data.write(generate_vx(v.index))
		data.write(struct.pack(">fff", x, z, y))
	
	return data.getvalue()
	
def get_pols(mesh):
	print("\t- Generating polygons for mesh '" + mesh.name + "'")
	data = io.BytesIO()
	data.write(b"FACE")
	for face in mesh.polygons:
		data.write(struct.pack(">H", face.loop_total))
		for v in range(face.loop_total - 1, -1, -1):
			data.write(generate_vx(face.vertices[v]))
	return data.getvalue()
	
def get_ptag(mesh, materials):
	print("\t- Generating per poly materials for mesh '" + mesh.name + "'")
	data = io.BytesIO()
	data.write(b"SURF")
	for poly in mesh.polygons:
		if mesh.materials:
			matindex = poly.material_index
			matname = mesh.materials[matindex].name
			surfindex = materials.index(matname)
			
			data.write(generate_vx(poly.index))
			data.write(struct.pack(">H", surfindex)) 
		else:
			data.write(generate_vx(poly.index))
			data.write(struct.pack(">H", 0)) 
	return data.getvalue()
	
def get_clip(image, clipid):
	print("\t\tImage: '" + image.name + "'")
	data = io.BytesIO()
	path = image.filepath
	path = path[0:2] + path.replace("\\", "/")[3:]
	image = lwopad(path)
	data.write(struct.pack(">L", clipid))
	data.write(b"STIL")
	data.write(struct.pack(">H", len(image)))
	data.write(bytes(image, 'UTF-8'))
	return data.getvalue()

def get_imap():
	data = io.BytesIO()
	data.write(struct.pack(">H", 0))
	data.write(b"CHAN")
	data.write(struct.pack(">H", 4))
	data.write(b"COLR")
	data.write(b"OPAC")
	data.write(struct.pack(">H", 8))
	data.write(struct.pack(">H", 0))
	data.write(struct.pack(">f", 1.0))
	data.write(struct.pack(">H", 0))
	data.write(b"ENAB")
	data.write(struct.pack(">HH", 2, 1))
	data.write(b"NEGA")
	data.write(struct.pack(">HH", 2, 0))
	data.write(b"AXIS")
	data.write(struct.pack(">HH", 2, 1))
	return data.getvalue()

def create_custom_surf(mat):
	print("\t\tCustom Surface: '" + mat.name + "'")
	data = io.BytesIO()
	data.write(bytes(lwopad(mat.name), 'UTF-8'))
	
	R = mat.diffuse_color.r
	G = mat.diffuse_color.g
	B = mat.diffuse_color.b
	Diff = mat.diffuse_intensity
	Lumi = mat.emit
	Spec = mat.specular_intensity

	data.write(b"COLR")
	data.write(struct.pack(">H", 0))
	
	data.write(b"COLR")
	data.write(struct.pack(">H", 14))
	data.write(struct.pack(">fffH", R, G, B, 0))

	data.write(b"DIFF")
	data.write(struct.pack(">H", 6))
	data.write(struct.pack(">fH", Diff, 0))

	data.write(b"LUMI")
	data.write(struct.pack(">H", 6))
	data.write(struct.pack(">fH", Lumi, 0))

	data.write(b"SPEC")
	data.write(struct.pack(">H", 6))
	data.write(struct.pack(">fH", Spec, 0))

	data.write(b"GLOS")
	data.write(struct.pack(">H", 6))
	gloss = 50 / (255/2.0)
	gloss = round(gloss, 1)
	data.write(struct.pack(">fH", gloss, 0))
	
	data.write(b"SMAN")
	data.write(struct.pack(">H", 4))
	data.write(struct.pack(">f", 1.5))
	
	if mat and mat.node_tree and mat.node_tree.nodes:
		for node in mat.node_tree.nodes:
			if node:
				if node.type == 'TEX_IMAGE':
					data.write(b"BLOK")
					data.write(struct.pack(">H", 104))
					imap = get_imap()
					data.write(b"IMAP")
					data.write(struct.pack(">H", len(imap)))
					data.write(imap)
					
					data.write(b"IMAG")
					data.write(struct.pack(">HH", 2, 1))
					data.write(b"PROJ")
					data.write(struct.pack(">HH", 2, 5))
					data.write(b"VMAP")
					uvname = lwopad(LWO_ACTIVE_UV)
					data.write(struct.pack(">H", len(uvname)))
					data.write(bytes(uvname, 'UTF-8'))
	
	data.write(b"VCOL")
	data.write(struct.pack(">H", 34))
	data.write(struct.pack(">fH4s", 1.0, 0, b"RGB ")) 
	data.write(bytes(lwopad(LWO_ACTIVE_VC), 'UTF-8'))
	return data.getvalue()

def create_vcolor_surf():
	print("\t\Vertex Color Surface")
	data = io.BytesIO()
	data.write(bytes(lwopad(LWO_VCOLOR_MATERIAL), 'UTF-8'))
	
	R = 1.0
	G = 1.0
	B = 1.0
	
	data.write(b"COLR")
	data.write(struct.pack(">H", 0))
	
	data.write(b"COLR")
	data.write(struct.pack(">H", 14))
	data.write(struct.pack(">fffH", R, G, B, 0))

	data.write(b"DIFF")
	data.write(struct.pack(">H", 6))
	data.write(struct.pack(">fH", 0.8, 0))

	data.write(b"LUMI")
	data.write(struct.pack(">H", 6))
	data.write(struct.pack(">fH", 0.0, 0))

	data.write(b"SPEC")
	data.write(struct.pack(">H", 6))
	data.write(struct.pack(">fH", 0.0, 0))

	data.write(b"SMAN")
	data.write(struct.pack(">H", 4))
	data.write(struct.pack(">f", 1.5))
	
	data.write(b"VCOL")
	data.write(struct.pack(">H", 34))
	data.write(struct.pack(">fH4s", 1.0, 0, b"RGB ")) 
	data.write(bytes(lwopad(LWO_ACTIVE_VC), 'UTF-8'))
	
	return data.getvalue()	

def create_surf(matname):
	print("\t\tDefault Surface: '" + matname + "'")
	data = io.BytesIO()
	data.write(bytes(lwopad(matname), 'UTF-8'))
	
	R = 0.9
	G = 0.9
	B = 0.9
	
	data.write(b"COLR")
	data.write(struct.pack(">H", 0))
	
	data.write(b"COLR")
	data.write(struct.pack(">H", 14))
	data.write(struct.pack(">fffH", R, G, B, 0))

	data.write(b"DIFF")
	data.write(struct.pack(">H", 6))
	data.write(struct.pack(">fH", 0.8, 0))

	data.write(b"LUMI")
	data.write(struct.pack(">H", 6))
	data.write(struct.pack(">fH", 0.0, 0))

	data.write(b"SPEC")
	data.write(struct.pack(">H", 6))
	data.write(struct.pack(">fH", 0.4, 0))

	data.write(b"GLOS")
	data.write(struct.pack(">H", 6))
	gloss = 50 / (255/2.0)
	gloss = round(gloss, 1)
	data.write(struct.pack(">fH", gloss, 0))
	
	data.write(b"SMAN")
	data.write(struct.pack(">H", 4))
	data.write(struct.pack(">f", 1.5))
	
	return data.getvalue()
	
def get_surf(matname):
	matid = bpy.data.materials.find(matname)
	if matid != -1:
		return create_custom_surf(bpy.data.materials[matid])
	else:
		if matname == LWO_VCOLOR_MATERIAL:
			return create_vcolor_surf()
	return create_surf(matname)
	
def get_header(chunks):
	print("\t- Generating headers..")
	data = io.BytesIO()
	chunk_sizes = map(len, chunks)
	chunk_sizes = reduce(operator.add, chunk_sizes)
	form_size = chunk_sizes + len(chunks)*8 + len("FORM")
	data.write(b"FORM")
	data.write(struct.pack(">L", form_size))
	data.write(b"LWO2")
	return data.getvalue()
	
def write_lwo(filepath):
	
	print("LWO exporter started..")
	start = time.time()
	
	try:
		bpy.ops.object.mode_set()
	except:
		pass
	
	if not bpy.context.selected_objects:
		print("Please select an object before exporting.")
		return
	
	meshes = []
	pivots = {}
	vgroups = {}
	activeshapekey = {}
	shapekeys = {}
	
	print("\t- Collecting meshes..")
	for obj in bpy.context.selected_objects:
		if obj.type == 'MESH':
			mesh = obj.to_mesh(bpy.context.scene, True, 'PREVIEW')
			mesh.transform(obj.matrix_world)
			mesh.calc_normals()
			meshes.append(mesh)
			pivots[mesh] = obj.location
			if obj.vertex_groups:
				vgroups[mesh] = obj.vertex_groups
			if not obj.modifiers and obj.active_shape_key:
				activeshapekey[mesh] = obj.active_shape_key.name
				shapekeys[mesh] = obj.data.shape_keys
			print("\t\tMesh: " + mesh.name)
	del obj

	if not meshes:
		print("\tThere are no meshes in your selection.")
		return
	
	materials = get_materials(meshes)
	tags = get_tags(materials)
	
	chunks = [tags]
	meshdata = io.BytesIO()
	meshindex = 0
	
	for mesh in meshes:
		layr = get_layr(mesh, meshindex, pivots[mesh])
		meshindex += 1
		meshdata.write(b"LAYR")
		meshdata.write(struct.pack(">L", len(layr)))
		meshdata.write(layr)
		chunks.append(layr)
		
		pnts = get_pnts(mesh)
		meshdata.write(b"PNTS")
		meshdata.write(struct.pack(">L", len(pnts)))
		meshdata.write(pnts)
		chunks.append(pnts)
		
		bbox = get_bbox(mesh)
		meshdata.write(b"BBOX")
		meshdata.write(struct.pack(">L", len(bbox)))
		meshdata.write(bbox)
		chunks.append(bbox)
		
		pols = get_pols(mesh)
		meshdata.write(b"POLS")
		meshdata.write(struct.pack(">L", len(pols)))
		meshdata.write(pols)
		chunks.append(pols)
		
		ptag = get_ptag(mesh, materials)
		meshdata.write(b"PTAG")
		meshdata.write(struct.pack(">L", len(ptag)))
		meshdata.write(ptag)
		chunks.append(ptag)

		if mesh.uv_layers:
			print("\t- Generating UV layers for mesh '" + mesh.name + "'")
			for uvlayer in mesh.uv_layers:
				vmad = get_vmad_uv(mesh, uvlayer, mesh.uv_layers.active.name)
				meshdata.write(b"VMAD")
				meshdata.write(struct.pack(">L", len(vmad)))
				meshdata.write(vmad)
				chunks.append(vmad)
		
		clipid = 1
		
		print("\t- Generating Images for mesh '" + mesh.name + "'")
		for mat in mesh.materials:
			if mat and mat.node_tree and mat.node_tree.nodes:
				for node in mat.node_tree.nodes:
					if node:
						if node.type == 'TEX_IMAGE':
							clip = get_clip(node.image, clipid)
							meshdata.write(b"CLIP")
							meshdata.write(struct.pack(">L", len(clip)))
							meshdata.write(clip)
							chunks.append(clip)
							clipid += 1
		
		if mesh.vertex_colors:
			print("\t- Generating Vertex colors for mesh '" + mesh.name + "'")
			for vclayer in mesh.vertex_colors:
				vmad = get_vmad_vc(mesh, vclayer, mesh.vertex_colors.active.name)
				meshdata.write(b"VMAD")
				meshdata.write(struct.pack(">L", len(vmad)))
				meshdata.write(vmad)
				chunks.append(vmad)
		
		if mesh in vgroups:
			if len(vgroups[mesh]):
				print("\t- Generating Vertex groups for mesh '" + mesh.name + "'")
				for vglayer in vgroups[mesh]:
					vmap = get_vmap_vg(mesh, vglayer,  vgroups[mesh].active.name)
					meshdata.write(b"VMAP")
					meshdata.write(struct.pack(">L", len(vmap)))
					meshdata.write(vmap)
					chunks.append(vmap)
		
		if mesh in activeshapekey:
			if mesh in shapekeys:
				print("\t- Generating Key Blocks for mesh '" + mesh.name + "'")
				for emlayer in shapekeys[mesh].key_blocks:
					vmap = get_vmap_em(mesh, emlayer, activeshapekey[mesh])
					meshdata.write(b"VMAP")
					meshdata.write(struct.pack(">L", len(vmap)))
					meshdata.write(vmap)
					chunks.append(vmap)
					
		bpy.data.meshes.remove(mesh)
	
	print("\t- Generating Surfaces")
	for mat in materials:
		surf = get_surf(mat)
		meshdata.write(b"SURF")
		meshdata.write(struct.pack(">L", len(surf)))
		meshdata.write(surf)
		chunks.append(surf)
	
	header = get_header(chunks)
	
	if not filepath.lower().endswith('.lwo'): filepath += '.lwo'
	file = open(filepath, 'wb')
	
	file.write(header)
	
	file.write(b"TAGS")
	file.write(struct.pack(">L", len(tags)))
	file.write(tags)
	
	file.write(meshdata.getvalue())
	meshdata.close()
	
	file.close()
	
	print("LWO exporter ended in %.03f seconds!" % (time.time() - start))
	
from bpy_extras.io_utils import ExportHelper
from bpy.props import StringProperty, BoolProperty, EnumProperty

class EXPORT_OT_lwo(bpy.types.Operator, ExportHelper):
	bl_idname= "export_scene.lwo"
	bl_label= "Export LWO"
	bl_description= "Export a LightWave Object file"
	bl_options= {'REGISTER'}
	
	filename_ext = ".lwo"
	filter_glob = StringProperty(default = "*.lwo", options = {'HIDDEN'})
	
	def execute(self, context):
		write_lwo(self.filepath)
		return {'FINISHED'}
	
	def invoke(self, context, event):
		wm = context.window_manager
		wm.fileselect_add(self)
		return {'RUNNING_MODAL'}
	
def menu_func(self, context):
    self.layout.operator(EXPORT_OT_lwo.bl_idname, text="LightWave Object (.lwo)")


def register():
    bpy.utils.register_module(__name__)
    bpy.types.INFO_MT_file_export.append(menu_func)


def unregister():
    bpy.utils.unregister_module(__name__)
    bpy.types.INFO_MT_file_export.remove(menu_func)

if __name__ == "__main__":
    register()