#!BPY
# $Id$

"""
Name: 'MSP GL Mesh (.mesh)...'
Blender: 244
Group: 'Export'
"""

import sys
import math
import bpy
import Blender

def make_edge_key(i1, i2):
	return (min(i1, i2), max(i1, i2))

progress_range = (0.0, 1.0, "")

def set_progress(value):
	Blender.Window.DrawProgressBar(progress_range[0]+progress_range[1]*value, progress_range[2])

def set_progress_range(low, high, text):
	global progress_range
	progress_range = (low, high-low, text)
	set_progress(0.0)


class Edge:
	def __init__(self, me):
		if me.__class__==Edge:
			self._medge = me._medge
			self.v1 = me.v1
			self.v2 = me.v2
			self.smooth = me.smooth
		else:
			self._medge = me
			self.smooth = False
		self.faces = []

	def __getattr__(self, attr):
		return getattr(self._medge, attr)

	def check_smooth(self, limit):
		if len(self.faces)!=2:
			return

		d = Blender.Mathutils.DotVecs(self.faces[0].no, self.faces[1].no)
		if (d>limit and self.faces[0].smooth and self.faces[1].smooth) or d>0.999:
			self.smooth = True

	def other_face(self, f):
		if f.index==self.faces[0].index:
			if len(self.faces)>=2:
				return self.faces[1]
			else:
				return None
		else:
			return self.faces[0]


class Vertex:
	def __init__(self, mv):
		if mv.__class__==Vertex:
			self._mvert = mv._mvert
			self.no = mv.no
			self.uv = mv.uv
		else:
			self._mvert = mv
			self.uv = None
		self.flag = False
		self.faces = []
		self.tan = None
		self.bino = None

	def __getattr__(self, attr):
		return getattr(self._mvert, attr)

	def __cmp__(self, other):
		if other is None:
			return 1
		return cmp(self.index, other.index)

	def __str__(self):
		return "<Vert %d (%.4f, %.4f, %.4f) (%.4f, %.4f, %.4f)>"%(self.index, self.co[0], self.co[1], self.co[2], self.no[0], self.no[1], self.no[2])
	
	__repr__ = __str__


class Face:
	def __init__(self, mf):
		self._mface = mf
		self.edges = []
		self.verts = [v for v in mf.verts]
		self.flag = False

	def __getattr__(self, attr):
		return getattr(self._mface, attr)

	def __cmp__(self, other):
		if other is None:
			return 1
		return cmp(self.index, other.index)

	def __str__(self):
		return "<Face %d (%s)>"%(self.index, " ".join([str(v.index) for v in self.verts]))
	
	__repr__ = __str__

	def pivot_vertices(self, *vt):
		flags = [(v in vt) for v in self.verts]
		l = len(self.verts)
		for i in range(l):
			if flags[i] and not flags[(i+l-1)%l]:
				return self.verts[i:]+self.verts[:i]

	def get_edge(self, v1, v2):	
		key = make_edge_key(v1.index, v2.index)
		for e in self.edges:
			if e.key==key:
				return e
		raise KeyError, "No edge %s"%(key,)


class Line:
	def __init__(self, e):
		self.edge = e
		self.verts = [e.v1, e.v2]
		self.flag = False

	def __str__(self):
		return "<Line (%d %d)>"%(self.verts[0].index, self.verts[1].index)

	__repr__ = __str__


