#!${pyExecutable}
# encoding: utf-8
#
# vim: syntax=python
# portions © 2008 Václav Šmilauer <eudoxos@arcig.cz>
#  © 2017 William Chèvremont <william.chevremont@univ-grenoble-alpes.fr>
# This script is to be used with OAR task scheduler. May be an example to use use with other task scheduler for clusters
# Adapted from yade-batch

import os, sys, _thread, time, logging, shlex, socket, xmlrpc.client, re, shutil, random

# Add search path for yade Python-modules
# It allows to use both Yade-versions (packaged and self-compiled one).
# See LP:1254708 for more details
# (old site, fixed bug) https://bugs.launchpad.net/yade/+bug/1254708


#socket.setdefaulttimeout(10)

# Setup executable

## replaced by scons automatically
prefix,suffix='${runtimePREFIX}' if 'YADE_PREFIX' not in os.environ else os.environ['YADE_PREFIX'],'${SUFFIX}'

libPATH='${LIBRARY_OUTPUT_PATH}'
if (libPATH[1:] == '{LIBRARY_OUTPUT_PATH}'): libPATH='lib'

libDir=prefix+'/'+libPATH+'/yade'+suffix # run the batch always in non-debug mode (the spawned processes do honor debuggin flag, however)

sys.path.append(os.path.join(libDir,'py'))
executable=os.path.join(prefix,'bin','yade'+suffix)

import yade, yade.utils, yade.config, yade.remote
from shutil import copyfile


