#! /usr/bin/env python3

import sys
import os
import subprocess
import time

###############################################################################
def printUsageAndExit() :
  print("USAGE : %s (train | eval) (bash | oar | slurm) batchesDescription.py (--time nbHours)"%sys.argv[0], file=sys.stderr)
  exit(1)
###############################################################################

###############################################################################
def prepareExperiment(lang, template, expName) :
  subprocess.Popen("./prepareExperiment.sh %s %s %s"%(lang,template,expName),
      shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE).wait()
###############################################################################

###############################################################################
def launchTrain(mode, expName, arguments, launcher, nbHours, seed) :
  if launcher == "bash" :
    launchTrainBash(mode, expName, arguments, seed)
  elif launcher == "oar" :
    launchTrainOar(mode, expName, arguments, nbHours, seed)
  elif launcher == "slurm" :
    launchTrainSlurm(mode, expName, arguments, nbHours, seed)
  else :
    printUsageAndExit()
###############################################################################

###############################################################################
def launchTrainBash(mode, expName, arguments, seed) :
  subprocess.Popen("./train.sh %s bin/%s %s --silent --seed %d"%(mode,expName,arguments,seed),
    shell=True, stdout=open("%s.stdout"%expName,'w'), stderr=open("%s.stderr"%expName,'w'))
###############################################################################

###############################################################################
def nbMaxLongJobs() :
  return 2
###############################################################################

###############################################################################
def launchTrainOar(mode, expName, arguments, nbHours, seed) :
  bestEffort = getOarNbLongJobs() >= nbMaxLongJobs()

  command = "oarsub"
  command += " -t besteffort" if bestEffort else ""
  command += " -t idempotent" if bestEffort else ""
  command += " -n train:%s"%expName
  command += " -E %s.stderr"%expName
  command += " -O %s.stdout"%expName
  command += " -p \"gpu IS NOT NULL%s\""%getBestHostConstraint()
  command += " -l walltime=%d:00:00"%nbHours
  command += " \'" + "./train.sh %s bin/%s %s --silent --seed"%(mode,expName,arguments,seed) + "\'"

  subprocess.Popen(command, shell=True).wait()
###############################################################################

###############################################################################
def launchTrainSlurm(mode, expName, arguments, nbHours, seed) :
  filename = "train.{}.slurm".format(expName)
  sFile = open(filename, "w")

  print("""#! /usr/bin/env bash

#SBATCH --job-name=train:{}
#SBATCH --output={}.stdout
#SBATCH --error={}.stderr
#SBATCH --open-mode=append
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=10
#SBATCH --gres=gpu:1
#SBATCH --hint=nomultithread
#SBATCH --partition=gpu_p1
#SBATCH --qos={}
#SBATCH --time={}:00:00

module purge
module load gcc/9.1.0
module load python/3.7.5

./train.sh {} bin/{} {} --silent --seed {}
""".format(expName, expName, expName, "qos_gpu-t4" if nbHours > 20 else "qos_gpu-t3", nbHours, mode, expName, arguments, seed), file=sFile)
  sFile.close()

  subprocess.Popen("sbatch {}".format(filename), shell=True).wait()
###############################################################################

###############################################################################
def launchEval(mode, expName, launcher, nbHours) :
  if launcher == "bash" :
    launchEvalBash(mode, expName)
  elif launcher == "oar" :
    launchEvalOar(mode, expName, nbHours)
  elif launcher == "slurm" :
    launchEvalSlurm(mode, expName, nbHours)
  else :
    printUsageAndExit()
###############################################################################

###############################################################################
def launchEvalBash(mode, expName) :
  subprocess.Popen("./evaluate.sh %s bin/%s --silent"%(mode,expName),
    shell=True, stdout=open("%s.stdout"%expName,'a'), stderr=open("%s.stderr"%expName,'a'))
###############################################################################

###############################################################################
def launchEvalOar(mode, expName, nbHours) :
  bestEffort = getOarNbLongJobs() >= nbMaxLongJobs() and nbHours > 10

  command = "oarsub"
  command += " -t besteffort" if bestEffort else ""
  command += " -t idempotent" if bestEffort else ""
  command += " -n eval:%s"%expName
  command += " -E %s.stderr"%expName
  command += " -O %s.stdout"%expName
  command += " -p \"gpu IS NOT NULL%s\""%getBestHostConstraint()
  command += " -l walltime=%d:00:00"%nbHours
  command += " \"" + "./evaluate.sh %s bin/%s --silent"%(mode,expName) + "\""

  subprocess.Popen(command, shell=True).wait()
###############################################################################

