#!/usr/local/bin/bash
# preprocessing for DTI images acquired at FMRIB 
####################################################################
#This script performs registration of all the volumes
#in a given dti dataset to a reference non-diffusion-weighted image
#(scaling, translation along the PE direction and shearing)
#Brain Extraction is also performed
#with a fractional intensity treshold of 0.025
####################################################################

Usage() {
echo "REGSCRIPT - Preprocessing for DTI images acquired at FMRIB "
echo ""
echo "USAGE: regscript <username> <study_number> <output_dir> [number_acqs] [type_of_correction] [series_numb_dti] [zero_pad] [study_dir]"
echo ""
echo "default number_acqs is 1"
echo ""
echo "The FID files should normally be in"
echo "	    ~username/scratch/MRdata/study_number/"
echo ""
echo "type_of_correction (optional):"
echo "    mc - correct for eddy currents and motion (default)"
echo "         bvecs will only be rotated if number_acqs is 1"
echo "         as the amount of rotation corresponding to a given"
echo "         volume for different acquisitions would not be the same"
echo "         and it wouldn't be possible to simply average the data"
echo "    ed - correct only for eddy currents (3 degrees of freedom)"
echo ""
echo "series_numb_dti (optional)- default is 2"  
echo "                            e.g. series_2.1 - non-DW image (nodif)"
echo "                                 series_2.2 - DTI (acquisition 1)"
echo "                                 series_2.3 - DTI (acquisition 2) ..."
echo ""
echo "zero_pad (optional) - default is off"
echo "                      set it on to run load_varian"
echo "                      with the option -hkspad "
echo "                      you might wish to try this if the" 
echo "                      half k-space reconstruction has failed"
echo ""
echo "Indicate study_dir if the files are in a different directory"
echo ""
echo "OUTPUTS"
echo "-----------------------------------------------------------------------"
echo "        non-DWI: nodif"
echo "        DTI after registration:   data (if number_acqs equal to 1)"
echo "                                  data1, data2... (several repetitions)" 
echo "        Brain mask:               nodif_brain_mask"
echo "        T1-structural (optional - only if series_numb_dti is 2): struct"
echo "        T2-structural (optional - only if series_numb_dti is 2): T2-weighted"
echo "        Diffusion directions:     bvals and bvecs"                      
echo ""
exit
}

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

username=$1
study_number=$2
output_dir=`pwd`/$3

