#! /bin/env python
#
# Protein Engineering Analysis Tool Structure Analysis (PEATSA)
# Copyright (C) 2010 Michael Johnston & Jens Erik Nielsen
#
# Author: Michael Johnston
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Contact information:
# Email: Jens.Nielsen_at_gmail.com
# Normal mail:
# Jens Nielsen
# SBBS, Conway Institute
# University College Dublin
# Dublin 4, Ireland
#

'''Wrapper script which creates and submits a parallel ProteinDesignTool job

The script creates PBS scripts and submits them using qsub'''

import Environment, sys, os, ConfigParser, Utilities

def PBSScriptFromConfiguration(configuration, arguments, runDirectory=None, logDirectory=None):

	'''Creates a launch script using the information specified by the configuration object
	
	The configuration object specifies the following parameters
		
		numberOfNodes
		processorsPerNode
		runDir
		logDir
		executable
		queue
		keep
		
	However runDir and logDir can be overridden	
		
	The arguments to the executable must be passed to this function'''	

	numberOfNodes = configuration.get('PARALLEL', 'numberOfNodes')
	processorsPerNode = configuration.get('PARALLEL', 'processorsPerNode')
	executable = configuration.get('PARALLEL', 'executable')
	
	if runDirectory is None:
		runDirectory = configuration.get('PARALLEL', 'runDirectory')
	if logDirectory is None:	
		logDirectory = configuration.get('PARALLEL', 'logDirectory')
	
	try:
		queue = configuration.get('PARALLEL', 'queue')
	except ConfigParser.NoOptionError:
		queue = "short"
		
	try:
		keep = configuration.get('PARALLEL', 'keep')
	except ConfigParser.NoOptionError:
		keep = False	

	pbsScript = CreatePBSLaunchScript(executable=executable,
			arguments=arguments,
			numberOfNodes=numberOfNodes,
			processorsPerNode=processorsPerNode,
			queue=queue,
			runDirectory=runDirectory,
			logDirectory=logDirectory,
			keep=keep)
	
	return pbsScript

def CreatePBSLaunchScript(executable, arguments, numberOfNodes=4, processorsPerNode=2, queue="short", runDirectory=os.getcwd(), logDirectory=os.getcwd(), keep=True):
	
	'''Creates a PBS script for a parallel job submission'''
	
	argString = " ".join(arguments)
	runString = "%s %s" % (executable, argString)
	
	pbsString = "## Autogenerated script for submitting PEAT_SA to the PBS queue\n\n" 
	pbsString = pbsString + "#PBS -l nodes=%s:ppn=%s\n" % (numberOfNodes, processorsPerNode)
	pbsString = pbsString + "#PBS -j oe\n"
	pbsString = pbsString + "#PBS -q %s\n" % queue
	pbsString = pbsString + "#PBS -d %s\n" % os.path.abspath(runDirectory)
	pbsString = pbsString + "#PBS -o %s\n" % os.path.abspath(logDirectory)
	if keep is True:
		pbsString = pbsString + "#PBS -k oe\n"
	
	pbsString = pbsString + "\nmpirun -np %d python %s\n" % (int(numberOfNodes)*int(processorsPerNode), runString)

	return pbsString

if __name__ == "__main__":

	parser = Utilities.CommandLineParser()
	parser.parseCommandLine()

	configuration = Environment.Configuration(filename=parser.configurationFile())
	pbsScript = PBSScriptFromConfiguration(configuration, sys.argv[1:])
			
	stream = open('PEATSA.pbs', 'w+')
	stream.write(pbsScript)
	stream.close()

	#Run process		
	os.system("qsub PEATSA.pbs")
	print 'Submitted job-script PEATSA.pbs'
	
