'''
Hack to get shave in .vrscenes
Zach Gray
SVAD Animation Prof
zcgray@southern.edu


Basic idea is to write out .dra files which contain all the baked shave info
Then, hook those files up to the geo/shader/anim later in the .vrscene
'''



'''
resets callbacks to shave defaults. this will cause issues if you have other callbacks running
shave has a render time preKeyframe callback shaveRender -createVrayKeyframe that needs to be disabled with \
another hacky pre frame callback that clears that one. terrible, I know.
'''
#shaveRenderOptions(clear = True)

'''
use pre calculated .dra files made with shaveDraWriter
adds a callback to clear the per frame shave calc
'''
#shaveRenderOptions(writeDra= False)

'''
Write out dra files during render
Attempts to write out the files as it's writing the .vrscene
using a pre frame callback
not very stable
'''
#shaveRenderOptions(writeDra=True)

'''
Utility to write out dra files manually based on .vrscene path
'''
#shaveDraWriter() or shaveDraWriter(start=1, end=10)



import os as os
import pymel.core as pm
import maya.cmds as mc
import maya.mel as mel
from pprint import pprint
from functools import partial


def shaveRenderGetPaths(nameFromScene = True):
    '''
    utility to get paths to dra files
    makes a directory based on .vrscene location if it doesn't exist
    '''
    vrscene = mc.getAttr('vraySettings.vrscene_filename')
    path = os.path.dirname(vrscene)
    vrsceneFile = os.path.basename(vrscene)
    shavePath = path + '/shave_dra/'
    
    if not os.path.isdir(shavePath):
        try:
            os.mkdir(shavePath)
        except:
            mc.warning ("Shave DRA writer: cannot make directory")
                        
    name = 'shave_dra.'
    if nameFromScene:
        name = vrsceneFile.rsplit('.',1)[0] + '.'
    ext = '.dra'
    
    paths = {}
    paths['vrscene'] = vrscene
    paths['shavePath'] = shavePath
    paths['shaveName'] = name
    paths['shaveExt'] = ext
    
    return paths

def shaveRenderOptions(writeDra = True, processScene=True, disableShave = True, clear = False):
    '''
    sets and resets the render callbacks
    doesn't play nice with other callbacks (need to fix because it's causing problems)
    '''
    #Get the attributes from the vray settings

    p = shaveRenderGetPaths()
       
    
    if processScene:
        postMelOrig = mc.getAttr('defaultRenderGlobals.postMel')
        mc.setAttr('defaultRenderGlobals.postMel', 'shaveVrayPostRender; python("shaveDraVrscene()");', type = 'string')
    else:
        mc.setAttr('defaultRenderGlobals.postMelOrig', 'shaveVrayPostRender;', type = 'string')
        
        
        
    
    if disableShave:
        postFrameMel = 'setAttr -type "string" vraySettings.preKeyframeMel "";'
        mc.setAttr('defaultRenderGlobals.postRenderMel', postFrameMel, type='string')
    else:
        mc.setAttr('defaultRenderGlobals.postRenderMel', '', type = 'string')
        
    
    if writeDra:
        preMel = 'currentTime -e (getAttr("defaultRenderGlobals.startFrame") -1); int $f = `currentTime -query`; shaveRender -createDRAFile ("%s%s" + ($f+1) + "%s");' % (p['shavePath'], p['shaveName'], p['shaveExt'])
        preFrameMel = 'int $f = `currentTime -query`; shaveRender -createDRAFile ("%s%s" + ($f+1) + "%s");' % (p['shavePath'], p['shaveName'], p['shaveExt'])
        mc.setAttr('defaultRenderGlobals.preRenderMel', preFrameMel, type = 'string')
        mc.setAttr('defaultRenderGlobals.preMel', ('shaveVrayPreRender;' +  preMel), type = 'string')
    else:
         mc.setAttr('defaultRenderGlobals.preRenderMel', '', type = 'string')
         mc.setAttr('defaultRenderGlobals.preMel', 'shaveVrayPreRender;', type = 'string')
         
    if clear:
        mc.setAttr('defaultRenderGlobals.postMel', 'shaveVrayPostRender;', type = 'string')
        mc.setAttr('defaultRenderGlobals.preMel', 'shaveVrayPreRender;', type = 'string')
        mc.setAttr('defaultRenderGlobals.preRenderMel', '', type = 'string')
        mc.setAttr('defaultRenderGlobals.postRenderMel', '', type = 'string')    
        
    
        

      

