''' 
threeDS.py - Python Module for loading and displaying 3DStudio format modules
Copyright (C) 2000 Jason Petrone <jp@demonseed.net>

The 3DS format is very robust, and much of it is not implemented in this
module.  Currently, 3DS geometry, materials(including textures), and undirected
lights are supported.  Keyframing, directed lighting, and everything else is
currently unimplemented.

'''

from struct import *
import os

dispListCnt = 0

# generic chunks used throughout
RGB1           = 0x0010  # 3 floats of RGB 
RGB2           = 0x0011  # 3 bytes of RGB 
IPERCENT       = 0x0030  # integer percentage
FPERCENT       = 0x0031  # float percentage

OBJMESH        = 0x3D3D

# geometry chunks
OBJBLOCK       = 0x4000
TRIMESH        = 0x4100
VERTLIST       = 0x4110
FACELIST       = 0x4120
FACEMAT        = 0x4130
MAPLIST        = 0x4140
TRMATRIX       = 0x4160
LIGHT          = 0x4600
SPOTLIGHT      = 0x4610
CAMERA         = 0x4700
MAIN           = 0x4D4D

# material chunks
AMBIENT        = 0xA010
DIFFUSE        = 0xA020
SPECULAR       = 0xA030
DOUBLESIDED    = 0xA081
TEXTURE        = 0xA200
MATNAME        = 0xA000
MAT_MAPNAME    = 0xA300
MATMAPTILING   = 0xA351
TEXBLUR        = 0xA353
MAT_MAP_USCALE = 0xA354
MAT_MAP_VSCALE = 0xA356
MATERIAL       = 0xAFFF

KEYFRAMER      = 0xB000

# global textures(dont want to load them more than once!)
tex_data = {}

def readStr(f):
  txt = ""
  [c] = unpack('c', f.read(1) )
  while c != '\0':
    txt = txt + c
    [c] = unpack('c', f.read(1) )
  return txt

def readHdr(f):
  [hdr] = unpack( '<H', f.read(2) )
  [size] = unpack( '<I', f.read(4) )
  return (hdr, size)

class TexImage:
  data = None
  sizex = 0
  sizey = 0

class Texture:
  def __init__(self, f, myLen):
    self.imgfile = None
    while f.tell() < myLen:
      [hdr, size] = readHdr(f)
      if hdr == MAT_MAPNAME:
        global tex_data
        self.imgfile = readStr(f)
#        print "texture: " + self.imgfile
        try: 
          from PIL import Image
        except: 
          print 'Can\'t load texture without PIL!!'
          continue

        if not tex_data.has_key(self.imgfile):
          try:
            img = Image.open(self.imgfile)

            
          except IOError, e:
            print 'Cannot load texture!'
            print str(e)
            self.imgfile = None
            continue
              
          img = img.convert('RGBA', colors=32)
          # texture size must be aligned by 2^n
          w = 0; h = 0
          [w, h] = img.size
          power = 2
          while w > power:
            power = power * 2
          w = power
          power = 2
          while h > power:
            power = power * 2
          h = power
          img = img.resize((w, h))

          # now create storage class
          teximg = TexImage()
          teximg.sizex, teximg.sizey = w, h
          teximg.data = img.tostring('raw', 'RGBA', 0, -1)
          tex_data[self.imgfile] = teximg

      elif hdr == IPERCENT:
        [self.percent] = unpack( '<h', f.read(2) )
      elif hdr == FPERCENT:
        [self.percent] = unpack( '<f', f.read(4) )
      elif hdr == MATMAPTILING:
        # not sure what to do with this
        [self.tiling] = unpack( '<h', f.read(2) )
      elif hdr == TEXBLUR:
        [self.blur] = unpack( '<f', f.read(4) )
      else:
#        print 'hdr: ' + hex(hdr)
        f.read(size - 6)

    return None

  def set(self):
    global tex_data
    if self.imgfile is None: return
    from OpenGL.GL import *
    teximg = tex_data[self.imgfile]
    glPixelStorei(GL_UNPACK_ALIGNMENT, 1)
    glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT)
    glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)
    glTexImage2D(GL_TEXTURE_2D, 0, 4, teximg.sizex, teximg.sizey, 0,
                   GL_RGBA, GL_UNSIGNED_BYTE, teximg.data)

  def unset(self):
    glDisable(GL_TEXTURE_2D)

class Material:
  "All shading values are stored as floats"
  def __init__(self, f = None, myLen = None):

    self.texture = None
    self.doublesided = None
    if f == None:
      # set defaults
      self.ambient = [1.0, 1.0, 1.0, 1.0]
      self.diffuse = [1.0, 1.0, 1.0, 1.0]
      self.specular = [1.0, 1.0, 1.0, 1.0]
      return

    self.ambient = []
    self.diffuse = []
    self.specular = []
    self.name = ""
    while f.tell() < myLen:
      [hdr, size] = readHdr(f)
      if hdr == MATNAME:
        self.name = readStr(f)
