#!/bin/sh

# bedpostx_mgh
# (This script was modified from FSL's bedpostx by Anastasia Yendiki)

#   Copyright (C) 2004,2013 University of Oxford
#
#   Part of FSL - FMRIB's Software Library
#   http://www.fmrib.ox.ac.uk/fsl
#   fsl@fmrib.ox.ac.uk
#
#   Developed at FMRIB (Oxford Centre for Functional Magnetic Resonance
#   Imaging of the Brain), Department of Clinical Neurology, Oxford
#   University, Oxford, UK
#
#
#   LICENCE
#
#   FMRIB Software Library, Release 6.0 (c) 2018, The University of
#   Oxford (the "Software")
#
#   The Software remains the property of the Oxford University Innovation
#   ("the University").
#
#   The Software is distributed "AS IS" under this Licence solely for
#   non-commercial use in the hope that it will be useful, but in order
#   that the University as a charitable foundation protects its assets for
#   the benefit of its educational and research purposes, the University
#   makes clear that no condition is made or to be implied, nor is any
#   warranty given or to be implied, as to the accuracy of the Software,
#   or that it will be suitable for any particular purpose or for use
#   under any specific conditions. Furthermore, the University disclaims
#   all responsibility for the use which is made of the Software. It
#   further disclaims any liability for the outcomes arising from using
#   the Software.
#
#   The Licensee agrees to indemnify the University and hold the
#   University harmless from and against any and all claims, damages and
#   liabilities asserted by third parties (including claims for
#   negligence) which arise directly or indirectly from the use of the
#   Software or the sale of any products based on the Software.
#
#   No part of the Software may be reproduced, modified, transmitted or
#   transferred in any form or by any means, electronic or mechanical,
#   without the express permission of the University. The permission of
#   the University is not required if the said reproduction, modification,
#   transmission or transference is done without financial return, the
#   conditions of this Licence are imposed upon the receiver of the
#   product, and all original and amended source code is included in any
#   transmitted product. You may be held legally responsible for any
#   copyright infringement that is caused or encouraged by your failure to
#   abide by these terms and conditions.
#
#   You are not permitted under this Licence to use this Software
#   commercially. Use for which any financial return is received shall be
#   defined as commercial use, and includes (1) integration of all or part
#   of the source code or the Software into a product for sale or license
#   by or on behalf of Licensee to third parties or (2) use of the
#   Software or any derivative of it for research with the final aim of
#   developing software products for sale or license to a third party or
#   (3) use of the Software or any derivative of it for research with the
#   final aim of developing non-software products for sale or license to a
#   third party, or (4) use of the Software to provide any service to an
#   external organisation for which payment is received. If you are
#   interested in using the Software commercially, please contact Oxford
#   University Innovation ("OUI"), the technology transfer company of the
#   University, to negotiate a licence. Contact details are:
#   fsl@innovation.ox.ac.uk quoting Reference Project 9564, FSL.

# Check if this is being run on PBS but isn't a job itself (AY)
if [ -d /pbs ] && [ -z $PBS_JOBID ] ; then
    dopbsjobs=1;
else
    dopbsjobs=0;
fi

BPDIR=`dirname $0`

export LC_ALL=C

Usage() {
    echo ""
    echo "Usage: bedpostx <subject directory> [options]"
    echo ""
    echo "expects to find bvals and bvecs in subject directory"
    echo "expects to find data and nodif_brain_mask in subject directory"
    echo "expects to find grad_dev in subject directory, if -g is set"
    echo "options (old syntax)"
    echo "-n (number of fibres per voxel, default 3)"
    echo "-w (ARD weight, more weight means less secondary fibres per voxel, default 1)"
    echo "-b (burnin period, default 1000)"
    echo "-j (number of jumps, default 1250)"
    echo "-s (sample every, default 25)"
    echo "-model (Deconvolution model. 1: with sticks, 2: with sticks with a range of diffusivities (default), 3: with zeppelins)"
    echo "-g (consider gradient nonlinearities, default off)"
    #echo "-c do not use CUDA capable hardware/queue (if found)"
    echo ""
    echo ""
    echo "ALTERNATIVELY: you can pass on xfibres options onto directly bedpostx"
    echo " For example:  bedpostx <subject directory> --noard --cnonlinear"
    echo " Type 'xfibres --help' for a list of available options "
    echo " Default options will be bedpostx default (see above), and not xfibres default."
    echo ""
    echo "Note: Use EITHER old OR new syntax."
    echo ""
    echo "This is a modified version of FSL's bedpostx that is compatible with"
    echo "the PBS queueing system running on launchpad and can be used by MGH"
    echo "Martinos Center users to parallelize computation."
    echo ""
    echo "This script runs the exact same algorithm as bedpostx, but it submits"
    echo "the processing for each slice as a separate PBS job. Thus,"
    echo "bedpostx_mgh will give you the same result as bedpostx, but do it"
    echo "an order of magnitude faster."
    echo ""
    echo "Command line options and the structure of the input directory are the"
    echo "same as for the original bedpostx."
    echo ""
    echo "Do not submit bedpostx_mgh as a job on launchpad. Run it directly at"
    echo "the launchpad prompt and it will submit the jobs for you. Note that"
    echo "this will generate a large number of jobs, so don't launch it on more"
    echo "than one data set at the same time. You may get some PBS emails about"
    echo "jobs being held that you have to release - you can ignore these"
    echo "emails as this script will take care of that."
    exit 1
}

