#!/usr/bin/env fslpython

import sys,os,glob,subprocess,shutil
import argparse, textwrap

FSLDIR = os.getenv('FSLDIR')
FSLbin = os.path.join(FSLDIR, 'bin')
FSLINFMRIB = os.getenv('FSLINFMRIB')
ptxbin_gpu = os.path.join(FSLbin, 'probtrackx2_gpu') # Location of probtrackx2_gpu binary
ptxbin = os.path.join(FSLbin, 'probtrackx2') # Location of probtrackx2 binary
datadir = os.path.join(FSLDIR, 'data', 'xtract_data') # Location of xtract data

# some useful functions
def errchk(errflag):
    if errflag:
        print("Exit without doing anything..")
        quit()

def imgtest(fname):
    r = subprocess.run([f'{os.path.join(FSLbin, "imtest")} {fname}'], capture_output=True, text=True, shell=True)
    return int(r.stdout)

def applyisoxfm(fin, out, res, refimg, interp):
    subprocess.run([os.path.join(FSLbin, 'flirt'), '-in', fin, '-out', out, '-applyisoxfm', str(res), '-ref', refimg, '-interp', interp])

def remove_ext(p):
    exts = ['.nii.gz', '.nii', '.surf.gii', '.gii', '.asc']
    for e in exts:
        p = p.replace(e, '')
    return p

def conv_ascii(fin, roi, dirname):
    fout = os.path.basename(fin)
    fout = remove_ext(os.path.basename(fout))
    fout += '.asc'
    fout = os.path.join(dirname, fout)
    subprocess.run([os.path.join(FSLbin, "surf2surf"), '-i', fin, '-o', fout, f'--values={roi}', '--outputtype=ASCII'])
    return fout

class MyParser(argparse.ArgumentParser):
    def error(self, message):
        sys.stderr.write('error: %s\n' % message)
        self.print_help()
        sys.exit(2)

class Usage(Exception):
    def __init__(self, msg):
        self.msg = msg


splash = r"""
__  _______ ____      _    ____ _____ _     _                       _       _
\ \/ /_   _|  _ \    / \  / ___|_   _| |__ | |_   _  ___ _ __  _ __(_)_ __ | |_
 \  /  | | | |_) |  / _ \| |     | | | '_ \| | | | |/ _ \ '_ \| '__| | '_ \| __|
 /  \  | | |  _ <  / ___ \ |___  | | | |_) | | |_| |  __/ |_) | |  | | | | | |_
/_/\_\ |_| |_| \_\/_/   \_\____| |_| |_.__/|_|\__,_|\___| .__/|_|  |_|_| |_|\__|
                                                        |_|
 """
print(splash)

parser = MyParser(prog='XTRACT Blueprint',
                  description='xtract_blueprint: for generating connectivity blueprints from xtract output',
                  formatter_class=argparse.RawDescriptionHelpFormatter,
                  epilog=textwrap.dedent('''Example usage:
                        xtract_blueprint -bpx /data/Diffusion.bedpostX -out /data/blueprint -xtract /data/xtract
                                         -seeds l.white.surf.gii,r.white.surf.gii -target WM_segmentation.nii.gz
                                         -rois l.medwall.shape.gii,r.medwall.shape.gii
                                         -warps ref.nii.gz xtract2diff.nii.gz diff2xtract.nii.gz -gpu
                  '''))

required = parser.add_argument_group('Required arguments')
optional = parser.add_argument_group('Optional arguments')

# Compulsory arguments:
required.add_argument("-bpx", metavar='<folder>', help="Path to bedpostx folder", required=True)
required.add_argument("-xtract", metavar='<folder>', help="Path to xtract folder", required=True)
required.add_argument("-seeds", metavar='<list>', help="Comma separated list of seeds for which a blueprint is requested (e.g. left and right cortex in standard space)", required=True)
required.add_argument("-target", metavar='<mask>', help="A whole brain/WM binary target mask in the same space as the seeds", required=True)

# Either standard space or native space?
optional.add_argument("-warps", metavar='<path>', nargs=3, help="<ref> <xtract2diff> <diff2xtract> Standard space reference image, and transforms between xtract space and diffusion space")
optional.add_argument("-native", action="store_true", default=False, help="Run tractography in native (diffusion) space")

