#!BPY

################################################################
# Half-Life 2 Static SMD Importer Blender Script               #
# Released 22. August 2007                                     #
################################################################
# ***** BEGIN GPL LICENSE BLOCK *****
#
# Copyright (C) 2007 Jon Moen Drange - jondrange [A] gmail.com #
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
#
# ***** END GPL LICENCE BLOCK *****

"""
Name: 'Half-Life 2 (.SMD) ...'
Blender 244
Group: 'Import'
Submenu: 'Static Mesh ...' import_staticmesh
Submenu: 'Static Mesh with Textures ...' import_staticmesh_tex
Tooltip: 'Import SMD Mesh and Armature/Bones'
"""

__author__  = "Jon Moen Drange"
__url__     = ("http://folk.uio.no/jonmd")
__version__ = "v0.1b"
__bpydoc__  = """This script imports static mesh with or without textures from the
Half-Life 2 ASCII format SMD. The script needs Blender 2.44 or
higher to work.

Static Mesh
-----------
This option imports the mesh from the SMD with fully functional
bones/armatures. The SMD file allways contains at least one bone,
and for that reason there will allways be imported one armature.
For this option, the textures will be ignored.

Static Mesh w/ Textures
-----------------------
Same as the one above, only this one also reads the sticky coordinates
and loads the images/textures used by the model into blender. In
order to do so, the files containing the textures should allways be
present in the same directory as the SMD file.

Other Notes
-----------
I don't know what is "front" in the Half-Life2 coordination system, and
my script doesn't rotate anything imported mesh or armature. Be aware
that imported models might face the wrong directions.

Future Expansion (just thoughts)
--------------------------------
- Currently all bones get the same length. A future improvement might 
  be place the tail of a bone with only one child at the same position
  as the childs head and connect them. I wrote a small piece of code
  that did this, but it doesn't work correctly in all cases. Therefore,
  it's "commented away". (Found around line 345)
- Support for importing animations.
- Better import for textures.

Bugs
----
Bugs found and fixed:
- From Valve developer wiki:
  "<bone ID> <weight> will be looped at the end of each vertex
  <# of links> times. If the weights don't add up to 1.0 the rest will
  be assigned to the <Parent> bone. Same if no <# of links> is defined
  or if it is 0. Each triangle contains 3 vertex records."
  This was implemented inaccuratly. Now updated and fixed.
  (Discovered by "xz64")
- Lineendings on linux weren't treated properly. Now fixed.
  (Discovered by "J. Snake")

Credits
-------
The first import script for Blender was made by Michael Fuller. His
script doesn't work anymore, because of updates in the Blender Python
API. This inspired me to write my own script. All code in this file
is written by Jon Moen Drange (username JMD at the blenderartists.org
forum). The script can be found at:
http://folk.uio.no/jonmd/blender/smdimport/
"""


############################ CODE STARTS HERE ############################

import Blender, meshtools, bpy, math, time
from Blender import *
from Blender.Mathutils import *
import string


########################################################
##### Exceptionclasses #################################

class ParsingSMDFileError (Exception):
   def __init__(self,value):
      Exception.__init__(self,value)


########################################################
##### Small classes for various purposes ###############

class Bone:
   def __init__(self,number,name,bone,parentnumber,bones):
      self.childs   = []
      self.index    = number
      self.name     = name
      self.bone     = bone
      self.bone.name= self.name
      self.indexPar = parentnumber
      self.parent   = None
      if (self.indexPar > -1):
         self.parent   = bones[self.indexPar]
         #self.bone.parent = bones[self.indexPar].bone
      self.matrix   = None

   def __repr__(self):
      return '['+str(self.index)+' "'+self.name+'" Par: '+str(self.indexPar)+']'
   def __str__(self):
      return self.__repr__()


class Vertex:
   def __init__(self,position=None,index=None,vertex=None,tex=None,texco=None,
                vertgroups=None,vertgroupsWeights=None):
      self.pos    = position
      self.index  = index # need this?
      self.v      = vertex
      self.texture= tex
      self.texco  = texco
      if (vertgroups == None or vertgroupsWeights == None):
         self.vgs    = []
         self.vgsW   = []
      else:
         self.vgs  = vertgroups
         self.vgsW = vertgroupsWeights
      
   def __repr__(self):
      return '['+str(self.pos)+' '+str(self.norms)+']'
   def __str__(self):
      return self.__repr__()


class Face:
   def __init__(self,face=None,texture=None):
      self.texname = texture
      self.v1 = None
      self.v2 = None
      self.v3 = None
      self.uv1 = None
      self.uv2 = None
      self.uv3 = None
      self.f = face

   def __repr__(self):
      return '['+str(self.vertices)+' Tex="'+self.texname+'" ]'
   def __str__(self):
      return self.__repr__()