###############################################################################
def launchEvalSlurm(mode, expName, nbHours) :
  filename = "eval.{}.slurm".format(expName)
  sFile = open(filename, "w")

  print("""#! /usr/bin/env bash

#SBATCH --job-name=eval:{}
#SBATCH --output={}.stdout
#SBATCH --error={}.stderr
#SBATCH --open-mode=append
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=10
#SBATCH --gres=gpu:1
#SBATCH --hint=nomultithread
#SBATCH --partition=gpu_p1
#SBATCH --time={}:00:00

module purge
module load gcc/9.1.0
module load python/3.7.5

./evaluate.sh {} bin/{} --silent
""".format(expName, expName, expName, nbHours, mode, expName), file=sFile)
  sFile.close()

  subprocess.Popen("sbatch {}".format(filename), shell=True).wait()
###############################################################################

###############################################################################
def getOarNbLongJobs() :
  return int(subprocess.Popen('oarstat -u | grep "Q=long" | wc -l',
    shell=True, stdout=subprocess.PIPE).stdout.read())
###############################################################################

###############################################################################
def getOarNbGpuPerNode() :
  l = subprocess.Popen("oarnodes | grep gpunum=. | grep -o 'host=[^,]*' | cut -f2 -d= | sort | uniq -c", shell=True, stdout=subprocess.PIPE).stdout.read().decode('utf8').split('\n')

  res = {}
  for line in l :
    splited = line.split()
    if len(splited) != 2 :
      continue
    res[splited[1]] = int(splited[0])

  return res
###############################################################################

###############################################################################
def getOarNbUsedGpuPerNode() :
  l = subprocess.Popen("oarstat -f | grep 'assigned_hostnames =\|propert\|wanted_resources' | grep -i 'gpu is not null' -B 2 | grep [^-]", shell=True, stdout=subprocess.PIPE).stdout.read().decode("utf8").split('\n')

  res = {}

  for i in range(len(l)//3) :
    ressources = l[3*i]
    hostname = l[3*i+1].split()[-1]
    cores = 1
    gpunum = 1
    if "core=" in ressources :
      coresStr=""
      coresStrBase = ressources.split("core=")[-1]
      for symbol in coresStrBase :
        if symbol < '0' or symbol > '9' :
          break
        coresStr = coresStr + symbol
      cores = int(coresStr)
    if "gpunum=" in ressources :
      gpunum = int(ressources.split("gpunum=")[-1].split(',')[0])

    if hostname not in res :
      res[hostname] = 0
    res[hostname] += cores * gpunum

  return res
###############################################################################

###############################################################################
def getOarNotAliveNodes() :
  res =  subprocess.Popen("oarnodes | grep -B 2 'state : [^A]' | grep 'network_address' | sort --unique | awk '{print $3}'", shell=True, stdout=subprocess.PIPE).stdout.read().decode("utf8").split('\n')
  return [node for node in res if len(node) > 0]
###############################################################################

###############################################################################
def getOarNbFreeGpuPerNode() :
  gpus = getOarNbGpuPerNode()
  notAlive = getOarNotAliveNodes()
  usedGpus = getOarNbUsedGpuPerNode()

  for gpu in gpus :
    gpus[gpu] -= usedGpus[gpu] if gpu in usedGpus else 0

  for host in notAlive :
    gpus[host] = 0

  return gpus
###############################################################################

###############################################################################
def getBestHostConstraint() :
  freeGpus = getOarNbFreeGpuPerNode()

  if freeGpus["diflives1"] > 0 or freeGpus["lisnode2"] > 0 or freeGpus["lisnode3"] > 0 :
    return " and host!='lifnode1' and host!='adnvideo1' and host!='asfalda1' and host!='see4c1' and host!='sensei1'"
  return ""
###############################################################################

###############################################################################
if __name__ == "__main__" :

  if len(sys.argv) < 4 :
    printUsageAndExit()

  mode = sys.argv[1]
  launcher = sys.argv[2]
  batchesDescription = sys.argv[3]
  nbHours = 92

  if len(sys.argv) > 4 :
    if sys.argv[4] == "--time" :
      if 5 not in range(4,len(sys.argv)) :
        printUsageAndExit()
      nbHours = int(sys.argv[5])
    else :
      printUsageAndExit()

  if mode not in ["train","eval"] or launcher not in  ["bash","oar","slurm"] :
    printUsageAndExit()

  desc = __import__(os.path.splitext(batchesDescription)[0])

  for lang in desc.langs :
    for xp in desc.templatesExperiments :
      for i in range(desc.nbReplicas) :
        xp['lang'] = lang
        xp['expName'] = xp['expName'].split('.')[0]+"."+lang+"."+str(i)
        if mode == "train" :
          prepareExperiment(xp['lang'],xp['template'],xp['expName'])
          launchTrain(xp['mode'],xp['expName'],xp['arguments'],launcher,nbHours,seed=100+i)
        else :
          launchEval(xp['mode'],xp['expName'],launcher,nbHours)

###############################################################################