#!/bin/sh

##
##  Copyright 2012-2014 SRI International
##
##  This file is part of the Computational Morphometry Toolkit.
##
##  http://www.nitrc.org/projects/cmtk/
##
##  The Computational Morphometry Toolkit is free software: you can
##  redistribute it and/or modify it under the terms of the GNU General Public
##  License as published by the Free Software Foundation, either version 3 of
##  the License, or (at your option) any later version.
##
##  The Computational Morphometry Toolkit is distributed in the hope that it
##  will be useful, but WITHOUT ANY WARRANTY; without even the implied
##  warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
##  GNU General Public License for more details.
##
##  You should have received a copy of the GNU General Public License along
##  with the Computational Morphometry Toolkit.  If not, see
##  <http://www.gnu.org/licenses/>.
##
##  $Revision$
##
##  $LastChangedDate$
##
##  $LastChangedBy$
##

export CMTK_BINARY_DIR=${CMTK_BINARY_DIR:-@CMTK_BINARY_DIR_CONFIG@}
 
# Get the cmtk_functions.sh script from the scripts/ directory in the CMTK source tree
. ${CMTK_BINARY_DIR}/cmtk_functions.sh
 
# Check for command line arguments, print help if none given
if test $# -lt 2; then
    echo "Correct distortion and subject motion in diffusion-weighted MR images."
    echo
    echo "For this script to be applicable, the b=0 image must have been acquired twice, "
    echo "once with standard phase encoding direction and once with phase encoding reversed."
    echo
    echo "USAGE: $0 outdir b0Reverse b0Forward bX1 [bX2 ...]"
    exit 2
fi

outdir=$1
b0Rev=$2
b0Fwd=$3

if ! cmtk geomatch -v ${b0Rev} ${b0Fwd}; then
    echo "ERROR: the two b=0 images have mismatched geometries or spatial coordinated:"
    echo "  ${b0Rev}"
    echo "  ${b0Fwd}"
    exit 1
fi

shift 3
bXlist="$*"

#
# First, do eddy currents by registering all bX images to forward b=0 image
#
for bX in ${bXlist}; do
    base=`basename $bX`
    pref=`echo ${base} | sed 's/\..*//g'`

    xform=${outdir}/eddy/b0_${pref}.xform
    if CMTK_needs_update_and_lock ${xform} ${b0Fwd} ${bX}; then
	echo "Computing eddy current correction for ${base}..."
	cmtk registrationx --pad-flt -1 --restrict-in-plane xy --cubic --dofs 9,12 --auto-multi-levels 2 --nmi -o ${xform} ${b0Fwd} ${bX}
	CMTK_lockfile_delete ${xform}
    fi
done
 
b0FwdCorr=${outdir}/`basename ${b0Fwd}`
b0RevCorr=${outdir}/`basename ${b0Rev}`
if CMTK_needs_update_and_lock ${outdir}/dfield_fwd.nrrd ${b0Fwd} ${b0Rev}; then
    echo "Computing deformation..."
    cmtk epiunwarp --no-flip --write-jacobian-fwd ${outdir}/jacobian.nii --smooth-sigma-max 16 --smooth-sigma-min 0 --smooth-sigma-diff 0.25 --iterations 5 --smoothness-constraint-weight 1e4 --phase-encode-ap ${b0Fwd} ${b0Rev} ${b0FwdCorr} ${b0RevCorr} ${outdir}/dfield_fwd.nrrd
    CMTK_lockfile_delete ${outdir}/dfield_fwd.nrrd
fi