class Mesh:
	def __init__(self, m):
		self._mesh = m
		self.verts = [Vertex(v) for v in m.verts]
		self.faces = [Face(f) for f in m.faces]
		self.materials = m.materials[:]

		for f in self.faces:
			for i in range(len(f.verts)):
				f.verts[i] = self.verts[f.verts[i].index]
				f.verts[i].faces.append(f)

		self.edges = dict([(e.key, Edge(e)) for e in m.edges])
		for f in self.faces:
			for k in f.edge_keys:
				e = self.edges[k]
				e.faces.append(self.faces[f.index])
				f.edges.append(e)

		self.lines = [Line(e) for e in self.edges.itervalues() if not e.faces]

		if m.mode&Blender.Mesh.Modes.AUTOSMOOTH:
			smooth_limit = math.cos(m.degr*math.pi/180)
		else:
			smooth_limit = -1

		for e in self.edges.itervalues():
			e.v1 = self.verts[e.v1.index]
			e.v2 = self.verts[e.v2.index]
			e.check_smooth(smooth_limit)

	def __getattr__(self, attr):
		return getattr(self._mesh, attr)

	def splice(self, other):
		material_map = []
		for m in other.materials:
			if m in self.materials:
				material_map.append(self.materials.index(m))
			else:
				material_map.append(len(self.materials))
				self.materials.append(m)

		offset = len(self.verts)
		for v in other.verts:
			v.index += offset
			self.verts.append(v)

		offset = len(self.faces)
		for f in other.faces:
			f.index += offset
			f.mat = material_map[f.mat]
			self.faces.append(f)

		for e in other.edges.itervalues():
			e.key = make_edge_key(e.v1.index, e.v2.index)
			self.edges[e.key] = e

		self.lines += other.lines
	
	def generate_material_uv(self):
		for f in self.faces:
			f.uv = ([(f.mat+0.5)/len(self.materials), 0.5],)*len(f.verts)
		self.faceUV = True

	def split_vertices(self, find_group_func, debug):
		groups = []
		for i in range(len(self.verts)):
			v = self.verts[i]
			for f in v.faces:
				f.flag = False

			vg = []
			for f in v.faces:
				if not f.flag:
					vg.append(find_group_func(v, f))

			groups.append(vg)

			set_progress(i*0.5/len(self.verts))

		for i in range(len(self.verts)):
			if len(groups[i])==1:
				continue

			if debug:
				print "Vertex %s has %d groups"%(self.verts[i], len(groups[i]))

			for g in groups[i][1:]:
				v = Vertex(self.verts[i])
				v.index = len(self.verts)
				self.verts.append(v)

				if debug:
					print "  -> %d %s"%(v.index, [f.index for f in g])

				for f in g:
					for j in range(len(f.edges)):
						e = f.edges[j]

						if e.v1!=self.verts[i] and e.v2!=self.verts[i]:
							continue

						if debug:
							print "  Splitting edge %s with faces %s"%(e.key, e.faces)

						if e.other_face(f) not in g and len(e.faces)>=2:
							k = e.faces.index(f)
							e.faces.remove(f)
							e = Edge(e)
							f.edges[j] = e
							e.faces.append(f)
						else:
							del self.edges[e.key]

						if e.v1==self.verts[i]:
							e.v1 = v
						elif e.v2==self.verts[i]:
							e.v2 = v

						e.key = make_edge_key(e.v1.index, e.v2.index)
						self.edges[e.key] = e

					self.verts[i].faces.remove(f)
					f.verts[f.verts.index(self.verts[i])] = v
					v.faces.append(f)

			set_progress(0.5+i*0.5/len(self.verts))

	def split_smooth(self, debug = False):
		self.split_vertices(self.find_smooth_group, debug)

	def split_uv(self, debug = False):
		self.split_vertices(self.find_uv_group, debug)

	def find_smooth_group(self, vert, face):
		face.flag = True
		queue = [face]

		for f in queue:
			for e in f.edges:
				other = e.other_face(f)
				#if not other or other.index not in face_indices:
				if other not in vert.faces:
					continue

				if e.smooth:
					if not other.flag:
						other.flag = True
						queue.append(other)

		return queue

	def find_uv_group(self, vert, face):
		uv = face.uv[face.verts.index(vert)]
		face.flag = True
		group = [face]
		for f in vert.faces:
			if not f.flag and f.uv[f.verts.index(vert)]==uv:
				f.flag = True
				group.append(f)
		return group

	def compute_normals(self):
		for v in self.verts:
			if v.faces:
				v.no = Blender.Mathutils.Vector()
				for f in v.faces:
					v.no += f.no
				v.no.normalize()
			else:
				# XXX Should use edges to compute normal
				v.no = Blender.Mathutils.Vector(0, 0, 1)

	def compute_uv(self):
		for v in self.verts:
			if v.faces:
				v.uv = v.faces[0].uv[v.faces[0].verts.index(v)]

	def compute_tbn(self):
		for v in self.verts:
			v.tan = Blender.Mathutils.Vector()
			v.bino = Blender.Mathutils.Vector()
			for f in v.faces:
				fverts = f.pivot_vertices(False, v)
				v1 = fverts[1]
				v2 = fverts[-1]
				du1 = v1.uv[0]-v.uv[0]
				du2 = v2.uv[0]-v.uv[0]
				dv1 = v1.uv[1]-v.uv[1]
				dv2 = v2.uv[1]-v.uv[1]
				div = du1*dv2-du2*dv1
				edge1 = fverts[1].co-fverts[0].co
				edge2 = fverts[-1].co-fverts[0].co
				v.tan += (edge1*dv2-edge2*dv1)/div
				v.bino += (edge2*du1-edge1*du2)/div
			v.tan.normalize()
			v.bino.normalize()

	def create_strip(self, face, max_len, debug):
		edge = None
		for e in face.edges:
			other = e.other_face(face)
			if other and not other.flag:
				edge = e
				break

		if not edge:
			return None

		if debug:
			print "Starting strip from %s, edge %s"%([v.index for v in face.verts], (edge.v1.index, edge.v2.index))

		verts = face.pivot_vertices(edge.v1, edge.v2)
		if len(verts)==3:
			result = [verts[-1], verts[0]]
		else:
			result = [verts[-2], verts[-1]]

		while 1:
			if debug:
				print "  Adding %s"%face

			verts = face.pivot_vertices(*result[-2:])
			k = len(result)%2

			face.flag = True
			if len(verts)==4 and not k:
				result.append(verts[3])
			result.append(verts[2])
			if len(verts)==4 and k:
				result.append(verts[3])

			if len(result)>=max_len:
				if debug:
					print "  Max length exceeded"
				break

			edge = face.get_edge(*result[-2:])

			if debug:
				print "  Next edge is %s"%(edge.key, )

			next = edge.other_face(face)
			if not next or next.flag:
				break
			face = next

		if debug:
			print "  %s"%[v.index for v in result]

		return result


