#!BPY

""" Registration info for Blender menus:
Name: 'LDraw...'
Blender: 234
Group: 'Import'
Tip: 'Import an LDraw (.dat,.ldr) 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 MPD support.
#  Add loading of objects from Blender (children, templates)?
#  Implement material calculation?
#  Complete detection of file type (pieces)
#  Add BFC support and quad winding fixes
#
#  Flatten works, but fix it to fold in primitives and flatten parts
#  based on which search path they came from.

ldrawdir='/home/yann/Blender/ldraw/ldraw'

from Blender.NMesh import Vert, Face
from Blender.Mathutils import Matrix, Vector, MatMultVec, CopyMat
import Blender

# This will not remove *all* duplicate vertices, but enough to
# make a difference (over 2/3 of all vertices). Blender's built-in
# Remove Doubles is better, wish I could just do mesh.RemoveDoubles(0.001).
# When the emesh module is done, I guess I will be able to.
# Turning this off will speed the import up, break autosmoothing,
# and increase file size.
RemoveDoubles=True

models={}
col2bmat={}
# Attempt to load material-color mappings from a Text.
try:
  for line in Blender.Text.Get('LDrawMaterials').asLines():
    words=line.split()
    try:
      number=int(words[0])
      name=words[1]
      col2bmat[number]=Blender.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=Blender.Material.New('LDraw%d'%(number))
    try:
      text=Blender.Text.Get('LDrawMaterials')
    except NameError:
      text=Blender.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

class datfile:
  def __init__(self, file):
    self.name=Blender.sys.basename(file.name)
    self.file=file
  def __getitem__(self, name=None):
    if name==self.name or name is None:
      return self.file
    else:
      raise KeyError, name
  def has_key(self, name):
    return name==self.name

class mpdfile:
  def __init__(self, file):
    self.files={}
    first=True
    for line in file:
      if len(line)>7 and line[:7]=='0 FILE ':
	part=[]
	name=line[7:].rstrip()
	if first:
	  self.firstfile=part
	  first=False
	self.files[name]=part
      else:
	part.append(line)
  def __getitem__(self, name=None):
    if name is None:
      return self.firstfile
    else:
      return self.files[name]
  def has_key(self, name):
    return self.files.has_key(name)

class directory:
  def __init__(self, path):
    self.path=path
  def __getitem__(self, name):
    fullname=Blender.sys.join(self.path,name)
    if Blender.sys.exists(fullname)==1:
      return file(fullname)
    else:
      raise KeyError, name
  def has_key(self, name):
    fullname=Blender.sys.join(self.path,name)
    return Blender.sys.exists(fullname)==1

class datfinder:
  def __init__(self):
    self.finders=[]
  def __getitem__(self, name=None):
    for finder in self.finders:
      try:
        return finder[name]
      except KeyError:
        pass
    raise KeyError, name
  def has_key(self, name):
    for finder in self.finders:
      if finder.has_key():
        return True
    return False
  def insertpath(self, path):
    '''Add a path (directory, dat, mpd) to the search path. NOTE: This puts it first in the search path.'''
    typ=Blender.sys.exists(path)
    if typ==1:
      # Include the directory for searching
      dir=Blender.sys.dirname(path)
      self.finders.insert(0, directory(dir))
      # Check if it's an MPD
      fil=file(path,'rU')
      head=fil.read(7)
      fil.seek(0)
      if head=='0 FILE ':
        # It is
	self.finders.insert(0, mpdfile(fil))
      else:
        # Nope. Assume it's a DAT file.
	self.finders.insert(0, datfile(fil))
    elif typ==2:
      # It's a directory
      self.finders.insert(0, directory(path))
    else:
      raise ValueError, path

datlib=datfinder()
for name in 'p', 'parts':
  path=Blender.sys.join(ldrawdir,name)
  datlib.insertpath(path)

def getmdl(name):
  '''Gets a legomesh instance for the given model name. May load different files'''
  # TODO: Determine model type by path. Load paths from environment.
  # Etc etc etc...
  fname=name.replace('\\', Blender.sys.sep)
  fname=fname.replace('/', Blender.sys.sep)
  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

# TODO: binary search/sorted list for the vertices!

def place(l, v):
  '''Returns the index of the second argument in the first, appending it if necessary.'''
  # For duplicate vertex elimination. Far too slow.
  #for i in xrange(len(verts)):
  #  v=verts[i]
  #  if abs(v[0]-vert[0])>=0.001: continue
  #  if abs(v[1]-vert[1])>=0.001: continue
  #  if abs(v[2]-vert[2])>=0.001: continue
  #  return i
  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
    # Face: [colour,smooth,(x,y,z),(x,y,z)...]
    self.faces=[]
    self.verts=[]
    # Children will be a list of colour, matrix, legomesh
    self.children=[]
    self.colours=[16]
  def toBMesh(self):
    # Note: This means we'll never override a Blender Mesh.
    if self.name and self.mesh is None:
      self.mesh=Blender.NMesh.GetRaw(self.name)
    if self.mesh is None:
      mesh=Blender.NMesh.New(self.name)
      mesh.materials=[None]+map(material, self.colours[1:])
      # 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 in self.faces:
        bface=Face([mesh.verts[v] for v in face[2:]])
	bface.smooth=face[1]
	bface.mat=face[0]
	mesh.faces.append(bface)
      # Perhaps autosmooth threshold may need tuning for some objects.
      # By enabling it only for faces connected to 5 lines, we're pretty safe.
      mesh.setMode('TwoSided', 'AutoSmooth')
      self.mesh=mesh
    return self.mesh
  def toBObj(self, mycolour=7):
    '''Creates the Blender objects, including children'''
    obj=Blender.Object.New('Mesh', self.name)
    obj.link(self.toBMesh())
    obj.colbits=1
    obj.setMaterials([material(mycolour)])
    Blender.Scene.getCurrent().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, 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 flattenParts(self):
    '''Makes sure all Part models used by this model or its submodels are flattened'''
    for colour, 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, matrix, model in self.children:
      model.flatten()
      ci=[place(self.colours, c) for c in
          [colour]+model.colours[1:]]
      if RemoveDoubles:
	vi=[place(self.verts, v) for v in
	    [MatMultVec(matrix, Vector(list(v[0:3])+[1]))[0:3]
	    for v in model.verts]]
	self.faces.extend([
	    [ci[face[0]],face[1]] + [vi[v] for v in face[2:]]
	    for face in model.faces])
      else:
	vo=len(self.verts)
	self.verts.extend([MatMultVec(matrix, Vector(list(v[0:3])+[1]))[0:3] for v in model.verts])
	self.faces.extend([
	    [ci[face[0]],face[1]] + [v+vo for v in face[2:]]
	    for face in model.faces])
    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] 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, mat, childmdl))
      elif line[0] in ('3','4'):
	colour=place(self.colours, int(line[1]))
        face=[colour,0]
	for i in xrange(2, 2+int(line[0])*3, 3):
	  if RemoveDoubles:
	    face.append(place(self.verts, tuple(map(float, line[i:i+3]))))
	  else:
	    face.append(len(self.verts))
	    self.verts.append(tuple(map(float, line[i:i+3])))
	self.faces.append(face)
      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 face in self.faces:
	  if v[0] in face[2:] and v[1] in face[2:]:
	    face[1]=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, 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=Blender.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):
  datlib.insertpath(filename)
  name=Blender.sys.basename(filename)
  lm=maindat(name)
  # None will cause it to load the first part of an MPD.
  lm.fromFile(datlib[None])
  lm.flattenParts()
  scene=Blender.Scene.getCurrent()
  rc=scene.getRenderingContext()
  steps=len(lm.steps)
  rc.currentFrame(steps)
  rc.endFrame(steps)
  obj=lm.toBObj()
  scene.update()

from Blender.Window import FileSelector
FileSelector(loaddat)