#! /usr/bin/python3

import sys
import os
import time

#cluster = "LIS"
cluster = "MESOCENTRE"

prefix = "#SBATCH "
account = "-A b177"
shell = "#! /bin/bash"
partitionStr = "-p %s"
jobNameStr = "-J %s"
memPerCPU = "--mem-per-cpu=4GB"
nbCpusStr = "--ntasks-per-node=%d"
timeStr = "-t %s"
stdoutStr = "-o %x.stdout"
stderrStr = "-e %x.stderr"
modules = ["userspace/all", "boost/gcc72/1.65.1"]
configPrefix = "UD_ROOT=~/ud/ud-treebanks-all/"
commandStr = "time ./train.sh . %s %s --interactive 0 --printTime %s --seed %d"

def waitUntilDirExists(dirPath) :
  time.sleep(2.0)
  while not os.path.isdir(dirPath) :
    time.sleep(2.0)

def waitUntilFileExists(filePath) :
  time.sleep(2.0)
  while not os.path.isfile(filePath) :
    time.sleep(2.0)

def launchExperiment(langChanged, jobName,lang,template,seed,nbCpus=2,partition='skylake',time='7-00',moreArgs='') :

  if langChanged :
    print(configPrefix+lang+'/', file=open('config', 'w'))

    os.system("cd data && make clean && make && cd ..")

    waitUntilFileExists('data/tagger.as')

  jobName += '_' + lang

  if cluster == "MESOCENTRE" :
    slurmFile = open('script.slurm', 'w')
  
    print(shell, file=slurmFile)
    print(prefix+account, file=slurmFile)
    print(prefix+jobNameStr%jobName, file=slurmFile)
    print(prefix+partitionStr%partition, file=slurmFile)
    print(prefix+memPerCPU, file=slurmFile)
    print(prefix+nbCpusStr%nbCpus, file=slurmFile)
    print(prefix+timeStr%time, file=slurmFile)
    print(prefix+stdoutStr, file=slurmFile)
    print(prefix+stderrStr, file=slurmFile)
    print("", file=slurmFile)
    print("module purge", file=slurmFile)
    for module in modules :
      print("module load %s"%module, file=slurmFile)
  
    print("", file=slurmFile)
  
    print(commandStr%(template,jobName,moreArgs,seed), file=slurmFile)
  
    slurmFile.close()
  
    os.system("sbatch script.slurm")
  elif cluster == "LIS" :
    command = "oarsub"
    command += " -t besteffort"
    command += " -n %s"%jobName
    command += " -E %s.stderr"%jobName
    command += " -O %s.stdout"%jobName
    command += " -p \"gpu IS NULL\""
    command += " -l walltime=96:00:00"
    command += " \"" + commandStr%(template,jobName,moreArgs,seed) + "\""
    os.system(command)
  else :
    print("ERROR : not a valid cluster \'%s\'"%cluster)
    exit(1)

  waitUntilDirExists('bin/'+jobName+'/data/')


if __name__ == "__main__" :
  templatesExperiments = \
  [\
#    {'jobName' : 'tagparser','template' : 'tagparser','moreArgs' : '-n 10'},\
    #{'jobName' : 'tokeparser','template' : 'tokeparser','moreArgs' : '-n 15 --rawInput'},\
    {'jobName' : 'tokeparser_incremental_b0','template' : 'tokeparser_incremental','moreArgs' : '-n 15 --rawInput'},\
    {'jobName' : 'tokeparser_incremental_b1','template' : 'tokeparser_incremental','moreArgs' : '-n 15 --rawInput'},\
    {'jobName' : 'tokeparser_incremental_b2','template' : 'tokeparser_incremental','moreArgs' : '-n 15 --rawInput'},\
    {'jobName' : 'tokeparser_incremental_b3','template' : 'tokeparser_incremental','moreArgs' : '-n 15 --rawInput'},\
    {'jobName' : 'tokeparser_incremental_b4','template' : 'tokeparser_incremental','moreArgs' : '-n 15 --rawInput'},\
    {'jobName' : 'tokeparser_incremental_b5','template' : 'tokeparser_incremental','moreArgs' : '-n 15 --rawInput'},\
    {'jobName' : 'tokeparser_incremental_b6','template' : 'tokeparser_incremental','moreArgs' : '-n 15 --rawInput'},\
    #{'jobName' : 'tokeparser_sequential','template' : 'tokeparser_sequential','moreArgs' : '-n 15 --rawInput'},\
  ]\

  #langs = ["UD_French-GSD", "UD_Hebrew-HTB", "UD_Chinese-GSD", "UD_English-EWT", "UD_French-Spoken"]
  #langs = ["UD_French-GSD", "UD_Hebrew-HTB", "UD_Chinese-GSD", "UD_English-EWT", "UD_French-Spoken", "UD_Russian-SynTagRus", "UD_Arabic-PADT", "UD_Finnish-TDT", "UD_Turkish-IMST", "UD_Norwegian-Bokmaal", "UD_Romanian-RRT"]
  langs = ["UD_English-EWT"]

  nbReplicas = 3

  for lang in langs :
    langChanged = True
    for experience in templatesExperiments :
      for i in range(nbReplicas) :
        experience['lang'] = lang
        experience['langChanged'] = langChanged
        experience['seed'] = 100+i
        experience['jobName'] = experience['jobName'].split('.')[0]+"."+str(i)
        launchExperiment(**experience)
        langChanged = False