#!/usr/bin/env python

import os
import subprocess
import random

def send_sbatch(env_data):
  """
  This function creates and executes an sbatch script into SLURM
 
  It takes a dict with the required environment data and returns
  the jobid of the task
  """
  content = '#!/bin/bash\n'
  if 'job_name' in env_data:
    content += '#SBATCH -J '+env_data['job_name']+'\n'
  else:
    content += '#SBATCH -J myjob\n'
  if 'cpus' in env_data:
    content += '#SBATCH -c '+str(env_data['cpus'])+'\n'
    if 'mem_cpu' in env_data:
      content += '#SBATCH --mem-per-cpu='+env_data['mem_cpu']+'\n'
    else:
      content += '#SBATCH --mem-per-cpu=4G\n'
  if 'time' in env_data: content += '#SBATCH --time='+env_data['time']+'\n'
  content += '#SBATCH --mail-user='+os.environ.get('USER')+'\n'
  if 'output' in env_data: content += '#SBATCH -o '+env_data['output']+'-%j\n'
  if 'partition' in env_data: content += '#SBATCH -p '+partition+'\n'
  if 'command' in env_data:
    content += '#SBATCH --mail-type=FAIL,TIME_LIMIT,STAGE_OUT\n'
    content += env_data['command']+'\n'
  else:
    content += '#SBATCH --mail-type=END\n'
    content += ':\n'
  if 'filename' in env_data:
    filename = env_data['filename']
  else:
    filename = 'slurm{:03d}.sh'.format(random.randint(0,1000))
  f = open(filename, 'w')
  f.write(content)
  f.close()
  if 'dependency' in env_data:
    order = ['sbatch --parsable --dependency='+env_data['dependency']+' '+filename]
  else:
    order = ['sbatch --parsable '+filename]
  return int(subprocess.check_output(order, shell=True))

# Static covariables
stcov = [6, 7, 8]
# Variables a procesar
dynvars = list(range(9,16))

#Entorno de trabajo
time = '48:0:0'
cpus = 4
wdir = os.environ.get('PWD')

for v in dynvars:
  cdata = {}
  vwdir = wdir+'/var_{:02d}'.format(v)
  if not os.path.isdir(vwdir+'/slurm'): os.mkdir(vwdir+'/slurm')
  cdata['time'] = time
  cdata['cpus'] = cpus
  cdata['job_name'] = 'randomise'
  cdata['filename'] = vwdir+'/randomise_{:02d}.sh'.format(v)
  cdata['output'] = vwdir+'/slurm/randomise_{:02d}-%j.log\n'.format(v)
  cdata['command'] = 'cd '+vwdir+'\n'+'fslmaths '+vwdir+'/stats/GM_merg -s 4 '+vwdir+'/stats/GM_merg_s4'+'randomise -i '+vwdir+'/stats/GM_merg_s4.nii.gz -o '+vwdir+'/stats/fslvbm_{:02d} -m '.format(v)+vwdir+'/stats/GM_mask -d '+vwdir+'/design.mat -t '+vwdir+'/design.con -n 5000 -T -V\n'
  send_batch(cdata)
mdata = {}
mdata['job_name'] = 'randomise'
mdata['filename'] = wdir+'/randomise_end.sh'
mdata['output'] = wdir+'/randomise_end.log'
mdata['dependency'] = 'singleton'
send_batch(mdata)