class VertexCache:
	def __init__(self, size):
		self.size = size
		self.slots = [-1]*self.size

	def fetch(self, v):
		hit = v.index in self.slots
		if hit:
			self.slots.remove(v.index)
		self.slots.append(v.index)
		if not hit:
			del self.slots[0]
		return hit

	def fetch_strip(self, strip):
		hits = 0
		for v in strip:
			if self.fetch(v):
				hits += 1
		return hits

	def test_strip(self, strip):
		hits = 0
		for i in range(len(strip)):
			if i>=self.size:
				break
			if strip[i].index in self.slots[i:]:
				hits += 1
		return hits


class OutFile:
	def __init__(self, fn):
		if fn==None:
			self.file = sys.stdout
		else:
			self.file = open(fn, "w")
		self.indent = 0

	def make(self, kwd, *params):
		pstr = ""
		for p in params:
			if type(p)==float:
				pstr += " %.6g"%p
			else:
				pstr += " %s"%p
		return "%s%s"%(kwd, pstr)

	def write(self, kwd, *params):
		self.file.write("%s%s;\n"%('\t'*self.indent, self.make(kwd, *params)))

	def begin(self, kwd, *params):
		i = '\t'*self.indent
		self.file.write("%s%s\n%s{\n"%(i, self.make(kwd, *params), i))
		self.indent += 1

	def end(self):
		self.indent -= 1
		self.file.write("%s};\n"%('\t'*self.indent))