#        print 'NAME: ' + self.name
      elif hdr == AMBIENT:
        self.ambient = self.readShading(f, f.tell() + size - 6)
        self.ambient.append(0.7)
      elif hdr == DIFFUSE:
        self.diffuse = self.readShading(f, f.tell() + size - 6)
        self.diffuse.append(0.7)
      elif hdr == SPECULAR:
        self.specular = self.readShading(f, f.tell() + size - 6)
        self.specular.append(0.7)
      elif hdr == TEXTURE:
         self.texture = Texture(f, f.tell() + size - 6)
      elif hdr == DOUBLESIDED:
        self.doublesided = 1
      else:
        f.read(size - 6)
       
  def readShading(self, f, myLen):
    while f.tell() < myLen:
      [hdr, size] = readHdr(f)
      if hdr == RGB1:
        return list(unpack('<3f', f.read(12)))
      elif hdr == RGB2:
        [r, g, b] = unpack( '3B', f.read(3) )
        return [ float(r)/256.0, float(g)/256.0, float(b)/256.0 ]
      else:
        f.read(size - 6)

  def set(self):
    from OpenGL.GL import *
    glMaterialfv(GL_FRONT, GL_AMBIENT, self.ambient)
    glMaterialfv(GL_FRONT, GL_DIFFUSE, self.diffuse)
    glMaterialfv(GL_FRONT, GL_SPECULAR, self.specular)

  def setTex(self):
    if self.texture: self.texture.set()

class Light:
  color = [0, 0, 0]
  pos = [0, 0, 0]
  target = None

  def __init__(self, name, f, myLen):
    self.pos = unpack('<3f', f.read(12))
    while f.tell() < myLen:
      [hdr, size] = readHdr(f)
      if hdr == RGB1:
        self.color = unpack( '<3f', f.read(12) )
      elif hdr == RGB2:
        [r, g, b] = unpack( '3B', f.read(3) )
        self.color = [ float(r)/256.0, float(g)/256.0, float(b)/256.0 ]
      elif hdr == SPOTLIGHT:
        self.target = [0, 0, 0]
        self.target = unpack('<3f', f.read(12))
        self.hotspot = unpack('<1f', f.read(4))
        self.falloff = unpack('<1f', f.read(4))
        
      else:
#        print 'LIGHTCHUNK: ' + hex(hdr) + ', ' + str(size)
        f.read(size - 6)

class TriObj:
  def __init__(self, name, f, myLen):
#    self.name = readStr(f)
    self.name = name
    self.vertices = []
    self.normals = []
    self.maps = []
    self.traMatrix = []
    self.matName = ""
    self.faceMats = []
    self.faces = []
    self.displayList = None
    self.material = None
    self.materialList = None
    self.load(f, myLen)

  def load(self, f, end):
        while f.tell() < end:
          [hdr, size] = readHdr(f)

          if hdr == VERTLIST:
            [num] = unpack( '<h', f.read(2) )
            c = 0
            for v in range(0, num):
              self.vertices.append( unpack('<3f', f.read(12)) ) 
              c = c + 1
             
          elif hdr == FACELIST:
            [num] = unpack( '<h', f.read(2) )
            for v in range(0, num):
              self.faces.append( unpack('<4h', f.read(8)) )
              
          elif hdr == MAPLIST:
            [num] = unpack( '<h', f.read(2) )
            for v in range(0, num):
              self.maps.append( unpack('<ff0f', f.read(8)) )

          elif hdr == TRMATRIX:
            self.traMatrix = unpack( '<12f0f', f.read(48) )
          
          elif hdr == FACEMAT:
            self.matName = readStr(f)                
            [num] = unpack( '<h', f.read(2) )
            for v in range(0, num):
              [mat] = unpack( '<h', f.read(2) )
              self.faceMats.append(mat)

          else:
            f.read(size - 6)


  def draw(self):
    x = [0.0, 0.0, 1.0, 1.0]
    y = [0.0, 1.0, 1.0, 0.0]
    from OpenGL.GL import *
   
    glPushMatrix()
    glBegin(GL_TRIANGLES)
    i = 0
    for face in self.faces:
      face = face[:3]
      if len(face) != 3: print "BAD FACE " + str(len(face))
      if self.material: self.material.set()

      if self.normals[i]: glNormal3f(self.normals[i])
      i = i + 1
      j = 0
      for vertex in face:
        v = self.vertices[vertex]
        glVertex3f(v[0], v[2], v[1])
        j = j + 1

    glEnd()

    if self.material.texture: 
      self.material.setTex()
      glEnable(GL_TEXTURE_2D)
      glTexEnvf(GL_TEXTURE_ENV, GL_TEXTURE_ENV_MODE, GL_DECAL)
      glBegin(GL_QUADS)
      for face in self.faces:
        face = face[:3]
        if len(face) != 3: print "BAD FACE " + str(len(face))
        if self.material: self.material.set()
        i = i + 1
        j = 0
        for vertex in face:
          if self.material.texture:
            glTexCoord2f(self.maps[vertex][0], self.maps[vertex][1])
            v = self.vertices[vertex]
            glVertex3f(v[0], v[2], v[1])
            j = j + 1
      glEnd()

      glDisable(GL_TEXTURE_2D)
    glPopMatrix()

  def fastDraw(self):
    from OpenGL.GL import *
    if self.displayList is None:
