#!BPY

""" Registration info for Blender menus:
Name: 'LDraw...'
Blender: 234
Group: 'Import'
Tip: 'Import an LDraw (.dat,.ldr,.mpd) file.'
"""

# Warning: MPD support and file type detection are not written.
# Currently the first dat is imported with subobjects as joined meshes.
# TODO:
#  Fix matrix stuff. Rotations are somewhat broken for Objects. (Blender internal)
#  Add loading of objects from Blender (children, templates)?
#  Implement material calculation?
#  Complete detection of file type (pieces)
#
#  Flatten works, but fix it to fold in primitives and flatten parts
#  based on which search path they came from(?)

from ldraw import newdatlib

from Blender.NMesh import Vert, Face
from Blender.Mathutils import Matrix, Vector, MatMultVec, CopyMat, CrossVecs, DotVecs
from Blender.sys import basename, dirsep
from Blender import NMesh, Object, Text, Material, Scene, Ipo
import Blender

col2bmat={}
# Attempt to load material-color mappings from a Text.
try:
  for line in Text.Get('LDrawMaterials').asLines():
    words=line.split()
    try:
      number=int(words[0])
      name=words[1]
      col2bmat[number]=Material.Get(name)
    except IndexError:
      pass
except NameError:
  pass	# Well, no materials defined. What a shame.
def material(number):
  '''Return a Blender material for LDraw color'''
  if col2bmat.has_key(number):
    return col2bmat[number]
  # Material 16 (current colour) is useful for parts. The script shouldn't
  # replace it here as that would break editing of parts.
  #elif number==16:
  #  return material(7)
  else:
    mat=Material.New('LDraw%d'%(number))
    try:
      text=Text.Get('LDrawMaterials')
    except NameError:
      text=Text.New('LDrawMaterials')
    text.write('%d %s\n'%(number,mat.name))
    col2bmat[number]=mat
    # TODO: Fill in material properties. (Automix dithered, make transparent..)
    return mat

def getmdl(name):
  '''Gets a legomesh instance for the given model name. May load different files'''
  global models
  # TODO: Determine model type by path. Load paths from environment.
  # Etc etc etc...
  fname=name.replace('\\', dirsep)
  fname=fname.replace('/', dirsep)
  try:
    # Find if we already have a legomesh for that name
    mdl=models[name]
  except KeyError:
    try:
      fil=datlib[name]
    except KeyError:
      fname=fname.lower()
      fil=datlib[fname]
    mdl=legomesh(name)
    mdl.fromFile(fil)
    models[name]=mdl
  return mdl

def place(l, v):
  '''Returns the index of the second argument in the first, appending it if necessary.'''
  try:
    return l.index(v)
  except ValueError:
    l.append(v)
    return len(l)-1

