# LightWave .lwo file loader

from math import *
from struct import *
from raytracer import Triangle, NormalVertex, Material
from vector import dot

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