# Find FSL version (AY)
fslver=(`cat $FSLDIR/etc/fslversion | sed 's/\./ /g'`)
if [ -z ${fslver[1]} ]; then fslver[1]=0; fi
if [ -z ${fslver[2]} ]; then fslver[2]=0; fi
if [ ${fslver[0]} -lt 6 ]; then
    echo "Outdated FSL version (${fslver[0]}.${fslver[1]}.${fslver[2]})"
    echo "Use 6.0.0 or newer"
    exit 1
fi

monitor(){
    cat <<EOM > ${subjdir}.bedpostX/monitor
#!/bin/sh
echo
echo "----- Bedpostx Monitor -----"
finished=0
lastprinted=0
havedad=2
while [ \$finished -eq 0 ] ; do
    nfin=0
    slice=0
    errorFiles=\`ls ${subjdir}.bedpostX/logs/*.e* 2> /dev/null \`
    for errorFile in \$errorFiles
    do
        if [ -s \$errorFile ]; then
            echo An error ocurred. Please check file \$errorFile
            kill -9 $$
            exit 1
        fi
    done
    while [ \$slice -le $nslices ];do
        if [ -e ${subjdir}.bedpostX/logs/monitor/\$slice ]; then
            nfin=\$((\$nfin + 1))
        fi
        slice=\$((\$slice + 1))
    done
    newmessages=\$((\$nfin - \$lastprinted))
    while [ "\$newmessages" -gt 0 ];do
        lastprinted=\$((\$lastprinted + 1))
        echo \$lastprinted slices processed out of $nslices
        newmessages=\$((\$newmessages - 1))
    done
    if [ -f ${subjdir}.bedpostX/xfms/eye.mat ] ; then
        finished=1
        echo "All slices processed"
	exit 
    fi
    if [ ! \$havedad -gt 0 ]; then
       exit 0
    fi
    if [ "x$SGE_ROOT" = "x" ] && [ $dopbsjobs -eq 0 ]; then
        havedad=\`ps -e -o pid 2>&1 | grep "$$\\b" | wc -l\`        
    fi
    sleep 180;
done
EOM
    chmod +x ${subjdir}.bedpostX/monitor
}

make_absolute(){
    dir=$1;
    if [ -d ${dir} ]; then
	OLDWD=`pwd`
	cd ${dir}
	dir_all=`pwd`
	cd $OLDWD
    else
	dir_all=${dir}
    fi
    echo ${dir_all}
}

[ "$1" = "" ] && Usage

myopts="$@"

#parse option arguments
nfibres=3
fudge=1
burnin=1000
njumps=1250
sampleevery=25
model=2
gflag=0
nocuda=0
other=""

while [ ! -z "$1" ]
do
  if [ -d "$1" ]; then		# Allow multiple input directories (AY)
    dirlist="$dirlist $1"
  else
    case "$1" in
        -n) nfibres=$2;shift;;
        -w) fudge=$2;shift;;
        -b) burnin=$2;shift;;
        -j) njumps=$2;shift;;
        -s) sampleevery=$2;shift;;
        -model) model=$2;shift;;
        -g) gflag=1;;
        -c) nocuda=1;;
        *) other=$other" "$1;;
    esac
  fi
  shift
done
opts="--nf=$nfibres --fudge=$fudge --bi=$burnin --nj=$njumps --se=$sampleevery --model=$model"
defopts="--cnonlinear"
opts="$opts $defopts $other"
if [ -n "$SGE_ROOT" ]; then
	# Look for an SGE queue
	if [ -n "$FSLGECUDAQ" -a "$nocuda" -eq 0 ]; then
		# Test that the queue exists
		qconf -sq $FSLGECUDAQ 2>&1 >/dev/null
		if [ $? -eq 0 ]; then
			# CUDA queue exists
			if [ -x $0_gpu ]; then
				exec $0_gpu $myopts
			fi
		fi
	fi
fi

# Process multiple input directories, queue them up serially (AY)
for subjdir in $dirlist; do

subjdir=`make_absolute $subjdir`
subjdir=`echo $subjdir | sed 's/\/$/$/g'`

echo subjectdir is $subjdir

#check that all required files exist

if [ ! -d $subjdir ]; then
	echo "subject directory $1 not found"
	exit 1
fi

if [ ! -e ${subjdir}/bvecs ]; then
    if [ -e ${subjdir}/bvecs.txt ]; then
	mv ${subjdir}/bvecs.txt ${subjdir}/bvecs
    else
	echo "${subjdir}/bvecs not found"
	exit 1
    fi
fi

if [ ! -e ${subjdir}/bvals ]; then
    if [ -e ${subjdir}/bvals.txt ]; then
	mv ${subjdir}/bvals.txt ${subjdir}/bvals
    else
	echo "${subjdir}/bvals not found"
	exit 1
    fi
fi

if [ `${FSLDIR}/bin/imtest ${subjdir}/data` -eq 0 ]; then
	echo "${subjdir}/data not found"
	exit 1
fi

if [ ${gflag} -eq 1 ]; then
    if [ `${FSLDIR}/bin/imtest ${subjdir}/grad_dev` -eq 0 ]; then
	echo "${subjdir}/grad_dev not found"
	exit 1
    fi
fi

if [ `${FSLDIR}/bin/imtest ${subjdir}/nodif_brain_mask` -eq 0 ]; then
	echo "${subjdir}/nodif_brain_mask not found"
	exit 1
fi

if [ -e ${subjdir}.bedpostX/xfms/eye.mat ]; then
	echo "${subjdir} has already been processed: ${subjdir}.bedpostX." 
	echo "Delete or rename ${subjdir}.bedpostX before repeating the process."
	exit 1
fi

echo Making bedpostx directory structure

mkdir -p ${subjdir}.bedpostX/
mkdir -p ${subjdir}.bedpostX/diff_slices
mkdir -p ${subjdir}.bedpostX/logs
mkdir -p ${subjdir}.bedpostX/logs/monitor
mkdir -p ${subjdir}.bedpostX/xfms

mailto=`whoami`
echo Queuing preprocessing stages
if [ "x$mergeid" = "x" ] ; then
    preprocid=`${BPDIR}/fsl_sub_mgh -T 60 -m a -N bedpostx_pre -l ${subjdir}.bedpostX/logs ${FSLDIR}/bin/bedpostx_preproc.sh ${subjdir} ${gflag}`
else
    preprocid=`${BPDIR}/fsl_sub_mgh -j $mergeid -T 60 -m a -N bedpostx_pre -l ${subjdir}.bedpostX/logs ${FSLDIR}/bin/bedpostx_preproc.sh ${subjdir} ${gflag}`
fi

echo Queuing parallel processing stage
nslices=`${FSLDIR}/bin/fslval ${subjdir}/data dim3`
[ -f ${subjdir}.bedpostX/commands.txt ] && rm ${subjdir}.bedpostX/commands.txt

monitor
if [ "x$SGE_ROOT" = "x" ] && [ $dopbsjobs -eq 0 ]; then
    ${subjdir}.bedpostX/monitor&
fi

slice=0
while [ $slice -lt $nslices ]
do
    slicezp=`$FSLDIR/bin/zeropad $slice 4`
    if [ `$FSLDIR/bin/imtest ${subjdir}.bedpostX/diff_slices/data_slice_$slicezp/dyads1` -eq 1 ] && [ -f ${subjdir}.bedpostX/logs/monitor/$slice ];then
	echo "slice $slice has already been processed"
    else
        [ -f ${subjdir}.bedpostX/logs/monitor/$slice ] && rm ${subjdir}.bedpostX/logs/monitor/$slice
	if [ ${gflag} -eq 1 ]; then
	    gopts="$opts --gradnonlin=${subjdir}/grad_dev_slice_${slicezp}"
	else
	    gopts=$opts
	fi    
	echo "${FSLDIR}/bin/bedpostx_single_slice.sh ${subjdir} ${slice} $gopts" \
	    >> ${subjdir}.bedpostX/commands.txt
    fi
    slice=$(($slice + 1))
done
bedpostid=`${BPDIR}/fsl_sub_mgh -j $preprocid -l ${subjdir}.bedpostX/logs -m a -N bedpostx -t ${subjdir}.bedpostX/commands.txt`

echo Queuing post processing stage
mergeid=`${BPDIR}/fsl_sub_mgh -j $bedpostid -T 60 -m ae -N bedpostx_post -l ${subjdir}.bedpostX/logs ${FSLDIR}/bin/bedpostx_postproc.sh $subjdir`

echo $mergeid > ${subjdir}.bedpostX/logs/postproc_ID

if [ "x$SGE_ROOT" != "x" ] || [ $dopbsjobs -eq 1 ]; then
    echo
    echo Type ${subjdir}.bedpostX/monitor to show progress.
    echo Type ${subjdir}.bedpostX/cancel to terminate all the queued tasks.
    cat <<EOC > ${subjdir}.bedpostX/cancel
#!/bin/sh
qdel $mergeid $bedpostid $preprocid
EOC
    chmod +x ${subjdir}.bedpostX/cancel

    echo
    echo You will get an email at the end of the post-processing stage.
    echo
else
    sleep 190
fi

# Process multiple input directories, queue them up serially (AY)
done