if [ $# -ge 4 ]; then
    number_acqs=$4
    echo Assuming number_acqs is $number_acqs
else 
    number_acqs=1;
    echo Assuming number_acqs is 1
fi

if [ $# -ge 5 ]; then
    type_cor=$5
else     
    type_cor='mc'
fi

if [ $# -ge 6 ]; then
   series_numb=$6
   echo series_numb_dti is $series_numb
else 
   series_numb=2;
   echo series_numb_dti is $series_numb
fi

option='-hks'
if [ $# -ge 7 ]; then
    check=$7
    check_option=`echo $check`
    case $check_option
    in
	on) option='-hkspad';;
	off) option='-hks';;
	*) option='-hks';;
    esac
fi


echo load_varian will run with option $option

if [ $# -ge 8 ]; then
   study_dir=$8
else 
   study_dir=/usr/people/${username}/scratch/MRdata
fi

nodiffid=${study_dir}/${study_number}/series_${series_numb}.1
dtifid=${study_dir}/${study_number}/series_${series_numb}.
structfid=${study_dir}/${study_number}/series_3.1
T2fid=${study_dir}/${study_number}/series_4.1

if [ ! -e $output_dir ]; then
    mkdir $output_dir
fi

cd $output_dir
echo output_dir is $output_dir

cp ~dtiuser/etc/dof3 .

#Load FID files
if [ -e ${nodiffid}.fid ]; then
    echo Loading nodif
     load_varian $nodiffid $output_dir/nodif -auto -float -buo -egrv $option
else
    echo $nodiffid.fid: FID not found
    exit;
fi

nvols_nodif=`avwval $output_dir/nodif dim4`;

if [ $nvols_nodif -gt 1 ]; then
    echo nodif has $nvols_nodif volumes - file name will be changed to data
    echo first volume will be saved as nodif
    mv nodif.hdr dti1.hdr
    mv nodif.img dti1.img
    avwsplit dti1
    mv vol0000.hdr nodif.hdr
    mv vol0000.img nodif.img
    rm vol* -f	
    
    count=1
    image_numb=2
    scan_numb=2
    
else
    count=0
    image_numb=1
    scan_numb=2
fi

while [ $count -lt $number_acqs ]; do  
    if [ -e ${dtifid}${scan_numb}.fid ]; then
       echo Loading dti${image_numb}
      load_varian ${dtifid}${scan_numb}.fid ${output_dir}/dti${image_numb} -auto -float -buo -egrv $option 
    else
       echo ${dtifid}${image_numb}.fid: file not found
       exit;    
    fi
    count=`expr $count + 1`;
    image_numb=`expr $image_numb + 1`;
    scan_numb=`expr $scan_numb + 1`;
done

scan_numb=`expr $scan_numb - 1`;

if [ $series_numb -eq 2 ]; then
    if [ -e ${structfid}.fid ]; then
	echo Loading structural
	load_varian ${structfid} ${output_dir}/struct -auto -float
    fi
    if [ -e ${T2fid}.fid ]; then
	echo Loading T2-weighted
	load_varian $T2fid ${output_dir}/T2_weighted -auto -float 
    fi
fi

#Registers each volume in dti to the reference image nodif
nvols=`avwval dti1 dim4`;
echo $nvols gradient directions
count_acqs=0

correc=`echo $type_cor`;
case $correc
  in
    mc) echo Correcting for eddy currents and motion
        if [ ! -e matrices ]; then
	   mkdir matrices
        fi  

	while [ $count_acqs -lt $number_acqs ]; do
	    tmp_acqs=`expr $count_acqs + 1`; 
	    avwsplit dti${tmp_acqs}
		if [ $nvols -lt 9 ]; then
		    count=0
		    while [ ${count} != ${nvols} ];do
			echo FLIRT on volume ${count} 
			flirt  -in vol000$count -ref nodif -out rvol000$count -nosearch -omat matrices/mat
			
			avscale matrices/mat vol000$count >matrices/tmp
			head -4 matrices/tmp | tail -3 > matrices/MAT000${count}${tmp_acqs}
			count=`expr $count + 1`;
			done
		else
		    count=0
		    while [ $count != 10 ];do
			echo FLIRT on volume ${count}
			flirt  -in vol000$count -ref nodif -out rvol000$count -nosearch -omat matrices/mat
			avscale matrices/mat vol000$count >matrices/tmp
			head -4 matrices/tmp | tail -3 > matrices/MAT000${count}${tmp_acqs}
			count=`expr $count + 1`;
		    done
		    while [ ${count} != ${nvols} ];do
			    echo FLIRT on volume ${count}
			flirt -in vol00${count} -ref nodif -out rvol00${count} -nosearch -omat matrices/mat
			avscale matrices/mat vol00$count >matrices/tmp
			head -4 matrices/tmp | tail -3 > matrices/MAT00${count}${tmp_acqs}
			rm -f matrices/mat
			rm -f matrices/tmp
			count=`expr $count + 1`; 
		    done
         		fi
	        #Clean up unecessary files
		rm vol* -f
		rm dti${tmp_acqs}* -f

		if [ $number_acqs -eq 1 ]
		then
		    echo creating file: data
		    avwmerge -t data rvol00*.hdr
		else
		    echo creating file: data${tmp_acqs}
		    avwmerge -t data${tmp_acqs} rvol00*.hdr
		fi
	    
		#Clean up unecessary files
		rm rvol* -f
	    
		count_acqs=`expr $count_acqs + 1`;
	 done
	 rm -f matrices/mat;
	 rm -f matrices/tmp;;

    ed) echo Correcting only for eddy currents
	while [ $count_acqs -lt $number_acqs ]; do
	    tmp_acqs=`expr $count_acqs + 1`; 
	    avwsplit dti${tmp_acqs}
	    if [ $nvols -lt 9 ]; then
		count=0
		while [ ${count} != ${nvols} ];do
		    echo FLIRT on volume ${count} 
		    flirt -schedule dof3 -in vol000$count -ref nodif -out rvol000$count -nosearch 
		    count=`expr $count + 1`;
		done
	    else
		count=0
		while [ $count != 10 ];do
		    echo FLIRT on volume ${count}
		    flirt -schedule dof3 -in vol000$count -ref nodif -out rvol000$count -nosearch 
		    count=`expr $count + 1`;
		done
		while [ ${count} != ${nvols} ];do
		    echo FLIRT on volume ${count}
		    flirt -schedule dof3 -in vol00${count} -ref nodif -out rvol00${count} -nosearch 
		    count=`expr $count + 1`; 
		done
	    fi

	    #Clean up unecessary files
	    rm vol* -f
	    rm dti${tmp_acqs}* -f

	    if [ $number_acqs -eq 1 ]
	    then
		echo creating file: data
		avwmerge -t data rvol00*.hdr
	    else
		echo creating file: data${tmp_acqs}
		avwmerge -t data${tmp_acqs} rvol00*.hdr
	    fi

	    #Clean up unecessary files
	    rm rvol* -f
	    count_acqs=`expr $count_acqs + 1`;
	done;;

    *) echo Type_of_correction must be either mc or ed
       exit;;
