#----------------------------------------------------------------
# $Revision: #12 $
#----------------------------------------------------------------
# This file may be copied.  However, if you make improvements,
# please send me a copy.
#
# Brent Burton
# AustinMTB@yahoo.com
# July 2003
# RaceSimCentral: vizor_atx
#----------------------------------------------------------------
# DOFfile manipulation code.
#

import struct

#---------------- Module version info
version = 1.00

#---------------- Utility functions

# Write a 2-byte int value to the file, Little-Endian
def writeShortLE(file, val):
    file.write(struct.pack('<h', val))

# Write a 4-byte int value to the file, Little-Endian
def writeIntLE(file, val):
    file.write(struct.pack('<l', val))

# Write a 4-byte float value to the file
def writeFloat(file, val):
    file.write(struct.pack('f', val))

# Write a 4-byte float vector of 2 elements to the file
def writeFloat2(file, val2):
    file.write(struct.pack('ff', val2[0], val2[1]))

# Write a 4-byte float vector of 3 elements to the file
def writeFloat3(file, val3):
    file.write(struct.pack('fff', val3[0], val3[1], val3[2]))

# Write a 4-byte float vector of 4 elements to the file
def writeFloat4(file, val4):
    file.write(struct.pack('ffff', val4[0], val4[1], val4[2], val4[3]))

# Read a little endian 2-byte short
def readShortLE(file):
    return struct.unpack('<h', file.read(2))[0]

# Read a little endian 4-byte int
def readIntLE(file):
    return struct.unpack('<l', file.read(4))[0]

# Read a vector of little endian 4-byte ints
def readIntLEVec(file, num):
    iv = []
    for i in range(num): iv.append(readIntLE(file))
    return iv

# Read a 4-byte float value
def readFloat(file):
    return struct.unpack('f', file.read(4))[0]

# Read a 2 float vector, return as list
def readFloat2(file):
    f = [readFloat(file)]
    f.append(readFloat(file))
    return f

# Read a 3 float vector, return as list
def readFloat3(file):
    f = [readFloat(file)]
    f.append(readFloat(file))
    f.append(readFloat(file))
    return f

# Read a 4 float vector, return as list
def readFloat4(file):
    f = [readFloat(file)]
    f.append(readFloat(file))
    f.append(readFloat(file))
    f.append(readFloat(file))
    return f

# Read a QString type, return a QString object
def readQString(file):
    qs = QString("")
    qs.read(file)
    return qs

# Read a chunk tag, return the 4-byte string
def nextChunk(file):
    return file.read(4)

# Skip Over uninteresting data
def skipOver(file, dist):
    file.seek(dist, 1)                  # current(1) + dist

#---------------- Class definitions

#----------------
class DOFBase:
    def __init__(self):
        self._size = -1                 # -1 means recalculate

    def totalSize(self):
        # Total size is (4-byte tag + 4-byte int length + length bytes)
        return 8 + self.size()

#----------------
class DOF1(DOFBase):
    def __init__(self):
        self._size = -1
        self._MATS = MATS()
        self._GEOB = GEOB()

    def addMAT0(self, mat0):
        self._size = -1
        self._MATS.addMAT0(mat0)

    def numMAT0s(self):
        return self._MATS.numMAT0s()

    def getMAT0(self, idx):
        return self._MATS.getMAT0(idx)

    def addGOB1(self, gob1):
        self._size = -1
        self._GEOB.addGOB1(gob1)

    def numGOB1s(self):
        return self._GEOB.numGOB1s()

    def getGOB1(self, idx):
        return self._GEOB.getGOB1(idx)

    def size(self):
        if self._size == -1:
            self._size = self._MATS.totalSize() + self._GEOB.totalSize()
        return self._size

    def write(self, file):
        file.write("DOF1")
        writeIntLE(file, self.size())
        self._MATS.write(file)
        self._GEOB.write(file)
        file.write("EDOF")

    def read(self, file):
        length = readIntLE(file)
        done = 0
        while not done:
            chunk = nextChunk(file)
            if chunk == "MATS":
                self._MATS.read(file)
            elif chunk == "GEOB":
                self._GEOB.read(file)
            elif chunk == "EDOF":
                done = 1
            else:
                print "Unrecognized chunk in DOF1: %s" % (chunk)
                skipOver(file, readIntLE(file))