class legomesh:
  '''Class for holding a LDraw model'''
  def __init__(self, name=None):
    self.type=None
    self.name=name
    # If self.mesh is not None, the mesh exists in Blender.
    self.mesh=None
    # Faces contains vertex indicies
    self.faces=[]
    self.facecol=[]
    self.facesmooth=[]
    self.verts=[]
    # Children will be a list of colour, invert, matrix, legomesh
    self.children=[]
    self.colours=[16]
    # Winding order None, 'CW' or 'CCW'.
    # Note that we're wound counter-clockwise either way, this affects loading.
    self.windorder=None
    self.invertnext=False
  def toBMesh(self):
    # Note: This means we'll never override a Blender Mesh.
    if self.name and self.mesh is None:
      self.mesh=NMesh.GetRaw(self.name)
    if self.mesh is None:
      mesh=NMesh.New(self.name)
      mesh.materials=[None]+map(material, self.colours[1:])
      #self.removeDoubles()
      # If this fails, you're probably not using Python 2.3.
      # Change:
      #  Vert(*vert)
      # to:
      #  apply(Vert, vert)
      mesh.verts=[Vert(*vert) for vert in self.verts[:]]
      for face, smooth, colour in map(None, self.faces, self.facesmooth, self.facecol):
        bface=Face([mesh.verts[v] for v in face])
	bface.smooth=smooth
	bface.mat=colour
	mesh.faces.append(bface)
      # Autosmooth threshold will need tuning for some objects.
      # By enabling it only for faces connected to 5 lines, we're pretty safe.
      # Of course, the full job could be done by disconnecting faces at 2 lines
      # by using duplicate vertices, and using ordinary smoothing, I think.
      # Much easier to leave autosmoothing tuning for the user.
      if self.windorder:
	mesh.setMode('AutoSmooth')
      else:
	mesh.setMode('TwoSided', 'AutoSmooth')
      self.mesh=mesh
    return self.mesh
  def toBObj(self, mycolour=7):
    '''Creates the Blender objects, including children'''
    global scene
    obj=Object.New('Mesh', self.name)
    obj.link(self.toBMesh())
    obj.colbits=1
    obj.setMaterials([material(mycolour)])
    scene.link(obj)
    self.childrenBObj(obj, mycolour)
    return obj
  def childrenBObj(self, obj, mycolour=7, childlist=None):
    '''Creates the children for our object'''
    if not childlist:
      childlist=self.children
    children=[]
    for colour, invert, matrix, model in childlist:
      if colour==16:
        colour=mycolour
      child=model.toBObj(colour)
      matrix=CopyMat(matrix)
      matrix.transpose()
      # FIXME Blender seems to lose some rotations here.
      child.setMatrix(matrix)
      children.append(child)
    obj.makeParent(children, 0, 1)
    return children
  def fromFile(self, file):
    '''Loads LDraw file. Uses iteration.'''
    for line in file:
      self.loadline(line.split())
  def removeDoubles(self):
    if len(self.verts)<=1:
      return
    ind=range(len(self.verts))
    sv=zip(self.verts, ind)
    sv.sort()
    prev_v,first_i = sv.pop(0)
    verts=[prev_v]
    ind[first_i]=prev_newi=0
    for v,i in sv:
      # This is the place to add imprecision.
      # It's also a very cpu intensive loop.
      if v==prev_v:
	# Duplicate vertex
	ind[i]=prev_newi
      else:
	prev_newi+=1
	ind[i]=prev_newi
	prev_v=v
	verts.append(v)
    self.verts=verts
    self.faces=[[ind[v] for v in face] for face in self.faces]
  def flattenParts(self):
    '''Makes sure all Part models used by this model or its submodels are flattened'''
    for colour, invert, matrix, model in self.children:
      if model.type == 'Part' or model.faces:
        model.flatten()
      else:
        model.flattenParts()
  def flatten(self):
    '''Merges all children into the mesh. NOTE: also flattens all children!'''
    for colour, invert, matrix, model in self.children:
      model.flatten()
      if not model.faces:
        continue
      if not model.windorder:
        self.windorder=None
      ci=[place(self.colours, c) for c in
          [colour]+model.colours[1:]]
      voff=len(self.verts)
      if matrix.determinant()<0:
        invert=not invert
      self.facesmooth.extend(model.facesmooth)
      self.facecol.extend([ci[col] for col in model.facecol])
      if invert:
        for face in model.faces:
	  face=[v+voff for v in face]
	  face.reverse()
	  self.faces.append(face)
      else:
	self.faces.extend([[v+voff for v in face] for face in model.faces])
      self.verts.extend([
          tuple(MatMultVec(matrix, Vector(list(v[0:3])+[1]))[:3])
	  for v in model.verts])
    self.removeDoubles()
    self.children=[]
  def step(self):
    '''Records a building step. Does nothing in this class.'''
    pass
  def loadline(self, line):
    '''Loads a single line (passed as list of words) from an LDraw file'''
    if line:
      if line[0] == '0':
        try:
	  if line[1] == 'STEP':
	    self.step()
	  elif line[1] == 'BFC':
	    # Back face culling extension - we use this for winding order
	    if not (self.faces or self.children):
	      # Okay, this BFC command MAY be allowed to set winding order.
	      if 'CW' in line[2:]:
		self.windorder='CW'
	      else:
		self.windorder='CCW'
            if 'INVERTNEXT' in line[2:]:
	      self.invertnext=True
	  elif line[1] in ('LDRAW_ORG', 'Unofficial', 'Un-official'):
	    self.type=line[2]
	  elif tuple(line[1:3]) in (('Official', 'LCAD'), ('Official', 'LCad'), ('Original', 'LDraw')):
	    self.type=line[3]
	except IndexError:
	  pass
      elif line[0] == '1':
        # Inclusion of another file
	x,y,z,a,b,c,d,e,f,g,h,i=map(float, line[2:14])
	#mat=map(float,line[2:14])
	#mat=Matrix([a,b,c,x],[d,e,f,y],[g,h,i,z],[0,0,0,1])
	mat=Matrix([a,b,c,x],[d,e,f,y],[g,h,i,z],[0,0,0,1])
	# Lousy docs. MatMultVec only works if the Matrix is square;
	# otherwise, it will break, including crashes.
	#mat=Matrix(mat[3:6]+[mat[0]], mat[6:9]+[mat[1]], mat[9:]+[mat[2]],[0,0,0,1])
	name=line[14]
	colour=int(line[1])
	childmdl=getmdl(name)
	if childmdl.type in ('Primitive', 'Subpart'):
	  self.type='Part'
	self.children.append((colour, self.invertnext, mat, childmdl))
	self.invertnext=False
      elif line[0] in ('3','4'):
	verts=[map(float, line[i:i+3]) for i in range(2, 2+int(line[0])*3, 3)]
	if self.windorder=='CW':
	  verts.reverse()
	elif self.windorder is None and len(verts)==4:
	  # Fix bowtie quads
	  v=[Vector(v) for v in verts]
	  # This should be a diagonal of the quad
	  v2m0=v[2]-v[0]
	  # The normal of the plane
	  n=CrossVecs(v[1]-v[0], v2m0)
	  # If we want to verify that the quad is flat:
	  #k=DotVecs(v[0],n)
	  #if DotVecs(v[3],n) != k:
	  #  raise "Not Coplanar"
	  # Get a vector pointing along the plane perpendicular to the diagonal
	  n2=CrossVecs(n, v2m0)
	  # Calculate the position relative to the diagonal of three corners
	  k=[DotVecs(p,n2) for p in v[1:]]
	  # k[1] is *on* the diagonal, relative 0.
	  # Check if the corners are on opposite sides of the diagonal
	  if cmp(k[1],k[0]) == cmp(k[1],k[2]):
	    #print "Bowtie quad",v
	    verts=[verts[i] for i in 0,2,3,1]
	first=len(self.verts)
	face=range(first, first+len(verts))
	self.verts.extend(map(tuple, verts))
	self.faces.append(face)
	self.facecol.append(place(self.colours, int(line[1])))
	self.facesmooth.append(0)
      elif line[0] == '5':
        # Conditional line, used on curved surfaces.
	v=[tuple(map(float, line[2:5])), tuple(map(float, line[5:8]))]
	try:
	  v=[self.verts.index(x) for x in v]
	except ValueError:
	  # This line isn't among our faces
	  return
	# Find faces with that edge and set them smooth.
	for i in xrange(len(self.faces)):
	  face=self.faces[i]
	  if v[0] in face and v[1] in face:
	    self.facesmooth[i]=1

