# 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 = []

	(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