#!/bin/sh

Usage() { 
    echo "" 
    echo "Usage: dti_bedtrack.sh  <project> <name> <working dir> <network dir>"
    echo "" 
    echo "You must have FSL installed in order to run this script"
    echo ""
    exit 1
} 

[ "$3" = "" ] && Usage
debug=1
aseg_here=0
prj=$1
shift
pollo=$1
shift
w_dir=$1
shift
net_dir=$1
shift
td=${w_dir}'/.tmp_'${pollo}
bd=${td}'/bedpostx'
list=${w_dir}/../masks.list
seed_file=${w_dir}/../dti_track.seed
targets_file=${w_dir}/../dti_track.targets
if [ ! -d "$bd" ]; then
mkdir $bd 
fi 
echo "Copying files"
${FSLDIR}/bin/imcp ${w_dir}/${pollo}_dti_data ${bd}/data
${FSLDIR}/bin/imcp ${w_dir}/${pollo}_dti_brain_mask ${bd}/nodif_brain_mask
cp ${td}/bvecs ${bd}/bvecs
cp ${td}/bvals ${bd}/bvals
echo "Making bedpostx"
echo [`date`]
${FSLDIR}/bin/bedpostx_gpu ${bd}
echo "So far, so good"
echo [`date`]
###########################################
echo "Getting nodes and making masks"
if [ -e "${seed_file}" ]; then
	count=0;
	for seed in `awk NF ${seed_file}`;do
		if [[ ${seed} =~ ^[0-9]+$ ]]; then
			if [ ${aseg_here} = 0 ]; then
				${PIPEDIR}/bin/get_aseg.sh ${prj} ${pollo} ${w_dir};
				aseg_here=1
			fi;
			WarpImageMultiTransform 3 ${w_dir}/${pollo}_aseg.nii.gz ${td}/${pollo}_aseg_warped.nii.gz -R ${td}/hifi_b0.nii.gz ${td}/${pollo}_dti_ants_elast_t1_b0Warp.nii.gz ${td}/${pollo}_dti_ants_elast_t1_b0Affine.txt;
			${FSLDIR}/bin/fslmaths ${td}/${pollo}_aseg_warped -uthr ${seed} -thr ${seed} -div ${seed} ${td}/${pollo}_mask_${seed};
			if [ "${count}" -gt "0" ]; then
				${FSLDIR}/bin/fslmaths ${td}/${pollo}_mask_seed -add ${td}/${pollo}_mask_${seed} ${td}/${pollo}_mask_tmp;
				${FSLDIR}/bin/imcp ${td}/${pollo}_mask_tmp ${td}/${pollo}_mask_seed;
			else
				${FSLDIR}/bin/imcp ${td}/${pollo}_mask_${seed} ${td}/${pollo}_mask_seed;
			fi
			((count++))
		else
			x=`find ${net_dir} -name "${seed}"`;
			node=`${FSLDIR}/bin/remove_ext $(basename ${x})`;
			WarpImageMultiTransform 3 ${x} ${td}/${node}_warped.nii.gz -R ${w_dir}/${pollo}_t1_reoriented.nii.gz ${td}/${pollo}_dti_ants_elast_mni_t1Warp.nii.gz ${td}/${pollo}_dti_ants_elast_mni_t1Affine.txt;
			WarpImageMultiTransform 3 ${td}/${node}_warped.nii.gz ${td}/${pollo}_${node}.nii.gz -R ${td}/hifi_b0.nii.gz ${td}/${pollo}_dti_ants_elast_t1_b0Warp.nii.gz ${td}/${pollo}_dti_ants_elast_t1_b0Azffine.txt;
			if [ $count ]; then
				${FSLDIR}/bin/fslmaths ${td}/${pollo}_mask_seed -add ${td}/${pollo}_${node} ${td}/${pollo}_mask_tmp;
				${FSLDIR}/bin/imcp ${td}/${pollo}_mask_tmp ${td}/${pollo}_mask_seed;
			else
				${FSLDIR}/bin/imcp ${td}/${pollo}_${node} ${td}/${pollo}_mask_seed;
			fi
			((count++))
		fi;
	done;
	echo "${td}/${pollo}_mask_seed.nii.gz" > ${td}/${pollo}_seed.list;
else
	exit 1 "No seed file!";
fi
if [ -e "${targets_file}" ]; then
	for target in `awk NF ${targets_file}`; do
		if [[ ${target} =~ ^[0-9]+$ ]]; then
			if [ ! ${aseg_here} ]; then
				${PIPEDIR}/bin/get_aseg.sh ${prj} ${pollo} ${w_dir};
				aseg_here=1
			fi;
			WarpImageMultiTransform 3 ${w_dir}/${pollo}_aseg.nii.gz ${td}/${pollo}_aseg_warped.nii.gz -R ${td}/hifi_b0.nii.gz ${td}/${pollo}_dti_ants_elast_t1_b0Warp.nii.gz ${td}/${pollo}_dti_ants_elast_t1_b0Affine.txt;
			${FSLDIR}/bin/fslmaths ${td}/${pollo}_aseg_warped -uthr ${target} -thr ${target} -div ${target} ${td}/${pollo}_mask_${target};
			echo "${td}/${pollo}_mask_${target}.nii.gz" >> ${td}/${pollo}_targets.list;
		else
			x=`find ${net_dir} -name "${target}"`;
			node=$(basename ${x});
			WarpImageMultiTransform 3 ${x} ${td}/${node%.nii.gz}_warped.nii.gz -R ${w_dir}/${pollo}_t1_reoriented.nii.gz ${td}/${pollo}_dti_ants_elast_mni_t1Warp.nii.gz ${td}/${pollo}_dti_ants_elast_mni_t1Affine.txt;
			WarpImageMultiTransform 3 ${td}/${node%.nii.gz}_warped.nii.gz ${td}/${pollo}_${node}.gz -R ${td}/hifi_b0.nii.gz ${td}/${pollo}_dti_ants_elast_t1_b0Warp.nii.gz ${td}/${pollo}_dti_ants_elast_t1_b0Affine.txt;
			echo "${td}/${pollo}_${node}.gz" >> ${td}/${pollo}_targets.list;
		fi;
	done;
else
	exit 1 "No targets file!";
fi
###########################################		
echo "Doing probtrackx"
probtrackx2 --opd --os2t --forcedir -s ${bd}.bedpostX/merged -m ${w_dir}/${pollo}_dti_brain_mask -x ${td}/${pollo}_seed.list --targetmasks=${td}/${pollo}_targets.list --dir=${td}/probtrack_out
rm ${td}/${pollo}_seed.list;
rm ${td}/${pollo}_targets.list;
echo "Done"
echo [`date`]