class VertexGroup:
   def __init__(self,name,bone):
      self.name  = name
      self.bone  = bone
      self.verts = 0

   def __repr__(self):
      return '["'+self.name+'"  bone:'+str(self.bone)+']'
   def __str__(self):
      return self.__repr__()


########################################################
##### Base class for all importers #####################
class Importer:
   def __init__(self, filename, doImportMaterial):
      # making file ready for reading
      self.filename = filename
      self.file = open(filename, "rb")
      self._lines = self.file.readlines()
      self._lineno = 0

      # settings
      self.importmaterial = doImportMaterial

      # useful vars / pointers
      self.meshObj = None
      self.armObj  = None
      self.mesh    = None
      self.armature= None
      self.textures= []
      self.texnames= []
      self.bones   = []
      self.verts   = []
      self.vertgrps= []
      self.faces   = []
      self.coords  = {}
      self.scene   = Blender.Scene.GetCurrent()
      self.frame   = 0
      self.noBones = 0
      self.noVerts = 0
      self.noReadV = 0
      self.noFaces = 0
      self.time0   = 0
      self.scale   = 1.0 / 40.0
      self.scaleMat = Matrix([self.scale,0,0,0],
                             [0,self.scale,0,0],
                             [0,0,self.scale,0],
                             [0,0,0,         1])

      #TODO add more variables
      
      

   def printDebugInfo(self):
      print '\nDEBUGINFO for SMD Importer (fixme!)\n'


   def getLine(self):
      if (self._lineno < len(self._lines)):
         tmp = string.strip(self._lines[self._lineno].split("\r\n",1)[0])
         self._lineno += 1
         return tmp
      else:
         return ''