#----------------
class MATS(DOFBase):
    def __init__(self):
        self._size = -1
        self._MATS = []

    def addMAT0(self, mat0):
        self._MATS.append(mat0)

    def numMAT0s(self):
        return len(self._MATS)

    def getMAT0(self, idx):
        if idx >= len(self._MATS): return None
        return self._MATS[idx]

    def size(self):
        if self._size == -1:
            self._size = 4
            for mat0 in self._MATS:
                # Add the +4 here since each MAT0 appends a MEND tag
                self._size += mat0.totalSize() + 4
        return self._size

    def write(self, file):
        file.write("MATS")
        writeIntLE(file, self.size())
        writeIntLE(file, len(self._MATS))
        for mat in self._MATS:
            mat.write(file)

    def read(self, file):
        length = readIntLE(file)
        numMats = readIntLE(file)
        for i in range(numMats):
            chunk = nextChunk(file)
            if chunk == "MAT0":
                mat0 = MAT0()
                mat0.read(file)
                self.addMAT0(mat0)
            else:
                print "Unrecognized chunk in MATS block: %s" % (chunk)
                skipOver(file, readIntLE(file))

#----------------
class GEOB(DOFBase):
    def __init__(self):
        self._size = -1
        self._GEOB = []

    def addGOB1(self, gob1):
        self._GEOB.append(gob1)

    def numGOB1s(self):
        return len(self._GEOB)

    def getGOB1(self, idx):
        if idx >= len(self._GEOB): return None
        return self._GEOB[idx]

    def size(self):
        if self._size == -1:
            self._size = 4
            for gob in self._GEOB:
                # Add the +4 here since each GOB1 appends a GEND tag
                self._size += gob.totalSize() + 4
        return self._size

    def write(self, file):
        file.write("GEOB")
        writeIntLE(file, self.size())
        writeIntLE(file, len(self._GEOB))
        for gob in self._GEOB:
            gob.write(file)

    def read(self, file):
        self._size = readIntLE(file)
        numgobs = readIntLE(file)
        for i in range(numgobs):
            chunk = nextChunk(file)
            if chunk == "GOB1":
                gob1 = GOB1()
                gob1.read(file)
                self.addGOB1(gob1)
            else:
                print "Unrecognized chunk in GEOB block: %s" % (chunk)
                skipOver(file, readIntLE(file))
        
#----------------
class QString:
    def __init__(self, str):
        self._string = str

    def size(self):
        return len(self._string) + 2

    def write(self, file):
        writeShortLE(file, len(self._string))
        file.write(self._string)

    def read(self, file):
        length = readShortLE(file)
        self._string = file.read(length)

