import maya.cmds as mc
  
class Cache:
    def __init__(self, width):
        self.data = []
        self.width = width
        self.minX = 10000
        self.maxX = -10000
        self.minY = 10000
        self.maxY = -10000
        self.minZ = 10000
        self.maxZ = -10000
    def reset(self, width):
        self.data = []
        self.width = width
    def add(self, index, pos):
        if len(self.data) > index:
            xyz = pos[0:3]
            self.data[index].append(xyz[0])
            self.data[index].append(xyz[1])
            self.data[index].append(xyz[2])
        else:
            self.data.append(pos[0:3])
    def get(self, index):
        return self.data[index]
    def getFirstXYZ(self, index):
        x = self.data[index][0]
        y = self.data[index][1]
        z = self.data[index][2]
        return x,y,z
    def getLastXYZ(self, index):
        n = len(self.data[index])
        x = self.data[index][n - 3]
        y = self.data[index][n - 2]
        z = self.data[index][n - 1]
        return x,y,z
    def length(self):
        return len(self.data)
    def update(self, tnode):
        pnum = mc.particle(tnode, q = True, count = True)
        for n in range(pnum):
            pname = tnode + ".pt[%s]" % n
            pos = mc.getParticleAttr(pname,at = 'position')
            if pos[0] < self.minX:
                self.minX = pos[0]
            if pos[0] > self.maxX:
                self.maxX = pos[0]
            if pos[1] < self.minY:
                self.minY = pos[1]                
            if pos[1] > self.maxY:
                self.maxY = pos[1]
            if pos[2] < self.minZ:
                self.minZ = pos[2]
            if pos[2] > self.maxZ:
                self.maxZ = pos[2]
            self.add(n, pos) 
    def getBBox(self):
        return '#bbox: %1.3f %1.3f %1.3f %1.3f %1.3f %1.3f' % (self.minX, self.minY, self.minZ, self.maxX, self.maxY, self.maxZ)
cacheDB = { } 
#-----------------------------------------    
def getSceneName():
    name = mc.file(q = True, sceneName = True, shortName = True)
    if len(name) == 0:
        name = "untitled"
    else:
        name = name[:len(name) - 3]
    return name
#-----------------------------------------
def getDataDir():
    projPath = mc.workspace(q = True, rootDirectory = True)
    return projPath + "data"
#-----------------------------------------        
def bakeParticleCurves(tnode, start, end, width, curveType):
    global cacheDB
    if cacheDB.has_key(tnode):
        cache = cacheDB.get(tnode)
    else:
        cache = Cache(width)
        cacheDB[tnode] = cache
    currFrame = mc.currentTime(q = True)
    if currFrame == 1:
        cache.reset(width)
    if currFrame >= start and currFrame <= end:
        cache.update(tnode)
    pathToCurves = getDataDir() + "/" + getSceneName() + "." + tnode + "." + str(int(currFrame)) + ".rib";
    fileid = open(pathToCurves, 'w')
    fileid.write(cache.getBBox() + '\n')
    fileid.write('AttributeBegin\n')
    fileid.write('Attribute "dice" "hair" [1]\n')
    fileid.write('Attribute "stochastic" "int sigma" [1]\n')
    fileid.write('Basis "catmull-rom" 1 "catmull-rom" 1\n')
    for n in range(cache.length()):
        xyz = cache.get(n)
        if curveType == 'cubic' and len(xyz)/3 >= 4:
            rib = 'Curves "cubic" [%d] "nonperiodic" "P" [\n' % (len(xyz)/3)
            fileid.write(rib)
            for i in range(len(xyz)):
                fileid.write("%1.3f " % (xyz[i]))
            rib = '\n] "constantwidth" [%1.4f] \n' % width
            fileid.write(rib)
        elif curveType == 'linear' and len(xyz)/3 >= 2:
            rib = 'Curves "linear" [2] "nonperiodic" "P" ['
            rib = rib + '%1.3f %1.3f %1.3f ' % cache.getFirstXYZ(n)
            rib = rib + '%1.3f %1.3f %1.3f]' % cache.getLastXYZ(n)
            rib = rib + ' "constantwidth" [%1.4f]\n' % width
            fileid.write(rib)
    fileid.write('AttributeEnd\n')
    fileid.close()