# Optional arguments:
optional.add_argument("-out", metavar='<folder>', help="Path to output folder (default is <xtract_dir>/blueprint)")
optional.add_argument("-stage", choices=[0, 1, 2], default=0, type=int, help="What to run. 0:everythng (default), 1:matrix2, 2:blueprint")
optional.add_argument("-savetxt", action="store_true", default=False, help="Save blueprint as txt file (nseed by ntracts) instead of CIFTI")
optional.add_argument("-prefix", metavar='<string>', help="Specify a prefix for the final blueprint filename (e.g. <prefix>_BP.LR.dscalar.nii)")

# Optional tracking arguments:
optional.add_argument("-rois", metavar='<list>', help="Comma separated list (1 per seed): ROIs (gifti) to restrict seeding (e.g. medial wall masks)")
optional.add_argument("-stops", metavar='<stop.txt>', help="Text file containing line separated list of stop masks (see probtrackx usage for details)")
optional.add_argument("-wtstops", metavar='<wtstop.txt>', help="Text file containing line separated list of wtstop masks (see probtrackx usage for details)")
optional.add_argument("-exclusion", metavar='<nifti>', help="NIFTI containing a binary exclusion mask")
optional.add_argument("-subseed", metavar='<subseed.txt>', help="Text file containing line separated list of subcortical seeds")
optional.add_argument("-tract_list", metavar='<list>', help="Comma separated list of tracts to include (default = all found under -xtract <folder>)")
optional.add_argument("-thr", metavar='<float>', default=0.001, type=float, help="Threshold applied to XTRACT tracts prior to blueprint calculation (default = 0.001)")
optional.add_argument("-nsamples", metavar='<int>', default=1000, type=int, help="Number of samples per seed used in tractography (default = 1000)")
optional.add_argument("-res", metavar='<float>', default=3, type=float, help="Resolution of matrix2 output (Default = 3 mm). Set to 0 for standard space resolution.")
optional.add_argument("-ptx_options", metavar='<options.txt>', default=False, help="Pass extra probtrackx2 options as a text file to override defaults")

# Optional computing arguments:
optional.add_argument("-gpu", action="store_true", default=False, help="Use GPU version")
optional.add_argument("-keepfiles", action="store_true", default=False, help="Do not delete temp files")
optional.add_argument("-version", dest='version', action="store_true", default=False, help="List the package versions")

if FSLINFMRIB is not None:
    hidden = parser.add_argument_group('Hidden arguments')
    hidden.add_argument("-dev_create", metavar='<path>', default=None, help="Path to development create_blueprint script (default=$FSLDIR/bin/create_blueprint)")
    hidden.add_argument("-dev_subcort", metavar='<path>', default=None, help="Path to development make_subcort_bp script (default=$FSLDIR/bin/make_subcort_bp)")
    hidden.add_argument("-queue", metavar='<str>', default=None, help="For job schedulers, specify the queue fsl_sub submits to")
    hidden.add_argument("-submem", metavar='<int>', nargs=2, type=int, help="Memory (GB) requested for tracking and blueprinting cluster jobs (default=[40 100])")
    hidden.add_argument("-subtime", metavar='<int>', nargs=2, type=int, help="Time (minutes) requested for tracking and blueprinting cluster jobs (default=300 for both jobs, 1440 for tracking if not GPU)")
    hidden.add_argument("-distnorm", action="store_true", default=False, help="Distance normalisation in matrix2 tracking? (default=False)")
    hidden.add_argument("-job_wait", metavar='<jobID>', default=None, help="Set a job dependency (default=None)")

argsa = parser.parse_args()

# print versions (xtract, xtract_data, ptx2)
if argsa.version:
    verlist = subprocess.run([os.path.join(FSLbin, "fslversion"), '-p'], capture_output=True, text=True)
    verlist = verlist.stdout.split('\n')
    verlist = [t for t in verlist if 'fsl-xtract' in t or 'ptx2' in t]
    print('Versions:')
    verlist = [print(f'{t.split(" - ")[0].replace(" ", "")}: {t.split(" - ")[1].replace(" ", "")}') for t in verlist]
    print('\n')

# rename arguments for ease
bpx = argsa.bpx
out = argsa.out
xtract = argsa.xtract
seeds = argsa.seeds.split(",")
target = argsa.target
res = argsa.res
stage = argsa.stage
native = argsa.native
errflag = 0

if FSLINFMRIB is not None:
    if argsa.dev_create is None:
        cbp_path = os.path.join(FSLbin, "create_blueprint")
    else:
        cbp_path = argsa.dev_create
        print(f'create_blueprint path is: {cbp_path}')
    if argsa.dev_subcort is None:
        subcort_path = os.path.join(FSLbin, "make_subcort_bp")
    else:
        subcort_path = argsa.dev_subcort
        print(f'make_subcort_bp path is: {subcort_path}')
    job_wait = argsa.job_wait