#----------------
class MAT0(DOFBase):
    def __init__(self):
        self._size = -1
        self._mhdr = None
        self._mcol = None
        self._mcfl = None
        self._muvw = None
        self._mtra = None
        self._mtex = None
        self._msub = None

    # Set* methods
    def setMHDR(self, mhdr):
        self._mhdr = mhdr

    def setMCOL(self, mcol):
        self._mcol = mcol

    def setMCFL(self, mcfl):
        self._mcfl = mcfl

    def setMUVW(self, muvw):
        self._muvw = muvw

    def setMTRA(self, mtra):
        self._mtra = mtra

    def setMTEX(self, mtex):
        self._mtex = mtex

    def setMSUB(self, msub):
        self._msub = msub

    # Get* methods
    def getMTRA(self):
        return self._mtra

    def size(self):
        if self._size == -1:
            sz = 0
            if self._mhdr: sz += self._mhdr.totalSize()
            if self._mcol: sz += self._mcol.totalSize()
            if self._mcfl: sz += self._mcfl.totalSize()
            if self._muvw: sz += self._muvw.totalSize()
            if self._mtra: sz += self._mtra.totalSize()
            if self._mtex: sz += self._mtex.totalSize()
            if self._msub: sz += self._msub.totalSize()
            self._size = sz
        return self._size

    def write(self, file):
        file.write("MAT0")
        writeIntLE(file, self.size())
        # write each subobject, mhdr, etc
        # These are the same order as Racer's Modeler writes
        if self._mhdr: self._mhdr.write(file)
        if self._mcol: self._mcol.write(file)
        if self._muvw: self._muvw.write(file)
        if self._mtra: self._mtra.write(file)
        if self._mcfl: self._mcfl.write(file)
        if self._mtex: self._mtex.write(file)
        if self._msub: self._msub.write(file)
        file.write("MEND")

    def read(self, file):
        length = readIntLE(file)        # ignored
        done = 0
        while not done:
            chunk = nextChunk(file)
            if chunk == "MHDR":
                self._mhdr = MHDR()
                self._mhdr.read(file)
            elif chunk == "MCOL":
                self._mcol = MCOL()
                self._mcol.read(file)
            elif chunk == "MCFL":
                self._mcfl = MCFL()
                self._mcfl.read(file)
            elif chunk == "MUVW":
                self._muvw = MUVW()
                self._muvw.read(file)
            elif chunk == "MTRA":
                self._mtra = MTRA()
                self._mtra.read(file)
            elif chunk == "MTEX":
                self._mtex = MTEX()
                self._mtex.read(file)
            elif chunk == "MSUB":
                print "Submaterials ignored"
                skipOver(file, readIntLE(file))
            elif chunk == "MEND":
                done = 1
            else:
                print "Unknown chunk in MAT0...ignoring"
                skipOver(file, readIntLE(file))

#----------------
class GOB1(DOFBase):
    def __init__(self):
        self._size = -1
        self._ghdr = None
        self._indi = None
        self._vert = None
        self._vcol = None
        self._tver = None
        self._norm = None
        self._brst = None

    def size(self):
        if self._size == -1:
            sz = 0
            if self._ghdr: sz += self._ghdr.totalSize()
            if self._indi: sz += self._indi.totalSize()
            if self._vert: sz += self._vert.totalSize()
            if self._vcol: sz += self._vcol.totalSize()
            if self._tver: sz += self._tver.totalSize()
            if self._norm: sz += self._norm.totalSize()
            if self._brst: sz += self._brst.totalSize()
            self._size = sz
        return self._size

    def setGHDR(self, ghdr):
        self._ghdr = ghdr

    def setINDI(self, indi):
        self._indi = indi

    def setVERT(self, vert):
        self._vert = vert

    def setVCOL(self, vcol):
        self._vcol = vcol

    def setTVER(self, tver):
        self._tver = tver

    def setNORM(self, norm):
        self._norm = norm

    def setBRST(self, brst):
        self._brst = brst

    def getBRST(self):
        return self._brst

    def write(self, file):
        file.write("GOB1")
        writeIntLE(file, self.size())
        # write each subobject, ghdr, indi, vert, etc
        if self._ghdr: self._ghdr.write(file)
        if self._indi: self._indi.write(file)
        if self._vert: self._vert.write(file)
        if self._vcol: self._vcol.write(file)
        if self._tver: self._tver.write(file)
        if self._norm: self._norm.write(file)
        if self._brst: self._brst.write(file)
        file.write("GEND")

    def read(self, file):
        readIntLE(file)                 # size ignored
        self._size = -1
        done = 0
        while not done:
            chunk = nextChunk(file)
            if chunk == "GHDR":
                self._ghdr = GHDR()
                self._ghdr.read(file)
            elif chunk == "INDI":
                self._indi = INDI()
                self._indi.read(file)
            elif chunk == "VERT":
                self._vert = VERT()
                self._vert.read(file)
            elif chunk == "VCOL":
                self._vcol = VCOL()
                self._vcol.read(file)
            elif chunk == "TVER":
                self._tver = TVER()
                self._tver.read(file)
            elif chunk == "NORM":
                self._norm = NORM()
                self._norm.read(file)
            elif chunk == "BRST":
                self._brst = BRST()
                self._brst.read(file)
            elif chunk == "GEND":
                done = 1
            else:
                print "Unrecognized chunk in GOB1 block: %s skipped" % (chunk)
                skipOver(file, readIntLE(file))


