# LightWave .lwo file loader

from struct import *
from raytracer import Triangle, NormalVertex, Material

def read_int4(f):
	return unpack('>i', f.read(4))[0]

def read_int2(f):
	return unpack('>h', f.read(2))[0]

def read_float4(f):
	return unpack('>f', f.read(4))[0]

def read_string(f):
	res = ''
	b = f.read(1)
	l = 1
	while ( b != '\0' ):
		res += b
		b = f.read(1)
		l += 1
	if (l % 2 != 0):
		f.read(1)
		l += 1
	return (res,l)

def read_chunk(f):
	ID = f.read(4)
	if (ID == ''):
		return ('',0)
	size = read_int4(f)
	return (ID, size)

def read_subchunk(f):
	ID = f.read(4)
	size = read_int2(f)
	return (ID, size)

def read_lwo(filename):
	points = []
	faces = []
	tags = []
	surfaces = []
	f = file(filename, 'rb')
	(ID,size) = read_chunk(f)
	form = f.read(4)
	if (ID != 'FORM' or form != 'LWOB'):
		print 'not lwo file'
		return
	(ID,size) = read_chunk(f)
	while (ID != ''):
		#print ID,size
		if (ID == 'PNTS'):
			while (size > 0):
				p1 = read_float4(f)
				p2 = read_float4(f)
				p3 = read_float4(f)
				points.append((p1,p2,p3))
				size -= 12
		elif (ID == 'SRFS'):
			while (size > 0):
				(s,l) = read_string(f)
				size -= l
				tags.append(s)
		elif (ID == 'POLS'):
			while (size > 0):
				vertex_count = read_int2(f)
				size -= 2
				inds = []
				for i in range(vertex_count):
					index = read_int2(f)
					inds.append(index)
					size -= 2
				surf = read_int2(f)
				size -= 2
				faces.append((inds[0], inds[1], inds[2], surf))
				i = 0
				while (vertex_count > 3):
					faces.append((inds[0], inds[2+i], inds[3+i], surf))
					vertex_count -= 1
					i += 1
		elif (ID == 'SURF'):
			(name,l) = read_string(f)
			size -= l
			surf = {}
			while (size > 0):
				(subID,subsize) = read_subchunk(f)
				size -= 6
				if (subID == 'COLR'):
					col = f.read(subsize)
					surf['color'] = (unpack('BBB',col[:3]))
				else:
					print "Warning:", subID,"("+str(subsize),"B) ignored"
					f.read(subsize)
				size -= subsize
			surfaces.append(surf)
		else:
			print "Warning:", ID,"("+str(size),"B) ignored"
			f.read(size)
		(ID,size) = read_chunk(f)
	return (points, faces, tags, surfaces)

def LoadLightwaveLwoFile(rt, filename, defmat, smooth, scale=(1,1,1), trans=(0,0,0)):
	if (type(scale) == float or type(scale) == int):
		scale = (scale,)*3
	(points, faces, tags, surfaces) = read_lwo(filename)

	vertices = []
	normals = []
	vertex_face_num = []
	materials = []

	for surf in surfaces:
		mat = Material(colour=tuple(float(x)/255. for x in surf['color']))
		materials.append(mat)

	for point in points:
		v0 = scale[0]*point[0] + trans[0]
		v1 = scale[1]*point[1] + trans[1]
		v2 = scale[2]*point[2] + trans[2]
		vertices.append(NormalVertex((v2,v1,v0)))
		normals.append([0.,0.,0.])
		vertex_face_num.append(0)

	for f in faces:
		v = [vertices[x] for x in f[0:3]]
		matidx = f[3]-1
		face = Triangle(v[0], v[1], v[2], materials[matidx])
		n = face.getNormal()
		for x in f[0:3]:
			for i in range(3):
				normals[x][i] += n[i]
			vertex_face_num[x] += 1
		if (smooth):
			face.setSmooth()
		rt.addshape(face)

	# interpolate normals at vertices
	num = 0
	for vf in vertex_face_num:
		if (vertex_face_num[num] > 0):
			for i in range(3):
				normals[num][i] /= vertex_face_num[num]
		vertices[num].setNormal(tuple(normals[num]))
		num += 1