else:
    cbp_path = os.path.join(FSLbin, "create_blueprint")
    job_wait = None

if FSLINFMRIB is None:
    argsa.queue = None

if FSLINFMRIB is not None:
    distnorm = argsa.distnorm
else:
    distnorm = False
            
# check compulsory arguments
# input directories
if not os.path.isdir(bpx):
    print(f'Bedpostx folder {bpx} not found')
    errflag = 1
if not os.path.isdir(xtract):
    print(f'Bedpostx folder {xtract} not found')
    errflag = 1
# Check seeds/target
for seed in seeds:
    if os.path.isfile(seed) == False:
        print(f'Seed file {seed} not found')
        errflag = 1
if imgtest(target) == 0:
    print(f'Target file {target} not found.')
    errflag = 1

# check which space we are running tracking in    
if native is True and argsa.warps[0] is not None:
    print("You have selected the native space and ref space options")
    print("Must select EITHER '-native', '-ref <refimage> <diff2ref> <ref2diff>', OR use the default standard space")
    errflag=1
elif native is False and argsa.warps[0] is not None:
    print("Creating blueprint in standard space")    
    ref = argsa.warps[0]
    xtract2diff = argsa.warps[1]
    diff2xtract = argsa.warps[2]
    # check ref and warps
    if imgtest(ref) == 0:
        print(f'Reference file {ref} not found.')
        errflag = 1    
    if imgtest(xtract2diff) == 0:
        print(f'Warp file {xtract2diff} not found.')
        errflag = 1
    if imgtest(diff2xtract) == 0:
        print(f'Warp file {diff2xtract} not found.')
        errflag = 1   
elif native is True and argsa.warps[0] is None:
    print("Creating blueprint in native space")
errchk(errflag)

# check optional arguments
nseeds = len(seeds)
nrois, nwaypoints = 0, 0
if argsa.stops is not None:
    stops = argsa.stops
    if os.path.isfile(stops) == False:
        print(f'Stop file {stops} not found')
        errflag = 1
if argsa.wtstops is not None:
    wtstops = argsa.wtstops
    if os.path.isfile(wtstops) == False:
        print(f'Stop file {wtstops} not found')
        errflag = 1
if argsa.rois is not None:
    rois = argsa.rois.split(",")
    nrois = len(rois)
    if nrois != nseeds:
        print(f'Supplied {nseeds} seed but {nrois} ROIs')
        print('If using ROIs, must supply 1 per seed')
        errflag = 1
    for roi in rois:
        if os.path.isfile(roi) == False:
            print(f'ROI file {roi} not found')
            errflag = 1
if argsa.subseed is not None:
    if os.path.isfile(argsa.subseed) == False:
        print(f'Subcortical seed file {argsa.subseed} not found.')
        errflag = 1
    if nseeds == 2:
        seeds.append(argsa.subseed)
        foo = 0
    elif nseeds != 2:
        print('Using subcortical seed but only one surface seed')
        print('Subcortical seeding is currently limited to two surface seeds (whole-brain/two hemispheres)')
        errflag=1    
if argsa.exclusion is not None:
    exclusion = argsa.exclusion
    if exclusion == False:
        print(f'Exclusion mask file {exclusion} not found.')
        errflag = 1
errchk(errflag)

# GPU option
if argsa.gpu:
    ptxbin = ptxbin_gpu
    if FSLINFMRIB is not None:
        if argsa.submem is None:
            submem = [40, 100]
        else:
            submem = argsa.submem
        if argsa.subtime is None:
            subtime = [300, 300]
        else:
            subtime = argsa.subtime
    else:
        submem = [40, 100]
        subtime = [300, 300]
else:
    print('Warning: not using GPU mode - this may take a while. Consider downsampling the seed mask.')
    if FSLINFMRIB is not None:
        if argsa.submem is None:
            submem = [40, 100]
        else:
            submem = argsa.submem
        if argsa.subtime is None:
            subtime = [1440, 300]
        else:
            subtime = argsa.subtime
    else:
        submem = [40, 100]
        subtime = [1440, 300]

additional_opts = ''
if argsa.ptx_options:
    f = open(argsa.ptx_options, "r")
    additional_opts = f.read()
    if "\n" in additional_opts:
        print('Warning: new line(s) found in ptx_options. Removing.')
        additional_opts = additional_opts.strip()