########################################################
##### class for importing mesh with/without armature ###
class StaticMeshImporter (Importer):
   def __init__(self, filename, doImportMaterial):
      Importer.__init__(self, filename, doImportMaterial)


   def importSMD(self):
      line = self.getLine()
      if line != 'version 1':
         raise ParsingSMDFileError,'First line should be "version 1" ('+line+')'
      self.parseNodes()
      self.parseArmature()
      self.parseMesh()
      self.setupFinal()


   ### parsing the nodes data from SMD file
   def parseNodes(self):
      print 'Parsing Nodes ...'
      line = self.getLine()
      if line != 'nodes':
         raise ParsingSMDFileError,('Expecting section with nodes data'
                              +' at line '+str(self._lineno))
      
      line = self.getLine()
      while ( line != 'end'):
         line = line.split()
         boneindex   = int(line[0])
         parentindex = int(line[ len(line)-1 ])
         
         tmp = '' # finding the name of the bone
         for i in range (1,len(line)-1):
            if i == 1: tmp += str(line[i])
            else: tmp += ' ' + str(line[i])
         tmp = tmp[1:]; tmp = tmp[:-1]
         
         if (boneindex != self.noBones):
            raise ParsingSMDFileError,('Inconsistent bone-indices in SMD file'
                                 +' at line '+str(self._lineno))
         self.bones.append(Bone(self.noBones,tmp,
                                Armature.Editbone(),parentindex,self.bones))
         self.vertgrps.append(VertexGroup(self.bones[-1].name,self.bones[-1].bone))
         self.noBones += 1
         
         line = self.getLine()
      # end while

      if DEBUG: #printing bones
         print 'BONESDEBUG'
         for i in range (0,len(self.bones)):
            print str(i) + ' ->  ' + str(self.bones[i])

      print 'Parsing Nodes: OK'


   ### parsing the skeleton data from SMD file
   def parseArmature(self):
      print 'Parsing Armature ...'
      line = self.getLine()
      if (line != 'skeleton'):
         raise ParsingSMDFileError, ('Expecting section with skeleton data'
                                 +' at line: '+str(self._lineno))

      # setting up armature
      self.armature = Armature.New() #TODO name?
      self.armature.makeEditable()
      self.armature.drawType = Armature.STICK

      mode = 0
      line = self.getLine()
      while (line != 'end'):
         if ( (mode % (self.noBones+1)) == 0):
            #TODO make sure ONLY "time 0"
            if DEBUG: print line
            mode = 1; line = self.getLine()
            continue # to next iteration

         line = line.split()
         if (len(line) != 7):
            raise ParsingSMDFileError, ('Expected 7 numbers (1 int +'
                                 +' 6 floats) at line '+str(self._lineno))

         # reading from current line
         boneNo = int(line[0])
         posX = float(line[1])
         posY = float(line[2])
         posZ = float(line[3])
         rotX = math.degrees( float(line[4]) )
         rotY = math.degrees( float(line[5]) )
         rotZ = math.degrees( float(line[6]) )
         # creating some matrices for current bone
         current = self.bones[boneNo]
         currentVector = Vector([posX,posY,posZ,1]) * self.scaleMat
         currentRotMat = (RotationMatrix(rotX,4,'x')
                          * RotationMatrix(rotY,4,'y')
                          * RotationMatrix(rotZ,4,'z'))

         # calculating were the bone should be
         if (current.indexPar == -1):
            # calculating when has no parent
            matrix1 = currentRotMat * TranslationMatrix(currentVector)
            
            current.bone.head =(Vector([0,0,0,1])*matrix1).resize3D()
            current.bone.tail =(Vector([0.05,0,0,1])*matrix1).resize3D()
            
         else:
            # calculating when has parent
            parent  = current.parent
            matrix1 = currentRotMat * TranslationMatrix(currentVector)
            matrix1 = matrix1 * parent.matrix

            current.bone.head = (Vector([0,0,0,1])*matrix1).resize3D()
            current.bone.tail = (Vector([0.05,0,0,1])*matrix1).resize3D()
            
            current.bone.parent = self.bones[current.indexPar].bone
            current.parent.childs.append(current)

         # adding to my list of bones
         self.armature.bones[ current.name ] = current.bone
         current.matrix = matrix1
         
         mode += 1
         line = self.getLine() # end while

      # FOUND BUG IN IT... 
      # resizing bones, so that the ones with ONLY one
      # child stretches all the way to the head of the
      # child.  Not necessary, but gives a better look
      # of the skeleton. Connecting them, as well!
      #for bone in self.bones:
      #   if len(bone.childs) == 1:
      #      bone.bone.tail = bone.childs[0].bone.head
      #      bone.childs[0].bone.options = Armature.CONNECTED

      print 'Parsing Armature: OK'
      return
      # took out finalsetup armature


   ### parsing the mesh data from SMD file
   def parseMesh(self):
      print 'Parsing Mesh (all triangles) ...'
      line = self.getLine()
      if (line != 'triangles'):
         raise ParsingSMDFileError, ('Expecting section with triangles data'
                                 +' at line: '+str(self._lineno))

      if Window.EditMode():
         Window.EditMode(0) # make sure not in editmode
      self.mesh = Mesh.New('mesh') # TODO: need this allready?

      tempVerts = []
      texture = ''
      tex = None
      mode = 0
      
      line = self.getLine()
      while (line != 'end'):
         if ( (mode % 4) == 0): # first line in a "face-block"

            # check if i should import material
            if not self.importmaterial:
               tempVerts = [None,None,None]
               mode = 1
               line = self.getLine()
               continue

            # shall import material .. then do it!
            texture = line
            
            found = False
            for tx in self.texnames:
               if tx == texture:
                  found = True
            if not found:
               # temporary: adding texture to Blender if not found
               tmp = self.filename
               x = tmp.split('\\',)
               directory = ''
               for t in range(0,len(x)-1):
                  directory += x[t] + '\\'
               directory += texture
               try:
                  tex = Image.Load(directory)
               except:
                  raise IOError, ('(Importing w/tex):\nExpecting'+
                                  ' image "'+texture+'" to be '+
                                  'found in same directory as the SMD')
               self.textures.append(tex)
               self.texnames.append(texture)
               
            tempVerts = [None,None,None]
            mode = 1
            line = self.getLine()
            continue

         line = line.split()
         if (len(line) < 9):
            raise ParsingSMDFileError, ('Expecting 9 + (1 + 2*n) numbers'
                                 +' at line: '+str(self._lineno))

         vertex = Vertex()
         parent = int(line[0])
         c1 = float(line[1]) * self.scale # scaling down to match blender
         c2 = float(line[2]) * self.scale # and dragonlords export script
         c3 = float(line[3]) * self.scale # Standard: scale = 1.0 / 40.0
         vertex.pos = (c1,c2,c3)
         self.noReadV += 1

         # checking if (not) allready exists
         if not self.coords.has_key(vertex.pos):
            self.mesh.verts.extend(c1,c2,c3)
            vertex.v = self.mesh.verts[-1]
            #vertex.v.no = Vector([float(line[4]),float(line[5]),
            #                     float(line[6])])
            if self.importmaterial:
               vertex.texco = Vector([float(line[7]),float(line[8])])
            vertex.index = self.noVerts # TODO: need this?

            vertex.vgs  = []
            vertex.vgsW = []

            weight = 0.0
            if len(line) > 9:
               for i in range(0,int(line[9])):
                  vertex.vgs.append(int(line[10+i*2]))
                  vertex.vgsW.append(float(line[11+i*2]))
                  self.vertgrps[int(line[10+i*2])].verts += 1
                  weight += float(line[11+i*2])

            if parent != 0 and weight < 0.99:
               vertex.vgs.append(parent)
               vertex.vgsW.append(1.0-weight)
               self.vertgrps[parent].verts += 1

            self.coords[vertex.pos] = vertex
            self.verts.append(vertex)
            self.noVerts += 1

         # if vertex allready exists: new texco for the new face
         else:
            vertex = self.coords.get(vertex.pos) # exists
            if self.importmaterial:
               vertex.texco = Vector([float(line[7]),float(line[8])])

         # add v to list of v's used in current face
         tempVerts[mode-1] = vertex
                  
         if ( (mode % 4) == 3):
            self.mesh.faces.extend(tempVerts[0].v,
                                   tempVerts[1].v,
                                   tempVerts[2].v)
            # only if materials should be imported
            if self.importmaterial:
               self.mesh.faces[-1].image = tex
               self.mesh.faces[-1].uv = (tempVerts[0].texco,
                                         tempVerts[1].texco,
                                         tempVerts[2].texco)
               
            self.mesh.faces[-1].smooth = 1
            mode = -1 # don't need it ...
            self.noFaces += 1
            tempVerts = None # make empty

         mode += 1
         line = self.getLine() # end while
         
      print 'Parsing Mesh (all triangles): OK'


   ### final setup: fixing Vertgroups, linking things
   ### together
   def setupFinal (self):
      # setting up armature
      self.armObj = Object.New("Armature",'imported_armature') #FIXME?
      self.armObj.link(self.armature)
      # setting up mesh
      self.meshObj = Object.New('Mesh', 'imported_mesh') #FIXME?
      self.meshObj.link( self.mesh )

      # making mesh/armature ready for rigging
      self.armature.vertexGroups = True 
      self.armObj.makeParent([self.meshObj])
      newmod = self.meshObj.modifiers.append(Modifier.Types.ARMATURE)
      newmod[Modifier.Settings.OBJECT] = self.armObj

      #drawing up armature/mesh
      self.mesh.update()
      self.meshObj.makeDisplayList()
      self.scene.objects.link(self.meshObj)
      self.armature.update()
      self.armObj.makeDisplayList()
      self.scene.objects.link(self.armObj)

      # vertgroups
      for i in range(0,len(self.bones)):
         if self.vertgrps[i].verts > 0:
            self.mesh.addVertGroup(self.bones[i].name)

      for v in self.verts:
         for g in range(0,len(v.vgs)):
            a = v.vgs[g]
            self.mesh.assignVertsToGroup( self.bones[ a ].name,
                                          [ v.v.index ],
                                          v.vgsW[g],
                                          Mesh.AssignModes.ADD )

      # updating blender and redrawing
      bpy.data.scenes.active.update(1) 
      Blender.Redraw()

      # printing import statistics
      print '\n  # of Bones:            ' + str(self.noBones)
      print   '  # of Vertices in file: ' + str(self.noReadV)
      print   '  # of Actual Vertices:  ' + str(self.noVerts)
      print   '  # of Faces:            ' + str(self.noFaces) + '\n'
      # end setupfinal


   ### cleaning up my mess :-(
   def cleanup(self):
      print ' ' # todo: maybe fix me
      