unwarpedList="${b0FwdCorr}"
for bX in ${bXlist}; do
    base=`basename $bX`
    pref=`echo ${base} | sed 's/\..*//g'`
    xform=${outdir}/eddy/b0_${pref}.xform

    if CMTK_needs_update_and_lock ${outdir}/${base} ${bX} ${outdir}/dfield_fwd.nrrd ${xform}; then
	echo "Unwarping $bX"
	tmpdir=$(mktemp -d)
	# First, reslice using PEpolar unwarp and eddy current correction
	cmtk reformatx --sinc-cosine -o ${tmpdir}/reslice.nii --pad-floating -1 --floating ${bX} ${b0FwdCorr} ${outdir}/dfield_fwd.nrrd ${xform}
	# Second, multiply with Jacobian volume correction field
	cmtk imagemath --in ${outdir}/jacobian.nii ${tmpdir}/reslice.nii --mul --out ${tmpdir}/reslice_j.nii
	# Third, make a mask of bad slices (0 - bad slice, 1 - good slice), based on requirement that bad slices be marked as -1
	cmtk convertx --float --binarize-thresh -0.5 --byte ${bX} ${tmpdir}/mask.nii
	# Apply bad-slice mask back to resliced and Jacobian-corrected volume
	cmtk convertx --thresh-below 0 --set-padding -1 --mask ${tmpdir}/mask.nii ${tmpdir}/reslice_j.nii ${outdir}/${base}
	rm -rf ${tmpdir}
	CMTK_lockfile_delete ${outdir}/${base}
    fi

    unwarpedList="${unwarpedList} ${outdir}/${base}"

done

##
## Motion correction
##

average=${outdir}/bAverage.nii
if CMTK_needs_update_and_lock ${average} ${b0FwdCorr} ${b0RevCorr} ${unwarpedList}; then
    cmtk imagemath --in ${b0FwdCorr} ${b0RevCorr} ${unwarpedList} --average --out ${average}
    CMTK_lockfile_delete ${average}
fi

for bX in ${b0FwdCorr} ${bXlist}; do
    base=`basename $bX`
    bXUnwarp=${outdir}/${base}
    pref=`echo ${base} | sed 's/\..*//g'`

    if [ "${bX}" = "${b0FwdCorr}" ]; then
	if CMTK_needs_update_and_lock ${outdir}/motion/${base} ${b0FwdCorr}; then
	    cmtk convertx ${b0FwdCorr} ${outdir}/motion/${base}
	    CMTK_lockfile_delete  ${outdir}/motion/${base}
	fi
    else
	xform=${outdir}/motion/xforms/b0Unwarp_${pref}.xform
	if CMTK_needs_update_and_lock ${xform} ${average} ${bXUnwarp} ; then
	    cmtk registrationx --pad-flt -1 --linear --dofs 6 --auto-multi-levels 3 --ncc -o ${xform} ${average} ${bXUnwarp}
	    CMTK_lockfile_delete ${xform}
	fi

	if CMTK_needs_update_and_lock ${outdir}/motion/${base} ${bXUnwarp} ${xform}; then
	    echo "Reformatting motion-corrected $bX"
	    tmpdir=$(mktemp -d)
	    # First, re-reslice the unwarped image
	    cmtk reformatx --sinc-cosine -o ${tmpdir}/reslice.nii --pad-floating -1 --floating ${bXUnwarp} ${average} ${xform}
	    # Second, make a mask of bad slices (0 - bad slice, 1 - good slice), based on requirement that bad slices be marked as -1
	    cmtk convertx --float --binarize-thresh -0.5 ${bX} ${tmpdir}/mask.nii
	    # Third, reslice the mask using the same motion xform and interpolation as the image
	    cmtk reformatx --sinc-cosine -o ${tmpdir}/mask_reslice.nii --floating ${tmpdir}/mask.nii ${average} ${xform}
	    # Fourth, threshold resliced mask to include only those pixels that have more then 50% support
	    cmtk convertx --binarize-thresh 0.5 ${tmpdir}/mask_reslice.nii ${tmpdir}/mask_thresh.nii
	    # Fifth, create final image by masking out pixels with under 50% support
	    cmtk convertx --set-padding -1 --mask ${tmpdir}/mask_thresh.nii --thresh-below 0 ${tmpdir}/reslice.nii ${outdir}/motion/${base}
	    rm -rf ${tmpdir}
	    CMTK_lockfile_delete ${outdir}/motion/${base}
	fi
    fi
done

