#!usr/bin/python

#####################################################
# Author: Kenneth Stojevich                         #
# Program: Maya Command Line Batch Renderer         #
# Description: Allows the user to apply simple MEL  #
#    commands to a file and then save it.  You can  #
#    also do a command line render of the scene     #
#    supplied as a command line argument(s).  In    #
#    addition to rendering the user can give MEL    #
#    commands to the scene to be applied to it      #
#    either pre-render or pre-frame.                #
# Date: 03/11/2006                                  #
#####################################################

# Import the required modules needed for the script
import getopt, sys, re, os, commands

class MayaBatch:
    	##############################################
    	# Constructor                                #
    	##############################################
    	def __init__(self, mayaScene = None):
		self.MAYARENDER  = 0
        	self.dispRender  = 0
        
        	self.melCmd      = ""
        
        	self.firstFrame  = 1
        	self.lastFrame   = 10
        
        	self.blurFactor  = 0
        	self.blurOptions = ""
        
        	self.image       = "myImages"
        	self.imageDir    = "/images/"
        
        	self.scene       = mayaScene

    	##############################################
	# This is a setup function to read in all of #
	# the command line arguments the user gave so#
    	# we know what to do.			     #
    	##############################################
    	def setupMaya(self,arguments=[]):
		# Get our arguments and set our required arguments accordingly
		try:
			opts, args = getopt.getopt(arguments, "hm:rf:l:b:i:d", ["help", "mel=", "render", "first=", "last=", "blur=", "image=", "display"])
    		except getopt.GetoptError:
			print "Incorrect argument"
        		sys.exit()
	
		for o, a in opts:
			if (o == "-h" or o == "--help"):
		     		self.helpMe()
		      		sys.exit(0)
	    		elif (o == "-m" or o == "--mel"):
	      			self.melCommand = a
			elif (o == "-r" or o == "--render"):
				self.MAYARENDER = 1
			elif (o == "-f" or o == "--first"):
		        	self.firstFrame = a
			elif (o == "-l" or o == "--last"):
		        	self.lastFrame = a
			elif (o == "-b" or o == "--blur"):
		        	self.blurFactor = a
		        	self.blurOptions = "-mb 1 -mbf " + a
			elif (o == "-d" or o == "--display"):
		        	self.dispRender = 1
			elif (o == "-i" or o == "--image"):
				self.image = a

		        	dirInfo = os.path.split(self.image)

		        	self.imageDir = dirInfo[0]
		        	self.image    = dirInfo[1]

		self.scene = sys.argv.pop()

		if (self.scene == ""):
			print "Requires a Maya scene file to run"
			sys.exit()

		# Is the scene a valid Maya scene, if not we are unhappy
		self.validateScene()

	##############################################
	# Here is a simple function that is called if#
	# the user gives an invalid argument or      #
	# forgets something required for a command.  #
	# It then gives a basic help display on how  #
	# to use this script successfully.	     #
	##############################################
        def helpMe(self):
            print "Command options: [-m]    [-r]       [-f]      [-l]     [-b]     [-i]      [-d]       [-h]"
            print "                 [--mel] [--render] [--first] [--last] [--blur] [--image] [-display] [--help]"
            print ""
            print "USAGE: [-m] [--mel]     Use Mel command,  put inside \"\" "
            print "       [-r] [--render]  Render the scene"
            print "       [-f] [--first]   First frame (ex: 1)"
            print "       [-l] [--last]    Last frame  (ex: 10)"
            print "       [-b] [--blur]    Add motion blur to scene (ex: 0.5)"
            print "       [-i] [--image]   Image to render to (ex: \\dir\\myImage.tif)"
            print "       [-d] [--display] Execute fcheck to see your animaiton"
            print "       [-h] [--help]    You are reading it right now"
            print ""
            print "Example of use: "
            print 'Maya Batch: python Main.py -m -m \'-command \"setAttr \"lampShape.dropoff\" 350; file -f -save;\"\' -f 1 -l 5 -b 0.5 -i /home/images light2.mb'
	    print ""
            print 'Render Scene: python Main.py -m -r -f 1 -l 5 -b 0.5 -i /home/images -d light2.mb'
	    print ""
            print 'Render Scene with MEL: python Main.py -m -m \'-command \"setAttr \"lampShape.dropoff\" 350; file -f -save;\"\' -f 1 -l 5 -b 0.5 -i /home/images light2.mb'
        
	#################################################
	# Here we validate the directory the user gave  #
	# us exists, and if it does not it is then 	#
	# created.  					#
	#################################################
        def validateImageDir(self):
 		# See if the directory exists
    		if (os.path.exists(self.imageDir)):
        		# directory exists
        		print "Directory exists"
	    	else:
       		  	# Create directory if it doesn't exist
        		os.mkdir(self.imageDir)
	        	print "Directory created: " + self.imageDir

	#################################################
	# Here is where we validate that the scene is	#
	# a valide Maya scene file			#
	#################################################
   	def validateScene(self):
   		if (self.scene == None):
   	  		print "Maya scene required"
   			self.helpMe()
   	  	  	raise SystemExit

  		# Use a regular expression to check if the scene
		#    is a valid Maya scene file, being either a
   		#    *.ma (maya ASCII) or a *.mb (Maya binary)
   		valid = re.compile('\w+.mb|\w+.ma',re.IGNORECASE)
   		checkScene = valid.match(self.scene)
		try:
    			checkScene != None
  		except:
			self.scene = ""
	  		print "Invalid scene"
			self.helpMe()
	  		raise SystemExit
   		
	#################################################
	# Here is a call to execute the command line 	#
	# version of Maya.				#
	#################################################
	def batch(self):
		batchString = "maya -batch -file " + self.scene + " " + self.melCmd

	     	mayaCommand = commands.getoutput(batchString)
      		print mayaCommand
      	  
	##############################################
	# Here we do all the rendering of the Maya   #
	# scene.				     #
	##############################################
	def renderScene(self):
		if (self.blurFactor == 1):
      			renderMe = "Render -s " + self.firstFrame +  " -e " + self.lastFrame + " -im " + self.image + " -of tif -rd " + self.imageDir + " " + self.blurOptions + " "  + self.melCmd + " " + self.scene
		else:
			renderMe = "Render -s " + self.firstFrame +  " -e " + self.lastFrame + " -im " + self.image + " -of tif -rd " + self.imageDir + " " + self.melCmd + " " + self.scene

		render = commands.getoutput(renderMe)
		print render

		if (self.dispRender == 1):
			fCheck = "fcheck -S " + self.imageDir + "\/*.tif"
			fCheckOut = commands.getoutput(fCheck)
			print fCheckOut

	##############################################
	# Here is the bread and butter of this class #
	# as it calls and executes everything for the#
	# command line version of Maya. 	     #
	##############################################
	def execute(self):
		if (self.MAYARENDER == 1):
			self.renderScene()
        	else:
			self.batch()