# what are we running?
if stage == 1:
    print('Running tractography only')
elif stage == 2:
    print('Running blueprint processing only')
elif stage == 0:
    print('Running tractography and blueprint processing')

# make output directories
if out is None:
    out = os.path.join(xtract, 'blueprint')
if os.path.isdir(out):
    print('Warning: output directory already exists. XTRACT may overwrite existing content.')

os.makedirs(out, exist_ok=True)
os.makedirs(os.path.join(out, 'omatrix2'), exist_ok=True)

# Get tract list here
# if not provided by user, get list from under xtract folder
if stage != 1:
    if argsa.tract_list is not None:
        tracts = argsa.tract_list.split(',')
    else:
        tracts = glob.glob(os.path.join(xtract, 'tracts', '*', 'densityNorm.nii.gz'))
        tracts = [os.path.basename(os.path.dirname(t)) for t in tracts]
        tracts.sort()

# Resample ref to get target2
if res > 0:
    applyisoxfm(target, os.path.join(out, 'omatrix2', 'target2'), res, target, 'nearestneighbour')
else:
    foo = shutil.copyfile(target, os.path.join(out, 'omatrix2', 'target2.nii.gz'))


# Step 1 : Run omatrix2 for all seeds
ptxopts = f' --samples={os.path.join(bpx, "merged")} --mask={os.path.join(bpx, "nodif_brain_mask")} --target2={os.path.join(out, "omatrix2", "target2")}'
ptxopts += f' --omatrix2 --loopcheck --forcedir --opd --nsamples={argsa.nsamples} '

if native is False:
    ptxopts += f' --xfm={xtract2diff} --invxfm={diff2xtract} --seedref={ref} '
else:
    ptxopts += f' --seedref={os.path.join(out, "omatrix2", "target2")}'

if stage != 2:
    commands_1 = os.path.join(out, 'omatrix2', 'ptx_commands.txt')
    try:
        os.remove(commands_1)
    except OSError:
        pass

    working_dirs = []
    for COUNTER, seed in enumerate(seeds):
        r = remove_ext(os.path.basename(seed))
        print(f'\nSeed {COUNTER+1}:')
        print(f'seed {r}', end="")
        
        if seed is not argsa.subseed:
            working_dirs.append(os.path.join(out, 'omatrix2', f'omatrix2_{r}'))
            os.makedirs(working_dirs[COUNTER], exist_ok=True)

            s = seed
            # convert gii to ascii with medial wall, if using
            if nrois != 0:
                roi = rois[COUNTER]
                if not seed.endswith('.asc'):
                    s = conv_ascii(seed, roi, working_dirs[COUNTER])
                m = remove_ext(os.path.basename(roi))
                print(f': medial wall/ROI {m}', end="")

            hemi_opts = f'{ptxopts} --seed={s} --dir={working_dirs[COUNTER]}'
            # hemi_opts = f'{ptxopts} --seed={seed} --dir={working_dir}'
                    
            # if nwaypoints > 0:
            #     waypoint = waypoints[COUNTER]
            #     m = remove_ext(os.path.basename(waypoint))
            #     print(f': waypoint {m}', end="")
            #     hemi_opts += f' --waypoints={waypoint}'
                
            if argsa.stops is not None:
                m = remove_ext(os.path.basename(stops))
                print(f': stop {m}', end="")
                hemi_opts += f' --stop={stops}'
            if argsa.wtstops is not None:
                m = remove_ext(os.path.basename(wtstops))
                print(f': wtstop {m}', end="")
                hemi_opts += f' --wtstop={wtstops}'
            if argsa.exclusion is not None:
                m = remove_ext(os.path.basename(exclusion))
                print(f': exclusion {m}', end="")
                hemi_opts += f' --avoid={exclusion}'
            
            print('')
            # append extra options and save to commands file
            hemi_opts += f' {additional_opts}'
            
            with open(commands_1, 'a') as f:
                print(f'{ptxbin} {hemi_opts};', file=f)
        
        elif seed is argsa.subseed:
            working_dirs.append(os.path.join(out, 'omatrix2', 'subcortical'))
            os.makedirs(working_dirs[COUNTER], exist_ok=True)
            # call make_subcort_bp
            if argsa.keepfiles is False:
                cmd = f' {subcort_path} -xtract_folder {xtract} -out {working_dirs[COUNTER]} -subseed {seed} -tracts {",".join(tracts)} -thr {argsa.thr}'
            else:
                cmd = f' {subcort_path} -xtract_folder {xtract} -out {working_dirs[COUNTER]} -subseed {seed} -tracts {",".join(tracts)} -thr {argsa.thr} -keepfiles'

            with open(commands_1, 'a') as f:
                print(f'{cmd};', file=f)
    
    subprocess.run(['chmod', '+x', commands_1]);