esac

#Brain Extraction
echo BE
bet nodif mask -f 0.025 -m

rm mask.hdr -f
rm mask.img -f
mv mask_mask.hdr nodif_brain_mask.hdr
mv mask_mask.img nodif_brain_mask.img
rm dof3 -f 


#Copy relevant bvals and bvecs file
n_vols=`echo $nvols`;

case $n_vols
in
    7) cp ~dtiuser/etc/FMRIB_bvals_bvecs/bvals_6 bvals;
       cp ~dtiuser/etc/FMRIB_bvals_bvecs/bvecs_6 bvecs;
       filename='gdata7';;
    13) cp ~dtiuser/etc/FMRIB_bvals_bvecs/bvals_12 bvals;
	cp ~dtiuser/etc/FMRIB_bvals_bvecs/bvecs_12 bvecs;
	filename='gdata13';;
    63) cp ~dtiuser/etc/FMRIB_bvals_bvecs/bvals_60 bvals;
	cp ~dtiuser/etc/FMRIB_bvals_bvecs/bvecs_60 bvecs;
	filename='gdata63';;
    *) echo 'bvals and bvecs not available for that # of directions';
       exit;;
esac

#The bvecs will only be rotated for 1 acquisition
#The amount of rotation corresponding to a given volume for
#different acquisitions would be different and it wouldn't be possible
#to simply average the data

if [ $number_acqs -eq 1 ]; then
    echo Opening Matlab
    case $correc
    in
	mc)
	    echo "addpath('/usr/people/dtiuser/etc/FMRIB_bvals_bvecs/')">tmp1
	    if [ $nvols_nodif -gt 1 ]; then
		echo -n "rotate_bvecs($number_acqs,$n_vols,'$nodiffid')" >  tmp2
	    else
		echo -n "rotate_bvecs($number_acqs,$n_vols,'${dtifid}${scan_numb}.fid')" > tmp2;
	    fi
	    echo  >tmp3
	    cat tmp1 tmp2 tmp3 >tmp.m
	    cat tmp.m |matlab;
	    rm tmp1 tmp2 tmp3 tmp.m -f;;
	*)  exit;;
    esac
fi

echo Done!