#----------------
class MHDR(DOFBase):
    def __init__(self):
        self._size = -1
        self._name = QString("")
        self._classname = QString("Standard")

    def setName(self, name):
        self._name = QString(name)

    def setClassname(self, cname):
        self._classname = QString(cname)

    def size(self):
        return self._name.size() + self._classname.size()

    def write(self, file):
        file.write("MHDR")
        writeIntLE(file, self.size())
        self._name.write(file)
        self._classname.write(file)

    def read(self, file):
        self._size = readIntLE(file)
        self._name = readQString(file)
        self._classname = readQString(file)

#----------------
class MCOL(DOFBase):
    def __init__(self):
        self._size = -1
        self._ambient   = [0.0, 0.0, 0.0, 0.0]
        self._diffuse   = [0.0, 0.0, 0.0, 0.0]
        self._specular  = [0.0, 0.0, 0.0, 0.0]
        self._emission  = [0.0, 0.0, 0.0, 0.0]
        self._shininess = 0.0

    def size(self):
        return 17*4

    def setAmbient(self, amb):
        self._ambient = amb

    def setDiffuse(self, diff):
        self._diffuse = diff

    def setSpecular(self, spec):
        self._specular = spec

    def setEmission(self, emis):
        self._emission = emis

    def setShininess(self, shine):
        self._shininess = shine

    def write(self, file):
        file.write("MCOL")
        writeIntLE(file, self.size())
        writeFloat4(file, self._ambient)
        writeFloat4(file, self._diffuse)
        writeFloat4(file, self._specular)
        writeFloat4(file, self._emission)
        writeFloat(file,  self._shininess)

    def read(self, file):
        readIntLE(file)                 # ignored
        self._ambient = readFloat4(file)
        self._diffuse = readFloat4(file)
        self._specular = readFloat4(file)
        self._emission = readFloat4(file)
        self._shininess = readFloat(file)

#----------------
class MCFL(DOFBase):
    def __init__(self):
        self._flags = 0

    def setFlag(self, bit):
        self._flags = self._flags | (1 << bit)
        
    def size(self):
        return 4

    def write(self, file):
        file.write("MCFL")
        writeIntLE(file, self.size())
        writeIntLE(file, self._flags)

    def read(self, file):
        readIntLE(file)                 # size ignored
        self._flags = readIntLE(file)

#----------------
class MUVW(DOFBase):
    def __init__(self):
        self._UVoffset   = [0, 0]
        self._UVtiling   = [0, 0]
        self._angle      = 0
        self._blur       = 0
        self._blurOffset = 0

    def setUVOffset(self, uv):
        self._UVoffset = uv

    def setUVTiling(self, uv):
        self._UVtiling = uv

    def setAngle(self, a):
        self._angle = a

    def setBlur(self, b):
        self._blur = b

    def setBlurOffset(self, bo):
        self._blurOffset = bo

    def size(self):
        return 28

    def write(self, file):
        file.write("MUVW")
        writeIntLE(file, self.size())
        writeFloat(file, self._UVoffset[0])
        writeFloat(file, self._UVoffset[1])
        writeFloat(file, self._UVtiling[0])
        writeFloat(file, self._UVtiling[1])
        writeFloat(file, self._angle)
        writeFloat(file, self._blur)
        writeFloat(file, self._blurOffset)

    def read(self, file):
        readIntLE(file)                 # size ignored
        self._UVoffset = readFloat2(file)
        self._UVtiling = readFloat2(file)
        self._angle = readFloat(file)
        self._blur = readFloat(file)
        self._blurOffset = readFloat(file)