# Step 2 : Resample xtract, load all xtract thingies and matrix2 output and multiply them        
if stage != 1:
    commands_2 = os.path.join(out, 'bp_commands.txt')
    try:
        os.remove(commands_2)
    except OSError:
        pass    
    os.makedirs(os.path.join(out, 'temp_xtract'), exist_ok=True)
    
    for t in tracts:
        foo = shutil.copyfile(os.path.join(xtract, 'tracts', t, 'densityNorm.nii.gz'), os.path.join(out, 'temp_xtract', f'{t}.nii.gz'))

    cmd = ''
    # re-sample tracts if needed - assumes that all xtract tracts are equal dimensions - they should be!
    targ_form = subprocess.run([os.path.join(FSLbin, 'fslorient'), '-getsform', os.path.join(out, 'omatrix2', 'target2')], capture_output=True)
    tract_form = subprocess.run([os.path.join(FSLbin, 'fslorient'), '-getsform', os.path.join(out, 'temp_xtract', f'{tracts[0]}.nii.gz')], capture_output=True)
    if targ_form.stdout != tract_form.stdout:
        print('XTRACT resolution different from matrix2:')
        print(f' --- will resample xtract tracts to {res} mm')
        for t in tracts:
            cmd += f'{os.path.join(FSLbin, "flirt")} -in {os.path.join(out, "temp_xtract", t)} -out {os.path.join(out, "temp_xtract", t)} -ref {os.path.join(out, "omatrix2", "target2")} -applyisoxfm {res} -interp trilinear;'

    if argsa.thr > 0:
        for t in tracts:
            cmd += f'{os.path.join(FSLbin, "fslmaths")} {os.path.join(out, "temp_xtract", t)} -thr {argsa.thr} {os.path.join(out, "temp_xtract", t)};'
            
    # loop through seeds and prepare commands
    # start of python command
    cmd += f' {cbp_path} -xtract_folder {os.path.join(out, "temp_xtract")} -thr {argsa.thr}'
    
    working_dirs = []
    for seed in seeds:
        if seed is not argsa.subseed:
            r = remove_ext(os.path.basename(seed))
            working_dirs.append(os.path.join(out, 'omatrix2', f'omatrix2_{r}'))
        elif seed is argsa.subseed:
            working_dirs.append(os.path.join(out, 'omatrix2', 'subcortical'))

    py_args = f' -ptx_folder {",".join(working_dirs)}'
    py_args = f'{py_args} -seed_path {",".join(seeds)}' 

    if nrois != 0:
        py_args = f'{py_args} -roi_path {",".join(rois)}'

    cmd += f'{py_args} -tracts {",".join(tracts)}'
    if argsa.savetxt:
        cmd += f' -savetxt'
    if distnorm:
        cmd += f' -distnorm {distnorm}'
    if argsa.prefix is not None:
        cmd += f' -prefix {argsa.prefix}'
    if argsa.keepfiles is False:
        cmd += f'; rm -rf {os.path.join(out, "temp_xtract")}'
    
    with open(commands_2, 'a') as f:
        print(f'{cmd}', file=f)
    subprocess.run(['chmod', '+x', commands_2]);