def t2hhmmss(dt): return '%02d:%02d:%02d'%(dt//3600,(dt%3600)//60,(dt%60))


import sys,re,os
try:
	import argparse
except ImportError: # argparse not present, print error message
	raise RuntimeError("\n\nPlease install 'python-argparse' package.\n")


class JobInfo(object):
	def __init__(self,num,id,command,log,nCores,script,table,lineNo,affinity,oarprop):
		self.started,self.finished,self.duration,self.durationSec,self.exitStatus=None,None,None,None,None # duration is a string, durationSec is a number
		self.command=command; self.num=num; self.log=log; self.id=id; self.nCores=nCores; self.cores=set(); self.infoSocket=None
		self.script=script; self.table=table; self.lineNo=lineNo; self.affinity=affinity
		self.hasXmlrpc=False
		self.status='PENDING'
		self.threadNum=None
		self.oarprop=oarprop
		self.plotsLastUpdate,self.plotsFile=0.,yade.Omega().tmpFilename()+'.'+yade.remote.plotImgFormat


features,version='${CONFIGURED_FEATS}'.split(','),'${realVersion}'
if (features[0]==''): features=features[1:]

prog = os.path.basename(sys.argv[0])
parser=argparse.ArgumentParser(usage='%s --oar-script=[SCRIPT] [options] [ TABLE [SIMULATION.py] | SIMULATION.py[/nCores] [...] ]'%prog,description='%s runs yade simulation multiple times with different parameters.\n\nSee https://yade-dem.org/sphinx/user.html#batch-queuing-and-execution-yade-batch for details.\n\nBatch can be specified either with parameter table TABLE (must not end in .py), which is either followed by exactly one SIMULATION.py (must end in .py), or contains !SCRIPT column specifying the simulation to be run. The second option is to specify multiple scripts, which can optionally have /nCores suffix to specify number of cores for that particular simulation (corresponds to !THREADS column in the parameter table), e.g. sim.py/3.'%prog)
parser.add_argument('-v','--version',help='Print version and exit.',dest='version',action='store_true')
parser.add_argument('--job-threads',dest='defaultThreads',type=int,help="Default number of threads for one job; can be overridden by per-job with !THREADS (or !OMP_NUM_THREADS) column. Defaults to 1.",metavar='NUM',default=1)
parser.add_argument('--force-threads',action='store_true',dest='forceThreads',help=r'Force jobs to not use more cores than the maximum (see \-j), even if !THREADS colums specifies more.')
parser.add_argument('--log',dest='logFormat',help='Format of job log files: must contain a $, %% or @, which will be replaced by script name, line number or by description column respectively (default: $.%%.log)',metavar='FORMAT',default='$.%.log')
parser.add_argument('--global-log',dest='globalLog',help=r'Filename where to redirect output of yade-batch itself (as opposed to \-\-log); if not specified (default), stdout/stderr are used',metavar='FILE')
parser.add_argument('-l','--lines',dest='lineList',help='Lines of TABLE to use, in the format 2,3-5,8,11-13 (default: all available lines in TABLE)',metavar='LIST')
parser.add_argument('--nice',dest='nice',type=int,help='Nice value of spawned jobs (default: 10)',default=10)
parser.add_argument('--cpu-affinity',dest='affinity',action='store_true',help='Bind each job to specific CPU cores; cores are assigned in a quasi-random order, depending on availability at the moment the jobs is started. Each job can override this setting by setting AFFINE column.')
parser.add_argument('--executable',dest='executable',help='Name of the program to run (default: %s). Jobs can override with !EXEC column.'%executable,default=executable,metavar='FILE')
parser.add_argument('--dry-run',action='store_true',dest='dryRun',help='Do not actually run (useful for getting gnuplot only, for instance)',default=False)
parser.add_argument('--timing',type=int,dest='timing',default=0,metavar='COUNT',help='Repeat each job COUNT times, and output a simple table with average/variance/minimum/maximum job duration; used for measuring how various parameters affect execution time. Jobs can override the global value with the !COUNT column.')
parser.add_argument('--timing-output',type=str,metavar='FILE',dest='timingOut',default=None,help=r'With \-\-timing, save measured durations to FILE, instead of writing to standard output.')
parser.add_argument('--randomize',action='store_true',dest='randomize',help='Randomize job order (within constraints given by assigned cores).')
parser.add_argument('--disable-pynotify',action='store_true',dest='disablePynotify',help='Disable screen notifications')
parser.add_argument('--oar-project',dest='oar_p',help='Project name to pass to oarsub')
parser.add_argument('--oar-walltime',dest='oar_t',help='Walltime: max running time. May be overriden by !WALLTIME',default='2:00:00')
parser.add_argument('--oar-script',dest='oar_script',help='Script passed to oar-sub. Must contain __YADE_COMMAND__ string where the Yade command will be replaced.',default='')
parser.add_argument('--oar-besteffort',dest="oar_best",action='store_true',default=False,help="run jobs as besteffort");
parser.add_argument('--oar-properties',dest="oar_prop",help="Properties to pass to oarsub. May be overriden by !OARPROPERTIES",default='');
parser.add_argument('args',nargs=argparse.REMAINDER,help=argparse.SUPPRESS) # see argparse doc, par.disable_interspersed_args() from optargs module
#parser.add_argument('--serial',action='store_true',dest='serial',default=False,help='Run all jobs serially, even if there are free cores
opts=parser.parse_args()
args = opts.args

logFormat,lineList,nice,executable,dryRun,globalLog,project, walltime=opts.logFormat,opts.lineList,opts.nice,opts.executable,opts.dryRun,opts.globalLog,opts.oar_p,opts.oar_t

if opts.version:
	print('Yade version: %s, features: %s'%(version,','.join(features)))
	sys.exit(0)

if globalLog:
	print('Redirecting all output to',globalLog)
	sys.stderr=open(globalLog,"w")
	sys.stdout=sys.stderr

if len([1 for a in args if re.match(r'.*\.py(/[0-9]+)?',a)])==len(args) and len(args)!=0:
	# if all args end in .py, they are simulations that we will run
	table=None; scripts=args
elif len(args)==2:
	table,scripts=args[0],[args[1]]
elif len(args)==1:
	table,scripts=args[0],[]
else:
	parser.print_help()
	sys.exit(1)

if opts.oar_script == '' :
        print('--oar_script is mandatory');
        sys.exit(1);

pynotifyAvailable = False
n = 0;
if not opts.disablePynotify:
	try:
		import pynotify
		pynotifyAvailable = True
		pynotify.init("Yade")
	except ImportError:
		pynotifyAvailable = False

print("Will run simulation(s) %s using `%s', nice value %d"%(scripts,executable,nice))

if table:
	reader=yade.utils.TableParamReader(table)
	params=reader.paramDict()
	availableLines=list(params.keys())

	print("Will use table `%s', with available lines"%(table),', '.join([str(i) for i in availableLines])+'.')

	if lineList:
		useLines=[]
		def numRange2List(s):
			ret=[]
			for l in s.split(','):
				if "-" in l: ret+=list(range(*[int(s) for s in l.split('-')])); ret+=[ret[-1]+1]
				else: ret+=[int(l)]
			return ret
		useLines0=numRange2List(lineList)
		for l in useLines0:
			if l not in availableLines: logging.warn('Skipping unavailable line %d that was requested from the command line.'%l)
			else: useLines+=[l]
	else: useLines=availableLines
	print("Will use lines ",',\n '.join([str(i)+' (%s)'%params[i]['description'] for i in useLines])+'.')
else:
	print("Running %d stand-alone simulation(s) in batch mode."%(len(scripts)))
	useLines=[]
	params={}
	for i,s in enumerate(scripts):
		fakeLineNo=-i-1
		useLines.append(fakeLineNo)
		params[fakeLineNo]={'description':'default','!SCRIPT':s}
		# fix script and set threads if script.py/num
		m=re.match('(.*)(/[0-9]+)$',s)
		if m:
			params[fakeLineNo]['!SCRIPT']=m.group(1)
			params[fakeLineNo]['!THREADS']=int(m.group(2)[1:])
jobs=[]
executables=set()
logFiles=[]
for i,l in enumerate(useLines):
	script=scripts[0] if len(scripts)>0 else None
	envVars=[]
	nCores=opts.defaultThreads
	jobExecutable=executable
	jobAffinity=opts.affinity
	jobCount=opts.timing
	walltime = opts.oar_t
	oarprops = opts.oar_prop
	for col in list(params[l].keys()):
		if col[0]!='!': continue
		val=params[l][col]
		if col=='!OMP_NUM_THREADS' or col=='!THREADS': nCores=int(val)
		elif col=='!EXEC': jobExecutable=val
		elif col=='!SCRIPT': script=val
		elif col=='!AFFINITY': jobAffinity=eval(val)
		elif col=='!COUNT': jobCount=eval(val)
		elif col=='!WALLTIME' : walltime=val
		elif col=='!OARPROPERTIES' : oarprops=val
		else: envVars+=['%s=%s'%(head[1:],values[l][col])]
	if not script:
		raise ValueError('When only batch table is given without script to run, it must contain !SCRIPT column with simulation to be run.')
	logFile=logFormat.replace('$',script.split('/')[-1]).replace('%',str(l)).replace('@',params[l]['description'])

	if (len(logFile)>230):                              #For very long files, to prevent system errors.
		logFile = logFile[0:230]+'_'+str(l)+logFile[-4:]  #l is added to prevent name duplications.

	executables.add(jobExecutable)

	for j in range(0,jobCount if jobCount>0 else 1):
		jobNum=len(jobs)
		logFile2=logFile+('.%d'%j if opts.timing>0 else '')
		# append numbers to log file if it already exists, to prevent overwriting
		if logFile2 in logFiles:
			i=0;
			while logFile2+'.%d'%i in logFiles: i+=1
			logFile2+='.%d'%i
		logFiles.append(logFile2)
		env='YADE_WALLTIME=%s YADE_BATCH='%walltime
		if table: env+='%s:%d'%(os.path.abspath(table),l) # keep YADE_BATCH empty (but still defined) if running a single simulation
		cmd=' %s%s -j %i -x %s 2>&1'%(jobExecutable,' --nice=%s'%nice if nice!=None else '',nCores,os.path.abspath(script))
		desc=params[l]['description']
		if '!SCRIPT' in list(params[l].keys()): desc=script+'.'+desc # prepend filename if script is specified explicitly
		if opts.timing>0: desc+='[%d]'%j
		jobs.append(JobInfo(jobNum,desc,env+cmd,os.path.abspath('tmp/'+str(os.getpid())+'/'+logFile2),nCores,script=os.path.abspath(script),table=table,lineNo=l,affinity=jobAffinity,oarprop=oarprops))

if not os.path.exists('tmp/'+str(os.getpid())):
    os.makedirs('tmp/'+str(os.getpid()));

print("Job summary:")
i = 0
for job in jobs:
	print('   #%d (%s%s):'%(job.num,job.id,'' if job.nCores==1 else '/%d'%job.nCores),job.command)
	if not dryRun:
		# Create script file from model
		scriptname = 'tmp/'+str(os.getpid())+'/'+opts.oar_script.split('/')[-1]+'.'+str(i)
		i = i+1;

		with open(opts.oar_script) as infile, open(scriptname, 'w') as outfile:
			for line in infile:
				line = line.replace("__YADE_COMMAND__", job.command)
				line = line.replace("__YADE_LOGFILE__", job.log)
				line = line.replace("__YADE_ERRFILE__", job.log+'.err')
				line = line.replace("__YADE_JOBNO__", str(os.getpid())+'-'+str(job.num))
				line = line.replace("__YADE_JOBID__", str(job.id))
				outfile.write(line)

		prop='';

		if job.oarprop != '':
			prop='-p '+job.oarprop;

		if opts.oar_best:
			prop+=' -t besteffort'


		os.chmod(scriptname,0o744);
		#os.system('%s'%(scriptname));
		os.system('oarsub --project=%s %s -l/cpu=1/core=%i,walltime=%s -O %s -E %s.err %s\n'%(project, prop, job.nCores, walltime, job.log,job.log,scriptname));



sys.stdout.flush()

yade.Omega().exitNoBacktrace()