class maindat(legomesh):
  '''LDR loader - implements STEP support, but not for submodels.'''
  def __init__(self, name=None):
    legomesh.__init__(self, name)
    self.curframe=1
    self.steps=[]
  def fromFile(self, file):
    legomesh.fromFile(self, file)
    if self.children:
      self.step()
      self.curframe-=1
  def step(self):
    '''Records the children as belonging to one step'''
    if not self.children: return
    self.steps.append(self.children)
    self.children=[]
  def flattenParts(self):
    '''Makes sure all Part models used by this model or its submodels are flattened'''
    for children in self.steps:
      for colour, invert, matrix, model in children:
        # I don't know why, but some Parts just lack the Part word.
	if model.type in ('Part', 'update'):
	  model.flatten()
	else:
	  model.flattenParts()
  def childrenBObj(self, obj, mycolour=7):
    '''Creates the children for our object, including step animations'''
    allchildren=[]
    for step in xrange(1, len(self.steps)+1):
      ipo=Ipo.New('Object', 'step%d'%step)
      # Blender 2.34 can't edit an Object Ipo until it has been assigned to an Object.
      children=legomesh.childrenBObj(self, obj, mycolour, self.steps[step-1])
      for child in children:
	child.setIpo(ipo)
      allchildren.extend(children)
      ipo.addCurve('Layer')
      c=ipo.getCurves()[0]
      if step!=1:
	c.addBezier((1, 1))
      c.addBezier((step, 3))
    return allchildren

def loaddat(filename):
  global scene, models, datlib
  scene=Scene.getCurrent()
  datlib=newdatlib()
  datlib.insertpath(filename)
  models={}
  name=basename(filename)
  lm=maindat(name)
  # None will cause it to load the first part of an MPD.
  lm.fromFile(datlib[None])
  lm.flattenParts()
  steps=len(lm.steps)
  if steps:
    rc=scene.getRenderingContext()
    rc.currentFrame(steps)
    rc.endFrame(steps)
  obj=lm.toBObj()
  scene.update()
  # For some reason, the scene update doesn't do this.
  Blender.Redraw()

if __name__ == '__main__':
  from Blender.Window import FileSelector
  FileSelector(loaddat)