#!/bin/bash
#============================================================================
#
#  Program:     DTI ToolKit (DTI-TK)
#  Module:      $RCSfile: dti_diffeomorphic_reg,v $
#  Language:    bash
#  Date:        $Date: 2011/12/21 20:39:22 $
#  Version:     $Revision: 1.1.1.1 $
#
#  Copyright (c) Gary Hui Zhang (garyhuizhang@gmail.com).
#  All rights reserverd.
#
#  DTI-TK 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.
#
#  DTI-TK 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 DTI-TK.  If not, see <http://www.gnu.org/licenses/>.
#============================================================================

#
# deformable alignment of some subject to some template
#

#
# source PATH setting from ~.bashrc
# required for qsub to work
#
if [ -e ~/.bashrc ]
then
	. ~/.bashrc
elif [ -e ~/.bash_profile ]
then
	. ~/.bash_profile
fi

# check DTITK_ROOT variable
if [ -z "${DTITK_ROOT}" ]
then
	echo "Environment variable DTITK_ROOT is not defined"
	exit 1
fi

# source dtitk_common.sh
. ${DTITK_ROOT}/scripts/dtitk_common.sh

#
# utility: pwaAtvCGM
#
function pwaAtvCGM {
if [ $# -lt 11 ]
then
	echo "usage: $0 iter template_prefix template_base_prefix subject_prefix mask level_start level_end SMOption prior reg ftol"
	exit 1
fi

# parse the arguments
iter=$1
template_pref=$2
template_base_pref=$3
subject_pref=$4
mask=$5
level0=$6
levels=$7
SMOption=$8
prior=$9
reg=${10}
ftol=${11}

let piter=levels-level0
# run the incremental registration
piecewiseAtvCGM -template ${template_pref}.nii.gz -subject ${subject_pref}_diffeo.nii.gz -mask ${mask} -ftol ${ftol} -level0 ${level0} -iter ${piter} -SMOption ${SMOption} -prior ${prior} -reg ${reg} -out ${subject_pref}_to_${template_base_pref}.nii.gz
check_exit_code $?

let levels=levels-1

# compute the deformation field
dfpref=${subject_pref}_to_${template_base_pref}.${levels}
pwaToDisp -in ${dfpref}.pwa -vsize ${lengthscale} ${lengthscale} ${lengthscale} -out ${dfpref}.df.nii.gz
check_exit_code $?
dfToDiffeomorphic -in ${dfpref}.df.nii.gz -out ${dfpref}.df.nii.gz -smooth 2
check_exit_code $?

jac_mask=${subject_pref}_jac_mask.nii.gz
if [ ${iter} -eq 1 ]
then
	SVResample -in ${mask} -out ${jac_mask} -target ${dfpref}.df.nii.gz
	check_exit_code $?
	BinaryThresholdImageFilter ${jac_mask} ${jac_mask} 0.5 1 1 0
	check_exit_code $?
else
	dfComposition -df2 ${subject_pref}_diffeo.df.nii.gz -df1 ${dfpref}.df.nii.gz -out ${dfpref}.df.nii.gz
	check_exit_code $?
fi

# remove the old pwa
pwacount=${level0}
let levels=levels+1
while [ ${pwacount} -lt ${levels} ]
do
	rm -fr ${subject_pref}_to_${template_base_pref}.${pwacount}.pwa
	let pwacount=pwacount+1
done

## check the jacobian value
dfToJacobian -in ${dfpref}.df.nii.gz
check_exit_code $?
jacstats=`SVtool -in ${dfpref}.df_jac.nii.gz -stats -mask ${jac_mask} | grep "mean = "`
if [ ${iter} -eq 1 ]
then
	echo JACOBIAN STATISTICS: after current iteration  $jacstats
else
	echo JACOBIAN STATISTICS: after previous iteration `SVtool -in ${subject_pref}_diffeo.df_jac.nii.gz -stats -mask ${jac_mask} | grep "mean = "` after current iteration $jacstats
fi
jacmin=`echo $jacstats | awk '{ print $6 }'`
jacmax=`echo $jacstats | awk '{ print $9 }'`
isjacmingood=`echo "$jacmin > ${jaclimit[0]}" | bc`
isjacmaxgood=`echo "$jacmax < ${jaclimit[1]}" | bc`

if [ ${isjacmingood} -eq 1 -a ${isjacmaxgood} -eq 1 ]
then
	## create the new warped image
	deformationSymTensor3DVolume -in ${subject_pref}.nii.gz -out ${subject_pref}_diffeo_current.nii.gz -trans ${dfpref}.df.nii.gz
	check_exit_code $?
	
	# check whether to commit to another iteration or quit
	previous=`TVtool -SMOption DDS -in ${template_pref}.nii.gz -sm ${subject_pref}_diffeo.nii.gz |grep Similarity | awk '{print $3}'`
	current=`TVtool -SMOption DDS -in ${template_pref}.nii.gz -sm ${subject_pref}_diffeo_current.nii.gz |grep Similarity | awk '{print $3}'`
	echo IMAGE SIMILARITY: after previous iteration = $previous after current iteration = $current
	isminimizer=`echo "0 < ${previous} - ${current}" | bc`
else
	isminimizer=0
fi

if [ ${isminimizer} -eq 1 ]
then
	mv -f ${subject_pref}_diffeo_current.nii.gz ${subject_pref}_diffeo.nii.gz
	mv -f ${dfpref}.df.nii.gz ${subject_pref}_diffeo.df.nii.gz
	mv -f ${dfpref}.df_jac.nii.gz ${subject_pref}_diffeo.df_jac.nii.gz
	return 0
else
	rm -f ${subject_pref}_diffeo_current.nii.gz
	rm -f ${dfpref}.df.nii.gz ${dfpref}.df_jac.nii.gz
	rm -fr ${jac_mask}
	return 1
fi

}

#
# the main code
#

if [ $# -lt 6 ]
then
	echo "Deformable alignment of a DTI volume (the subject) to a DTI template"
	echo "Usage: `basename $0` template subject mask initial no_of_iter ftol"
	exit 1
else
	echo "registering $2 to $1 ..."
	echo "starting at `date`"
fi

template_pref=`getTVPrefix $1`
if [ $? -ne 0 ]
then
	echo "incompatible input: $1"
	exit 1
fi
template_basename=`basename $1`
template_base_pref=`getTVPrefix $template_basename`
 
subject_pref=`getTVPrefix $2`
if [ $? -ne 0 ]
then
	echo "incompatible input: $2"
	exit 1
fi

mask=$3
initial=$4
no_of_iter=$5
ftol=$6

if [ $initial -eq 1 ]
then
	rm -fr ${subject_pref}_diffeo.nii.gz
	cp ${subject_pref}.nii.gz ${subject_pref}_diffeo.nii.gz
fi

if [ $no_of_iter -gt 6 ]
then
	echo number of iterations cannot exceed 6
	exit 1
fi

# piecewise affine registration
iter_count=1
while [ $iter_count -le $no_of_iter ]
do
	echo
	echo iteration ${iter_count} begins ...
	pwaAtvCGM ${iter_count} ${template_pref} ${template_base_pref} ${subject_pref} ${mask} ${start[iter_count]} 6 DDS ${prior[iter_count]} ${reg[iter_count]} $ftol
	if [ $? -eq 0 ]
	then
		echo iteration ${iter_count} done
	else
		echo iteration ${iter_count} done, not accepted ...
		echo if you believe the registration terminates too early, you can relax the jacobian limit.
		echo default jacobian limits can be changed by modifying dtitk_common.sh in the scripts directory of DTI-TK.
		break
	fi
	let iter_count=iter_count+1
	echo
done

# clean up
rm -fr ${subject_pref}_jac_mask.nii.gz
rm -fr ${subject_pref}_diffeo.df_jac.nii.gz

# end
echo "ending at `date`"