#----------------
class MTRA(DOFBase):
    def __init__(self):
        self._transparency = 0
        self._blendmode = 0

    def setTransparency(self, t):
        self._transparency = t

    def setBlendmode(self, bm):
        self._blendmode = bm

    def size(self):
        return 8

    def write(self, file):
        file.write("MTRA")
        writeIntLE(file, self.size())
        writeFloat(file, self._transparency)
        writeIntLE(file, self._blendmode)

    def read(self, file):
        readIntLE(file)                 # size ignored
        self._transparency = readFloat(file)
        self._blendmode = readIntLE(file)

#----------------
class MTEX(DOFBase):
    def __init__(self):
        self._size = -1
        self._texnames = []

    def size(self):
        if self._size == -1:
            sz = 4
            for qstr in self._texnames:
                sz += qstr.size()
            self._size = sz
        return self._size

    def append(self, name):
        self._texnames.append(QString(name))

    def write(self, file):
        file.write("MTEX")
        writeIntLE(file, self.size())
        writeIntLE(file, len(self._texnames))
        for qstr in self._texnames:
            qstr.write(file)

    def read(self, file):
        self._size = readIntLE(file)
        numnames = readIntLE(file)
        self._texnames = []
        for i in range(numnames):
            self._texnames.append(readQString(file))

#----------------
class GHDR(DOFBase):
    def __init__(self, f=0, pf=0, mr=0):
        self._flags = f
        self._paintFlags = pf
        self._matRef = mr

    def size(self):
        return 12

    def write(self, file):
        file.write("GHDR")
        writeIntLE(file, self.size())
        writeIntLE(file, self._flags)
        writeIntLE(file, self._paintFlags)
        writeIntLE(file, self._matRef)

    def read(self, file):
        readIntLE(file)                 # size ignored
        self._flags = readIntLE(file)
        self._paintFlags = readIntLE(file)
        self._matRef = readIntLE(file)

#----------------
class INDI(DOFBase):
    def __init__(self):
        self._size = -1
        self._indices = []

    def size(self):
        # size = (1-short) * 2bytes * count + sizeof(count)
        return 2 * len(self._indices) + 4

    def append(self, index):
        self._indices.append(index)

    def write(self, file):
        file.write("INDI")
        writeIntLE(file, self.size())
        writeIntLE(file, len(self._indices))
        for idx in self._indices:
            writeShortLE(file, idx)

    def read(self, file):
        self._size = readIntLE(file)
        num = readIntLE(file)
        self._indices = []
        for i in range(num):
            self._indices.append(readShortLE(file))

#----------------
class VERT(DOFBase):
    def __init__(self):
        self._size = -1
        # array of 3-element XYZ vectors
        self._verts = []

    def size(self):
        # size = (3-floats)* 4bytes * count + sizeof(count)
        return len(self._verts) * 12 + 4

    def append(self, xyz):
        """Append new XYZ coords to this group"""
        self._verts.append(xyz)

    def write(self, file):
        file.write("VERT")
        writeIntLE(file, self.size())
        writeIntLE(file, len(self._verts))
        for v in self._verts:
            writeFloat3(file, v)

    def read(self, file):
        self._size = readIntLE(file)
        num = readIntLE(file)
        self._verts = []
        for i in range(num):
            self._verts.append(readFloat3(file))