class Exporter:
	def __init__(self, fn):
		self.out_file = OutFile(fn)
		self.use_strips = True
		self.use_degen_tris = True
		self.max_strip_len = 1024
		self.optimize_cache = False
		self.cache_size = 64
		self.export_lines = True
		self.tbn_vecs = False
		self.compound = False
		self.object = False
		self.material_tex = False
		self.debug = False
		self.strip_debug = False
		self.split_debug = False

	def stripify(self, mesh):
		for f in mesh.faces:
			f.flag = False

		faces_done = 0
		strips = []
		loose = []

		cache = None
		if self.optimize_cache:
			cache = VertexCache(self.cache_size)

		island = []
		island_strips = []
		while 1:
			if not island:
				queue = []
				for f in mesh.faces:
					if not f.flag:
						f.flag = True
						queue.append(f)
						break

				if not queue:
					break

				while queue:
					f = queue[0]
					del queue[0]
					island.append(f)

					for e in f.edges:
						other = e.other_face(f)
						if other and not other.flag:
							other.flag = True
							queue.append(other)

				for f in island:
					f.flag = False

			best = 5
			face = None
			for f in island:
				if f.flag:
					continue
				score = 0
				for e in f.edges:
					other = e.other_face(f)
					if other and not other.flag:
						score += 1
				if score>0 and score<best:
					face = f
					best = score

			if not face:
				while island_strips:
					best = 0
					if cache:
						best_hits = 0
						for i in range(len(island_strips)):
							hits = cache.test_strip(island_strips[i])
							if hits>best_hits:
								best = i
								best_hits = hits

					s = island_strips[best]
					del island_strips[best]
					strips.append(s)

					if cache:
						cache.fetch_strip(s)

				faces_done += len(island)
				set_progress(float(faces_done)/len(mesh.faces))

				loose += [f for f in island if not f.flag]
				for f in island:
					f.flag = True

				island = []
				island_strips = []
				continue

			strip = mesh.create_strip(face, self.max_strip_len, self.strip_debug)
			if strip:
				island_strips.append(strip)

		if self.debug:
			print "%d strips:"%len(strips)
			for i in range(len(strips)):
				print "  %d: %d indices"%(i, len(strips[i]))
			print "%d loose faces"%len(loose)
			nind = sum([len(s) for s in strips])+sum([len(f.verts) for f in loose])
			print "%d indices total"%nind

		if cache:
			cache = VertexCache(self.cache_size)
			total_hits = 0

		if self.use_degen_tris and strips:
			big_strip = []

			for s in strips:
				if big_strip:
					glue = [big_strip[-1], s[0]]
					if len(big_strip)%2:
						glue += [s[0]]

					big_strip += glue
					if cache:
						total_hits += cache.fetch_strip(glue)

				big_strip += s
				if cache:
					total_hits += cache.fetch_strip(s)

			for f in loose:
				if len(big_strip)%2:
					order = (-1, -2, 0, 1)
				else:
					order = (0, 1, -1, -2)
				verts = [f.verts[i] for i in order[:len(f.verts)]]

				if big_strip:
					glue = [big_strip[-1], verts[0]]
					big_strip += glue
					if cache:
						total_hits += cache.fetch_strip(glue)

				big_strip += verts
				if cache:
					total_hits += cache.fetch_strip(verts)

			strips = [big_strip]
			loose = []
			
			if self.debug:
				nind = len(big_strip)
				print "Big strip has %d indices"%nind
				if self.optimize_cache:
					print "%d cache hits"%total_hits

		if self.debug:
			ntris = sum([len(f.verts)-2 for f in mesh.faces])
			print "%.2f indices per triangle"%(float(nind)/max(ntris, 1))

		return strips, loose

	def export(self):
		scene = bpy.data.scenes.active

		objs = Blender.Object.GetSelected()
		if not objs:
			raise Exception, "Nothing to export"
		for o in objs:
			if o.getType()!="Mesh":
				raise Exception, "Can only export Mesh data"

		Blender.Window.DrawProgressBar(0.0, "Preparing");

		mesh = Blender.Mesh.New("export_tmp")
		mesh.getFromObject(objs[0])
		mesh = Mesh(mesh)
		if self.compound:
			# Must keep a ref to each Blender mesh
			bmeshes = []
			for o in objs[1:]:
				bmesh = Blender.Mesh.New("export_tmp")
				bmesh.getFromObject(o)
				bmeshes.append(bmesh)
				mesh.splice(Mesh(bmesh))

		if self.debug:
			ntris = sum([len(f.verts)-2 for f in mesh.faces])
			print "Starting with %d vertices, %d faces (%d triangles) and %d edges"%(len(mesh.verts), len(mesh.faces), ntris, len(mesh.edges))

		set_progress_range(0.05, 0.35, "Smoothing")
		mesh.split_smooth(self.split_debug)

		if self.debug:
			print "After smooth splitting %d vertices and %d edges"%(len(mesh.verts), len(mesh.edges))

		mesh.compute_normals()

		if self.material_tex:
			mesh.generate_material_uv()

		if mesh.faceUV:
			set_progress_range(0.35, 0.65, "Splitting UVs")
			mesh.split_uv(self.split_debug)
			if self.debug:
				print "After UV splitting %d vertices and %d edges"%(len(mesh.verts), len(mesh.edges))

			mesh.compute_uv()
			if self.tbn_vecs:
				mesh.compute_tbn()

		set_progress_range(0.65, 0.95, "Creating strips")
		strips = []
		loose = mesh.faces
		if self.use_strips:
			strips, loose = self.stripify(mesh)

		Blender.Window.DrawProgressBar(0.95, "Writing file");

		if self.object:
			self.out_file.begin("mesh")

		fmt = "NORMAL3"
		if mesh.faceUV:
			fmt += "_TEXCOORD2"
			if self.tbn_vecs:
				fmt += "_ATTRIB33_ATTRIB34"
		fmt += "_VERTEX3"
		self.out_file.begin("vertices", fmt)
		norm = None
		uv = None
		tan = None
		bino = None
		for v in mesh.verts:
			if v.no!=norm:
				self.out_file.write("normal3", *v.no)
				norm = v.no
			if v.uv!=uv:
				self.out_file.write("texcoord2", *v.uv)
				uv = v.uv
			if v.tan!=tan:
				self.out_file.write("attrib3", 3, *v.tan)
				tan = v.tan
			if v.bino!=bino:
				self.out_file.write("attrib3", 4, *v.bino)
				bino = v.bino
			self.out_file.write("vertex3", *v.co)
		self.out_file.end()
		for s in strips:
			self.out_file.begin("batch", "TRIANGLE_STRIP")
			indices = []
			n = 0
			for v in s:
				indices.append(v.index)
				if len(indices)>=32:
					self.out_file.write("indices", *indices)
					indices = []
			if indices:
				self.out_file.write("indices", *indices)
			self.out_file.end()

		if loose:
			self.out_file.begin("batch", "TRIANGLES")
			for f in loose:
				for i in range(2, len(f.verts)):
					self.out_file.write("indices", f.verts[0].index, f.verts[i-1].index, f.verts[i].index)
			self.out_file.end()

		if self.export_lines and mesh.lines:
			self.out_file.write("batch", "LINES")
			for l in mesh.lines:
				self.out_file.write("indices", l.verts[0].index, l.verts[1].index)
			self.out_file.end()

		if self.object:
			self.out_file.end()
			self.out_file.begin("technique")
			self.out_file.begin("pass", '""')
			if self.material_tex:
				self.out_file.begin("material")
				self.out_file.write("diffuse", 1.0, 1.0, 1.0, 1.0)
				self.out_file.end()
				self.out_file.begin("texunit", 0)
				self.out_file.begin("texture2d")
				self.out_file.write("min_filter", "NEAREST")
				self.out_file.write("mag_filter", "NEAREST")
				self.out_file.write("storage", "RGB", len(mesh.materials), 1)
				texdata = '"'
				for m in mesh.materials:
					texdata += "\\x%02X\\x%02X\\x%02X"%(int(m.R*255), int(m.G*255), int(m.B*255))
				texdata += '"'
				self.out_file.write("raw_data", texdata)
				self.out_file.end()
				self.out_file.end()
			elif mesh.materials:
				m = mesh.materials[0]
				self.out_file.begin("material")
				self.out_file.write("diffuse", m.R, m.G, m.B, 1.0)
				self.out_file.write("ambient", m.R*m.amb, m.G*m.amb, m.B*m.amb, 1.0)
				self.out_file.write("specular", m.specR*m.spec, m.specG*m.spec, m.specB*m.spec, 1.0)
				self.out_file.write("shininess", m.hard);
				self.out_file.end()
			self.out_file.end()
			self.out_file.end()

		Blender.Window.DrawProgressBar(1.0, "Done")