def shaveDraWriter(start=None,end=None):
    '''
    writes out dra files based on vrscene path
    able to set start and end, if not defaults to render settings
    '''
    
    p = shaveRenderGetPaths()
    
    gMainProgressBar = mel.eval('$tmp = $gMainProgressBar');
    
    
    if start is None:
        start = mc.getAttr('defaultRenderGlobals.startFrame')
    if end is None:
        end = mc.getAttr('defaultRenderGlobals.endFrame')
    
    progressControl = mc.progressBar(maxValue=abs(end-start), width=300)
    
    mc.progressBar( gMainProgressBar,
        edit=True,
        beginProgress=True,
        isInterruptable=True,
        status='"Exporting Dra ...',
        maxValue=abs(end-start) )
    
    
    #add step at some point
    for frame in range(start-1,end):
        mc.currentTime(frame, edit=True)
        mc.refresh;
        shaveRenderString = 'shaveRender -createDRAFile ("%s%s" + (%s) + "%s");' %(p['shavePath'], p['shaveName'], frame+1, p['shaveExt'])
        print shaveRenderString
        mel.eval(shaveRenderString)
        if mc.progressBar(gMainProgressBar, query=True, isCancelled=True ) :
            break
        mc.progressBar(gMainProgressBar, edit=True, step=1)
        
    cmds.progressBar(gMainProgressBar, edit=True, endProgress=True)
    

        
        
            
    
def shaveDraVrscene():    
    '''
    Runs as a post render callback
    looks for .dra files in path
    tries to put them in the right place in vrscene
    
    uses a messy nested loop for speed in reading big files
    '''
    p = shaveRenderGetPaths()
       
    
    #scene variables
    shaveHair = []
    firstFrame = True
    currentFrame = 0
    
    #don't indent this block!
    codeBlock = '''
  draFile=interpolate(
    (%s, "%s")
    );
}
'''
    
        
    if not os.path.isfile (p['vrscene']):
        return
    
    
    sourceName = p['vrscene']
    destName = p['vrscene']+'_temp'
    source = open(sourceName, 'r')
    dest = open(destName, 'w')
    
    block = False
    frameBlock = False
    
    
    for line in source:
    
        if '// Frame at time' in line and firstFrame:
            firstFrame = False
            currentFrame = line.rsplit(' ', 1)[1].rstrip()
            dest.write(line)
            continue
        #hair shader block    
        elif 'hairVRayShader' in line:
            print line
            if line not in shaveHair:
                shaveHair.append(line)
            
            #print '    block begin'
            block = True
            dest.write(line)
            continue
        if '}' in line and block:
            #print '    block end'
            block = False
            dest.write(line)
            continue    
        elif 'draFile=""' in line and block:
            draFile = (p['shavePath'] + p['shaveName'] + currentFrame + p['shaveExt'])
            if os.path.isfile(draFile) is None:
                mc.warning ("Shave DRA vrscene: file %s doesn't exist!") %(draFile)
            #print '        file'
            dest.write ('  draFile="' + draFile +'";\n')
            continue
        elif 'draData' in line and block:
            #print '        data'
            write = False
            continue
        elif 'draSize' in line and block:
            #print '        size'
            write = False
            continue
        
            
        if '// Frame at time' in line:
            frameBlock = True
            currentFrame = line.rsplit(' ', 1)[1].rstrip()
            dest.write(line)
            continue
            
            
            
        if frameBlock:
            currentIndex = int(currentFrame)
            draFile = (p['shavePath'] + p['shaveName'] + currentFrame + p['shaveExt'])
            if os.path.isfile (draFile) is None:
                mc.warning ("Shave DRA vrscene: file %s doesn't exist!") %(draFile)
             
            
            for plugin in shaveHair:
                dest.write('\n')
                dest.write(plugin.rstrip())
                dest.write(codeBlock % (currentIndex-1, draFile))
            dest.write('\n')
            frameBlock = False
            continue
                    
                
        dest.write(line)
    
    
    source.close()
    dest.close()
    
    sourceNameBak = sourceName+'_bak'
    if os.path.isfile (sourceName):
        if os.path.isfile (sourceNameBak):
            os.remove (sourceNameBak)
        os.rename(sourceName, sourceNameBak)
    if os.path.isfile (destName):
        os.rename(destName, sourceName)
        
        
        
     
    print 'Shave hair nodes: '
    print shaveHair

    mc.warning ("Shave DRA vrscene updated with file refs")