########################################################
##### Variouse Functions ###############################

def userAlert (msg):
   Blender.Draw.PupMenu ( 'An Error Occured%t|'+msg )


### Control-function for the import script #############
def myImport (filename):
   # Printing info in console
   print '\n\nHalf-Life 2 SMD Importer ('+__version__+')'
   print '**********************************************'
   
   # Chosing type of import
   options  = __script__[ 'arg' ]
   importer = None
   error = False

   try:
      if (options == 'import_staticmesh'):
         importer = StaticMeshImporter(filename,False)
      elif (options == 'import_staticmesh_tex'):
         importer = StaticMeshImporter(filename,True)
      else:
         raise Exception, 'No valid argument option chosen.'
      # starting parsing file in this call
      importer.time0 = time.time() # for taking time
      importer.importSMD()
      
      
   except ParsingSMDFileError, e:
      print '\nERROR (Parsing file):\n' + str(e)
      userAlert ('Import failed. Check console!')
      error = True
   except IOError, e:
      print '\nERROR (IOError):\n' + str(e)
      userAlert ('Import failed. Check console!')
      error = True
   except Exception, e:
      print '\nERROR (Unknown):\n' + str(e)
      userAlert ('Import failed. Check console!')
      error = True

   # Finishing
   if (error == True):
      #if importer != None:
         #print 'Trying to clean up unfinished imported objects...'
         #importer.cleanup()
      print '\nScript terminated. There might be unfinished'
      print 'imported objects in your current project.\n'
   else:
      print '\nImporting Completed!'
      print 'Total time taken: '+str(int(time.time()-importer.time0))+' seconds'
   print '\n**********************************************\n'

# Global variables
DEBUG = False  # whether to print debug information to console or not

##### Script starts with a fileselector
Blender.Window.FileSelector (myImport, 'Import SMD ...')