#----------------
class VCOL(DOFBase):
    def __init__(self):
        self._size = -1
        # array of 3-element RGB vectors
        self._vcols = []

    def size(self):
        # size = (3-float) * 4bytes * count + sizeof(count)
        return len(self._vcols) * 12

    def append(self, rgb):
        """Append new RGB color to this group"""
        self._vcols.append(rgb)

    def write(self, file):
        file.write("VCOL")
        writeIntLE(file, self.size())
        writeIntLE(file, len(self._vcols))
        for v in self._vcols:
            writeFloat3(file, v)

    def read(self, file):
        self._size = readIntLE(file)
        num = readIntLE(file)
        self._vcols = []
        for i in range(num):
            self._vcols.append(readFloat3(file))

#----------------
class TVER(DOFBase):
    def __init__(self):
        self._size = -1
        # array of 2-element UV vectors
        self._uv = []

    def size(self):
        # size = (2-float) * 4bytes * count + sizeof(count)
        return len(self._uv) * 8 + 4

    def append(self, uv):
        self._uv.append(uv)

    def write(self, file):
        file.write("TVER")
        writeIntLE(file, self.size())
        writeIntLE(file, len(self._uv))
        for uv in self._uv:
            writeFloat2(file, uv)

    def read(self, file):
        self._size = readIntLE(file)
        num = readIntLE(file)
        self._uv = []
        for i in range(num):
            self._uv.append(readFloat2(file))

#----------------
class NORM(DOFBase):
    def __init__(self):
        self._size = -1
        # array of 3-element IJK vectors
        self._norms = []

    def size(self):
        # size = (3-float) * 4bytes * count + sizeof(count)
        return len(self._norms) * 12 + 4

    def append(self, xyz):
        """Append new XYZ coords to this group"""
        self._norms.append(xyz)

    def write(self, file):
        file.write("NORM")
        writeIntLE(file, self.size())
        writeIntLE(file, len(self._norms))
        for n in self._norms:
            writeFloat3(file, n)

    def read(self, file):
        self._size = readIntLE(file)
        num = readIntLE(file)
        self._norms = []
        for i in range(num):
            self._norms.append(readFloat3(file))

#----------------
class BRST(DOFBase):
    def __init__(self):
        self._size = -1
        self._numBursts = 0
        self._burstStart = []
        self._burstCount = []
        self._burstMat   = []
        self._burstVperP = []

    def size(self):
        return self._numBursts * 4 * 4

    def addBurst(self, s, c, m, vpp):
        self._numBursts += 1
        self._burstStart.append(s)
        self._burstCount.append(c)
        self._burstMat.append(m)
        self._burstVperP.append(vpp)

    def addMaterialOffset(self, delta):
        """For each burst, add delta to the material index."""
        for i in range(len(self._burstMat)):
            self._burstMat[i] += delta

    def write(self, file):
        file.write("BRST")
        writeIntLE(file, self.size())
        writeIntLE(file, self._numBursts)
        for i in self._burstStart: writeIntLE(file, i)
        for i in self._burstCount: writeIntLE(file, i)
        for i in self._burstMat:   writeIntLE(file, i)
        for i in self._burstVperP: writeIntLE(file, i)

    def read(self, file):
        self._size = readIntLE(file)
        self._numBursts  = readIntLE(file)
        self._burstStart = readIntLEVec(file, self._numBursts)
        self._burstCount = readIntLEVec(file, self._numBursts)
        self._burstMat   = readIntLEVec(file, self._numBursts)
        self._burstVperP = readIntLEVec(file, self._numBursts)

#----------------
def read(file):
    """Read a DOF file and build internal rep of it.
    This function assumes there is only one DOF object in a file.
    """
    file.seek(0, 0)                     # go to file start
    chunk = nextChunk(file)
    if chunk == "DOF1":
        dof1 = DOF1()
        dof1.read(file)
        return dof1
    else:
        print "Unexpected chunk found: %s" % (chunk)

#----------------
if __name__ == "__main__":
    dof1 = DOF1()
    gob1 = GOB1()
    dof1.addGOB1(gob1)
    del gob1

    outf = file("sample.dof", "wb")
    dof1.write(outf)
    outf.flush()
    outf.close()

    del dof1
