#!/usr/bin/env fslpython
#   Copyright (C) 2016 University of Oxford
#   SHBASECOPYRIGHT

from __future__ import print_function

import errno
import os
import subprocess
import sys

fsldir = os.getenv('FSLDIR')
scriptdir = fsldir + '/python/mist'
sys.path.append(scriptdir)
import mist_script

# Input file format is CSV with double quotes, each line specifies a modality:
# Field 1: Arbitrary name for modality
# Field 2: Type specification (T1, T2, FA)
# Field 3: Relative filename
# Field 4: Voxel size (needed because images have already been resampled)

try:
    _, _, _, _, _, _ = mist_script.read_config()
except:
    print('Multimodal Image Segmentation Tool (MIST), part of FSL\n')
    print('Error: Unable to read configuration file\n')
    print('Please refer to the FSL wiki for usage instructions')

    sys.exit(1)

alldirs = mist_script.read_directories('mist_subjects')

extra = False
trainonly = False

args = sys.argv[1 :]

if len(args) > 0:
    if args[0] == '--extra':
        extra = True
        del args[0]
    elif args[0] == '--trainonly':
        trainonly = True
        del args[0]

if extra:
    with open('mist_out/mist_1a_extra_jobs', 'w') as f:
        f.writelines('fslpython ' + scriptdir + '/mist_script.py preproc ' + d + '\n' for d in alldirs if not os.path.exists(d + '/mist_t1_brain.nii.gz'))

    subprocess.check_output([fsldir + '/bin/fsl_sub', '-T', '2800', '-l', 'mist_out/log', '-t', 'mist_out/mist_1a_extra_jobs'])
else:
    if not trainonly:
        try:
            trainingdirs = mist_script.read_directories('mist_training_subjects')
        except IOError as e:
            if e.errno != errno.ENOENT:
                raise

            print('Using all subjects for training')

            trainingdirs = alldirs
            mist_script.write_directories('mist_training_subjects', trainingdirs)

        os.mkdir('mist_out')

        with open('mist_out/mist_1a_jobs', 'w') as f:
            f.writelines('fslpython ' + scriptdir + '/mist_script.py preproc ' + d + '\n' for d in alldirs)

        with open('mist_out/mist_1b_jobs', 'w') as f:
            f.write('fslpython ' + scriptdir + '/mist_script.py autosetup ' + ' '.join(args) + '\n')

    with open('mist_out/mist_1c_jobs', 'w') as f:
        workers = 40
        f.writelines('fslpython ' + scriptdir + '/mist_script.py train {0} {1}\n'.format(workers, i) for i in range(workers))

    with open('mist_out/mist_1d_jobs', 'w') as f:
        f.write('fslpython ' + scriptdir + '/mist_script.py merge \n')

    if not trainonly:
        jida = subprocess.check_output([fsldir + '/bin/fsl_sub', '-T', '2800', '-l', 'mist_out/log', '-t', 'mist_out/mist_1a_jobs'])
        jidb = subprocess.check_output([fsldir + '/bin/fsl_sub', '-T', '2800', '-l', 'mist_out/log', '-t', 'mist_out/mist_1b_jobs', '-j', str(jida.decode().split()[0])])
        jidc = subprocess.check_output([fsldir + '/bin/fsl_sub', '-T', '2800', '-l', 'mist_out/log', '-t', 'mist_out/mist_1c_jobs', '-j', str(jidb.decode().split()[0])])
    else:
        jidc = subprocess.check_output([fsldir + '/bin/fsl_sub', '-T', '2800', '-l', 'mist_out/log', '-t', 'mist_out/mist_1c_jobs'])

    subprocess.check_output([fsldir + '/bin/fsl_sub', '-T', '2800', '-l', 'mist_out/log', '-t', 'mist_out/mist_1d_jobs', '-j', str(jidc.decode().split()[0])])

print('All jobs submitted')

