#!/usr/bin/python3

"""
This program executes FSL randomise command over a given set of external
variables. You should supply the FSL template for the group of subjects,
the variables list, the full curated data and the design matrices. The 
script only  parallelize the set of randomise calls over the cluster.

"""

import os
import sys
import getopt
import csv
import re
import tempfile
import subprocess

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

#Entorno de trabajo
time = '12:0:0'
cpus = 4
wdir = os.environ.get('PWD')
outdir = wdir+'/slurm'
if not os.path.isdir(outdir): os.mkdir(outdir)

# Get CLI inputs
short_args = 'd:t:'
long_args = ['data=', 'template=']

try:
  args, values = getopt.getopt(sys.argv[1:], short_args, long_args)
except getopt.error as err:
  print(str(err))
  sys.exit(2)

for a,v in args:
  if a in ('--data', '-d'):
    datafile = v
  elif a in ('--template', '-t'):
    templatefile = v
for v in dynvars:
  ltw = ''
  with open(datafile) as f:
    reader = csv.reader(f)
    for row in reader:
      if not row[0] == 'Subject':
        for sv in stcov:
          ltw += row[sv-1]+' '
        ltw += row[v]+'\n'
  f.close()
  tempfile = 'tempdata_'+str(v)+'.dat'
  df = open(tempfile, 'w')
  df.write(ltw)
  df.close()
  os.system('Text2Vest '+df.name+' design_var'+str(v)+'.mat')
  bcontent = '#!/bin/bash\n'+'#SBATCH -J randomise\n'+'#SBATCH -c '+str(cpus)+'\n'+'#SBATCH --mem-per-cpu=4G\n'
  bcontent += '#SBATCH --time='+time+'\n'+'#SBATCH --mail-type=FAIL,TIME_LIMIT,STAGE_OUT\n'
  bcontent += '#SBATCH --mail-user='+os.environ.get('USER')+'\n'+'#SBATCH -o '+outdir+'/randomise-%j'+'\n'
  bcontent += 'randomise -i '+wdir+'/'+templatefile+' -o '+wdir+'/stats/fslvbm_var'+str(v)+' -m '+wdir+'/stats/GM_mask -d '
  bcontent += wdir+'/design_var'+str(v)+'.mat -t '+wdir+'/design.con -n 5000 -T\n'
  sbfilename = outdir+'/run_var_'+str(v)+'.sh'
  sbf = open(sbfilename, 'w')
  #bcontent = '#!/bin/bash\n:\n'
  sbf.write(bcontent)
  sbf.close()
  jobid = subprocess.getoutput('sbatch --parsable '+sbfilename)
  bcontent = '#!/bin/bash\n'+'#SBATCH -J randomise\n'+'#SBATCH -c '+str(cpus)+'\n'+'#SBATCH --mem-per-cpu=4G\n'
  bcontent += '#SBATCH --time='+time+'\n'+'#SBATCH --mail-type=FAIL,TIME_LIMIT,STAGE_OUT\n'
  bcontent += '#SBATCH --mail-user='+os.environ.get('USER')+'\n'+'#SBATCH -o '+outdir+'/cluster-%j'+'\n'
  bcontent += 'cluster --in='+wdir+'/stats/fslvbm_var'+str(v)+'_tfce_corrp_tstat1 --thresh=0.95 -o '+wdir+'/stats/clusters_var'+str(v)+' > '+wdir+'/stats/clusters_var'+str(v)+'_index.txt\n'
  sbfilename = outdir+'/cluster_var_'+str(v)+'.sh'
  sbf = open(sbfilename, 'w')
  sbf.write(bcontent)
  sbf.close()
  os.system('sbatch --dependency=afterok:'+jobid+' '+sbfilename)
bcontent = '#!/bin/bash\n'+'#SBATCH -J randomise\n'
bcontent += '#SBATCH --mail-type=END\n'
bcontent += '#SBATCH --mail-user='+os.environ.get('USER')+'\n'+'#SBATCH -o '+outdir+'/randomise-%j'+'\n'
bcontent += ':\n'
bfilename = outdir+'/run_var_'+str(v)+'.sh'
sbf = open(sbfilename, 'w')
sbf.write(bcontent)
sbf.close()
os.system('sbatch --dependency=singleton '+sbfilename)