# LightWave .lwo file loader

from math import *

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 'unknown format'

		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

				if surf < 0:

					# detail polygons

					surf = abs(surf)

					count = read_int2(f)

					size -= 2

					# ... FIXME

				#print size, vertex_count

				if vertex_count >= 3:

					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]))

				elif (subID == 'FLAG'):

					flags = read_int2(f)

					surf['luminous'] = (flags >> 0) & 1;

					surf['outline'] = (flags >> 1) & 1;

					surf['smooth'] = (flags >> 2) & 1;

					surf['color-highlights'] = (flags >> 3) & 1;

					surf['color-filter'] = (flags >> 4) & 1;

					surf['opaque-edge'] = (flags >> 5) & 1;

					surf['transparent-edge'] = (flags >> 6) & 1;

					surf['sharp-terminator'] = (flags >> 7) & 1;

					surf['double-sided'] = (flags >> 8) & 1;

					surf['additive'] = (flags >> 9) & 1;

				# Base Shading (float)

				elif (subID == 'VDIF'):

					surf['diffuse'] = read_float4(f)

				elif (subID == 'VSPC'):

					surf['specular'] = read_float4(f)

				elif (subID == 'VRFL'):

					surf['reflection'] = read_float4(f)

				elif (subID == 'VTRN'):

					surf['transparency'] = read_float4(f)

				# Base Shading (short)

				elif (subID == 'DIFF'):

					if not surf.has_key('diffuse'):

						surf['diffuse'] = read_int2(f)/255.

				elif (subID == 'SPEC'):

					if not surf.has_key('specular'):

						surf['specular'] = read_int2(f)/255.

				elif (subID == 'REFL'):

					if not surf.has_key('reflection'):

						surf['reflection'] = read_int2(f)/255.

				elif (subID == 'TRAN'):

					if not surf.has_key('transparency'):

						surf['transparency'] = read_int2(f)/255.

				elif (subID == 'RIND'):

					surf['refractive-index'] = read_float4(f)

				elif (subID == 'GLOS'):

					surf['glossiness'] = read_int2(f)

				elif (subID == 'SMAN'):

					surf['smoothing-max-angle'] = read_float4(f)

				else:

					print "Warning: SURF sub chunk", subID,"("+str(subsize),"B) ignored"

					f.read(subsize)

				size -= subsize

			surfaces.append(surf)

		else:

			print "Warning: chunk", ID,"("+str(size),"B) ignored"

			f.read(size)

		(ID,size) = read_chunk(f)

	return (points, faces, tags, surfaces)

def dot(a,b):

	sum = 0

	for i in range(min(len(a),len(b))):

		sum += a[i]*b[i]

	return sum

def LoadLightwaveLwoFile(rt, filename, 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_faces = []

	materials = []

	for surf in surfaces:

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

		if surf.has_key('smooth'):

			mat.setSmooth(surf['smooth'])

		diff = 1.

		if surf.has_key('diffuse'):

			diff = surf['diffuse']

		spec = 0.

		if surf.has_key('specular'):

			spec = surf['specular']

		gloss = 1.0

		if surf.has_key('glossiness'):

			gloss = surf['glossiness']

		mat.setPhong(0.1, diff, spec, gloss)

		refl = 0.

		if surf.has_key('reflection'):

			refl = surf['reflection']

		mat.setReflectivity(refl)

		transp = 0.

		if surf.has_key('transparency'):

			transp = surf['transparency']

		rindex = 1.0

		if surf.has_key('refractive-index'):

			rindex = surf['refractive-index']

		mat.setTransmissivity(transp, rindex)

		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_faces.append([])

	for f in faces:

		for x in f[0:3]:

			vertex_faces[x].append(f)

	# interpolate normals at vertices

	num = 0

	for vfaces in vertex_faces:

		vert = vertices[num]

		edges = {}

		N = [0,0,0]

		for f in vfaces:

			for fvert in f[0:3]:

				if edges.has_key(str(fvert)):

					edges[str(fvert)].append(f)

				else:

					edges[str(fvert)] = [f]

		for f in vfaces:

			vv = [vertices[x] for x in f[0:3]]

			fN = Triangle(vv[0], vv[1], vv[2], materials[f[3]-1]).getNormal()

			for i in range(3):

				N[i] += fN[i]

			surf = surfaces[f[3]-1]

			if not surf.has_key('smoothing-max-angle'):

				continue

			fNvert = list(fN)

			Ncount = 1

			copy_vertex = False

			for fvert in f[0:3]:

				for ef in edges[str(fvert)]:

					if ef == f:

						continue

					# f - ref. face; ef - other face

					vv = [vertices[x] for x in ef[0:3]]

					efN = Triangle(vv[0], vv[1], vv[2], materials[ef[3]-1]).getNormal()

					d = dot(fN, efN)

					if d > 1:

						d = 1

					if d < -1:

						d = -1

					if acos(d) < surf['smoothing-max-angle']:

						for i in range(3):

							fNvert[i] += efN[i]

						Ncount += 1

					else:

						copy_vertex = True

			# here fNvert is normal for num'th vertex in face f

			if copy_vertex:

				for i in range(3):

					fNvert[i] /= Ncount

				new_vert = NormalVertex(vert)

				new_vert.setNormal(tuple(fNvert))

				f.append(f[0])

				f.append(f[1])

				f.append(f[2])

				vertices.append(new_vert)

				for i in range(3):

					if f[i] == num:

						f[i+4] = len(vertices)-1

		if (len(vfaces) > 0):

			for i in range(3):

				N[i] /= len(vfaces)

		vertices[num].setNormal(tuple(N))

		num += 1

	for f in faces:

		if len(f) > 4:

			v = [vertices[x] for x in f[4:7]]

		else:

			v = [vertices[x] for x in f[0:3]]

		matidx = f[3]-1

		face = Triangle(v[0], v[1], v[2], materials[matidx])

		rt.addshape(face)