class FrontEnd:
	def __init__(self):
		self.config = Blender.Registry.GetKey('mspgl_export', True) or {}
		self.temp_config = Blender.Registry.GetKey('mspgl_export_temp') or {}

	def run(self):
		self.use_strips = Blender.Draw.Create(self.config.get('use_strips', True))
		self.use_degen_tris = Blender.Draw.Create(self.config.get('use_degen_tris', True))
		self.max_strip_len = Blender.Draw.Create(self.config.get('max_strip_len', 1024))
		self.optimize_cache = Blender.Draw.Create(self.config.get('optimize_cache', False))
		self.cache_size = Blender.Draw.Create(self.config.get('cache_size', 64))
		self.export_lines = Blender.Draw.Create(self.config.get('export_lines', False))
		self.tbn_vecs = Blender.Draw.Create(self.config.get('tbn_vecs', False))
		self.compound = Blender.Draw.Create(self.config.get('compound', False))
		self.object = Blender.Draw.Create(self.config.get('object', False))
		self.material_tex = Blender.Draw.Create(self.config.get('material_tex', False))
		self.debug = Blender.Draw.Create(self.config.get('debug', False))
		self.strip_debug = Blender.Draw.Create(self.config.get('strip_debug', False))
		self.split_debug = Blender.Draw.Create(self.config.get('split_debug', False))
		ret = Blender.Draw.PupBlock("Export MSP GL mesh",
			[("Use strips", self.use_strips, "Generage OpenGL triangle strips"),
				("Use degen tris", self.use_degen_tris, "Use degenerate triangles to combine triangle strips"),
				("Max strip len", self.max_strip_len, 4, 16384, "Maximum length of a triangle strip"),
				("Optimize cache", self.optimize_cache, "Optimize for vertex cache"),
				("Cache size", self.cache_size, 8, 1024, "Cache size to optimize for"),
				("Export lines", self.export_lines, "Export lone edges as lines"),
				("Compute T/B vecs", self.tbn_vecs, "Compute tangent/binormal vectors for bumpmapping"),
				("Compound", self.compound, "Create a compound mesh of all selected objects"),
				("Object", self.object, "Export as an object"),
				("Material texture", self.material_tex, "Create a texture from material data"),
				("Debugging options"),
				("Debug", self.debug),
				("Debug strips", self.strip_debug),
				("Debug splitting", self.split_debug)])
		if ret:
			dirname = self.temp_config.get("dirname", Blender.sys.dirname(Blender.Get("filename")))
			obj = Blender.Object.GetSelected()[0]
			ext = "mesh"
			if self.object.val:
				ext = "object"
			Blender.Window.FileSelector(self.export, "Export MSP GL mesh", "%s/%s.%s"%(dirname, obj.name, ext))

	def draw(self):
		pass

	def export(self, fn):
		self.config['use_strips'] = self.use_strips.val
		self.config['use_degen_tris'] = self.use_degen_tris.val
		self.config['max_strip_len'] = self.max_strip_len.val
		self.config['optimize_cache'] = self.optimize_cache.val
		self.config['cache_size'] = self.cache_size.val
		self.config['export_lines'] = self.export_lines.val
		self.config['tbn_vecs'] = self.tbn_vecs.val
		self.config['compound'] = self.compound.val
		self.config['object'] = self.object.val
		self.config['material_tex'] = self.material_tex.val
		self.config['debug'] = self.debug.val
		self.config['strip_debug'] = self.strip_debug.val
		self.config['split_debug'] = self.split_debug.val
		Blender.Registry.SetKey('mspgl_export', self.config, True)

		import os
		self.temp_config["dirname"] = os.path.dirname(fn)
		Blender.Registry.SetKey('mspgl_export_temp', self.temp_config)

		exp = Exporter(fn)
		exp.use_strips = self.use_strips.val
		exp.use_degen_tris = self.use_degen_tris.val
		exp.max_strip_len = self.max_strip_len.val
		exp.optimize_cache = self.optimize_cache.val
		exp.cache_size = self.cache_size.val
		exp.export_lines = self.export_lines.val
		exp.tbn_vecs = self.tbn_vecs.val
		exp.compound = self.compound.val
		exp.object = self.object.val
		exp.material_tex = self.material_tex.val
		exp.debug = self.debug.val
		exp.strip_debug = self.strip_debug.val
		exp.split_debug = self.split_debug.val
		exp.export()


if __name__=="__main__":
	fe = FrontEnd()
	fe.run()