#      print 'creating list: ' + self.name
      self.displayList = glGenLists(1)
      glNewList(self.displayList, GL_COMPILE)
      self.draw()
      glEndList()
#      print 'done'
    glCallList(self.displayList)


  def calcNormals(self):
    for face in self.faces:
      face = face[:3]
      if len(face) != 3: print "BAD FACE " + str(len(face))
      v0 = self.vertices[ face[0] ]
      v1 = self.vertices[ face[1] ]
      v2 = self.vertices[ face[2] ]
      a = [ v0[0] - v1[0], v0[1] - v1[1], v0[2] - v1[2] ]
      b = [ v1[0] - v2[0], v1[1] - v2[1], v1[2] - v2[2] ]
      c = [ a[1]*b[2] - b[1]*a[2], 
            b[0]*a[2] - a[0]*b[2],
            a[0]*b[1] - b[0]*b[1] ]

      import math
      l = math.sqrt(c[0]*c[0] + c[1]*c[1] + c[2]*c[2])
      if l == 0: 
        self.normals.append(None)
        continue
      c[0] = c[0] / l
      c[1] = c[1] / l
      c[2] = c[2] / l
      self.normals.append(c)


class Scene:
  materials = {}
  objects = []
  lights = []

  def __init__(self, filename):
    (path, filename) = os.path.split(filename)
    if path != '':
      os.chdir(path)
    f = open(filename, "rb")
    if f is None:
      print "error loading " + str(filename) + "."
      f.close()
      return

    # read file header
    [hdr, totLen] = readHdr(f)
    if hdr != MAIN:
      print str(filename) + ' is not a 3ds file'
      f.close()
      return
    
    end = f.tell() + totLen
    while f.tell() < totLen:
      [hdr, size] = readHdr(f)
      if hdr == OBJMESH:
        self.readMesh(f, f.tell() + size - 6)  
      elif hdr == KEYFRAMER:
        pass
      else:
        f.read(size - 6)   # skip to next chunk
    f.close()
    self.calcNormals()
    self.setMaterials()

  def readMesh(self, f, myLen):
    while f.tell() < myLen:
      [hdr, size] = readHdr(f)
      end = f.tell() + size - 6
      if hdr == MATERIAL:
        m = Material(f, end)
        self.materials[m.name] = m  # new material
      elif hdr == OBJBLOCK:
        self.readObj(f, end)
      else:
#        print 'MESH: ' + hex(hdr)
        f.read(size - 6)                         # skip to next block

  def readObj(self, f, myLen):
    while f.tell() < myLen:
      name = readStr(f)
      [hdr, size] = readHdr(f)
      end = f.tell() + size - 6
      if hdr == TRIMESH:
        self.objects.append(TriObj(name, f, end))      # new object
      elif hdr == LIGHT:
#        print 'size: ' + str(size)
        self.lights.append(Light(name, f, end))
        pass
      else: 
#        print hex(hdr)
        f.read(size - 6)

  def calcNormals(self):
    for m in self.objects:
      m.calcNormals()

  def setMaterials(self):
    for obj in self.objects:
      try:
        obj.material = self.materials[obj.matName]
      except Exception, e:
        # use default material
        print ' using default material!!!'
        obj.material = Material()
      

  def objNames(self):
    objs = []
    for obj in self.objects:
      objs.append(obj.name)
    return objs

  def getIndex(self, objName):
    for i in range(0, len(self.objects)):
      obj = self.objects[i]
      if obj.name == objName:
        return i

  def draw(self, index = None):
    from OpenGL.GL import *

    if index is None:
      for i in range(0, len(self.objects)):
        self.objects[i].fastDraw()

    else:
      self.objects[i].fastDraw()

  def lighting(self):
    from OpenGL.GL import *

    if len(self.lights) == 0: return None

    glLightModeli(GL_LIGHT_MODEL_LOCAL_VIEWER, GL_TRUE)

    cnt = 0
    for light in self.lights[:8]:
      ambient = [light.color[0] * .1, light.color[1] * .1, light.color[1] * .1]
      diffuse = [light.color[0] * .6, light.color[1] * .6, light.color[1] * .6]
      specular = [light.color[0] * .1, light.color[1] * .1, light.color[1] * .1]
      glLightf(GL_LIGHT0 + cnt , GL_AMBIENT, ambient)
      glLightf(GL_LIGHT0 + cnt , GL_DIFFUSE, diffuse)
      glLightf(GL_LIGHT0 + cnt , GL_SPECULAR, specular)
      glLightf(GL_LIGHT0, GL_POSITION, light.pos)
      cnt = cnt + 1 

    return 'true'