# Launch jobs
print('\nLaunching job(s)\n')
# function to run jobs, handling jodIDs and dependencies
def run_job(command_file, out, step, jobID=''):
    if step == 1:
        if job_wait is None:
            if argsa.gpu:
                if argsa.queue is None:
                    JID = subprocess.run([os.path.join(FSLbin, 'fsl_sub'), '-T', str(subtime[0]), '-R', str(submem[0]), '--coprocessor=cuda', '-l', os.path.join(out, "logs"), '-N', 'blueprint_track', '-t', command_file], capture_output=True, text=True)
                else:
                    JID = subprocess.run([os.path.join(FSLbin, 'fsl_sub'),'-q', argsa.queue, '-T', str(subtime[0]), '-R', str(submem[0]), '--coprocessor=cuda', '-l', os.path.join(out, "logs"), '-N', 'blueprint_track', '-t', command_file], capture_output=True, text=True)
            else:
                if argsa.queue is None:
                    JID = subprocess.run([os.path.join(FSLbin, 'fsl_sub'), '-T', str(subtime[0]), '-R', str(submem[0]), '-l', os.path.join(out, "logs"), '-N', 'blueprint_track', '-t', command_file], capture_output=True, text=True)
                else:
                    JID = subprocess.run([os.path.join(FSLbin, 'fsl_sub'),'-q', argsa.queue, '-T', str(subtime[0]), '-R', str(submem[0]), '-l', os.path.join(out, "logs"), '-N', 'blueprint_track', '-t', command_file], capture_output=True, text=True)
        else:
            if argsa.gpu:
                if argsa.queue is None:
                    JID = subprocess.run([os.path.join(FSLbin, 'fsl_sub'), '-T', str(subtime[0]), '-R', str(submem[0]), '-j', job_wait, '--coprocessor=cuda', '-l', os.path.join(out, "logs"), '-N', 'blueprint_track', '-t', command_file], capture_output=True, text=True)
                else:
                    JID = subprocess.run([os.path.join(FSLbin, 'fsl_sub'),'-q', argsa.queue, '-T', str(subtime[0]), '-R', str(submem[0]), '-j', job_wait, '--coprocessor=cuda', '-l', os.path.join(out, "logs"), '-N', 'blueprint_track', '-t', command_file], capture_output=True, text=True)
            else:
                if argsa.queue is None:
                    JID = subprocess.run([os.path.join(FSLbin, 'fsl_sub'), '-T', str(subtime[0]), '-R', str(submem[0]), '-j', job_wait, '-l', os.path.join(out, "logs"), '-N', 'blueprint_track', '-t', command_file], capture_output=True, text=True)
                else:
                    JID = subprocess.run([os.path.join(FSLbin, 'fsl_sub'),'-q', argsa.queue, '-T', str(subtime[0]), '-R', str(submem[0]), '-j', job_wait, '-l', os.path.join(out, "logs"), '-N', 'blueprint_track', '-t', command_file], capture_output=True, text=True)
    elif step == 2:
        if jobID != '':
            if argsa.queue is None:
                JID = subprocess.run([os.path.join(FSLbin, 'fsl_sub'), '-T', str(subtime[1]), '-R', str(submem[1]), '-j', jobID, '-l', os.path.join(out, "logs"), '-N', 'create_blueprint', 'sh', command_file], capture_output=True, text=True)
            else:
                JID = subprocess.run([os.path.join(FSLbin, 'fsl_sub'),'-q', argsa.queue, '-T', str(subtime[1]), '-R', str(submem[1]), '-j', jobID, '-l', os.path.join(out, "logs"), '-N', 'create_blueprint', 'sh', command_file], capture_output=True, text=True)
        else:
            if argsa.queue is None:
                JID = subprocess.run([os.path.join(FSLbin, 'fsl_sub'), '-T', str(subtime[1]), '-R', str(submem[1]), '-l', os.path.join(out, "logs"), '-N', 'create_blueprint', 'sh', command_file], capture_output=True, text=True)
            else:
                JID = subprocess.run([os.path.join(FSLbin, 'fsl_sub'),'-q', argsa.queue, '-T', str(subtime[1]), '-R', str(submem[1]), '-l', os.path.join(out, "logs"), '-N', 'create_blueprint', 'sh', command_file], capture_output=True, text=True)
    return JID.stdout.splitlines()[0]

hq = subprocess.run([os.path.join(FSLbin, 'fsl_sub'), '--has_queues'], capture_output=True, text=True)
if hq.stdout.rstrip() == 'Yes':
    if stage == 1:
        jobID = run_job(commands_1, out, stage)
        print(f'probtrackx job ID: {jobID}')
    elif stage == 2:
        jobID = run_job(commands_2, out, stage)
        print(f'blueprint job ID: {jobID}')
    else:
        jobID = run_job(commands_1, out, 1)
        print(f'probtrackx Job ID: {jobID}')
        jobID = run_job(commands_2, out, 2, jobID=jobID)
        print(f'blueprint job ID: {jobID}')
else:
    if stage == 1:
        subprocess.run([commands_1], shell=True)
    elif stage == 2:
        subprocess.run([commands_2], shell=True)
    else:
        subprocess.run([commands_1], shell=True)
        subprocess.run([commands_2], shell=True)

quit()
