#!/bin/bash

# Example script for registering individual T1w and DTI data to T1w and DTI multimodal templates
# The script requires ANTs and TORTOISE software packages. The registration parameters and the number of iterations (-n) could be tuned as needed. 


if [[ -z ${ANTSPATH} ]] ; then
  echo "Environment variable ANTSPATH must be defined"
  exit 1
fi
if [[ ! -f "${ANTSPATH}/ANTS" ]] ; then
  echo "Cannot find the ANTS program. Please \(re\)define \$ANTSPATH in your environment."
  exit 1
fi


if [ $# -lt 10 ]
then
    echo 'Usage: Register T1w and DTI data from a single individual to T1w and DTI multimodal templates'
    echo 'Example: bash reg_to_MIITRA_T1_and_DTI.sh -t1 <individual_T1w_image> -dti <individual_DTI_data> -i <Type_of_transformation_from_DTI_to_T1w_of_the_individual> -t  -T1 <T1w_template> -DT <DTI_template> -n <Number_of_iterations> -o <Output_folder> -m <Boolean_for_saving_intermediate_files>'
    echo 'OPTIONS:'
    echo '   -t1: T1w data of an individual (required, in the format used by ANTs)       '
    echo '   -dt: DTI data of the same individual (required, in the format used by TORTOISE)   '
    echo '   -i: Type of transformation from DTI to T1w of the individual (required)'   
    echo '       1: The transformation will be provided by user via -t option (default)'
    echo '       2: Do rigid-body registration of Trace map of the DTI to T1w image'
    echo '       3: Do affine registration of Trace map of the DTI to T1w image'
    echo '   -t: Initial transformation from DTI to T1w of the individual (required only if -i=1, one linear or nonlinear file in the format used by ANTs; provide an identity matrix if no transform needed)'
    echo '   -T1: T1w template  (required, in the format used by ANTs, e.g. MIITRA v1.5 T1w brain only)'
    echo '   -DT: DTI template (required, in the format used by TORTOISE, e.g. MIITRA v1.5 DTI)'
    echo '   -n: Number of iterations alternating between modalities (required, default 3)'
    echo '   -o: Output folder (required)'
    echo '   -m: Boolean for saving intermediate files of all iterations (optional)'
    echo '       0: Only final transformed T1w and DT images and overall transformations are saved in the output folder (default)'
    echo '       1: Intermediate files and transformations of all iterations are saved in the registration_dir_T1filename folder'

    exit
fi

#################################initializing variables 
type_trans=1
num_iterations=3 
save_intermediate=0

#################################reading command line arguments
while getopts "t1:dt:i:t:T1:DT:n:o:m:" OPT
do
  case $OPT in
      t1) # individual T1w image
    raw_t1=$OPTARG
   ;;
      dt) # individual DT image
    raw_dti=$OPTARG
   ;;
      i) # Type of DTI to T1w alignment (default = 1)
    type_trans=$OPTARG
    if [ ${type_trans} -lt 1 ] || [ ${type_trans} -gt 3 ]
    then
        echo "Type of transformation could be (1=provided via -t, 2=rigid-body, 3=affine), you passed ${type_trans}. Exiting..."
        exit 1
    fi
   ;;
      t) # Initial transformation from raw DTI to raw T1w provided by user
    initial_trans=$OPTARG
   ;;
      T1) # T1w template
    t1_template=$OPTARG
   ;;
      DT) # DTI template
    dti_template=$OPTARG
   ;;
      n) # Number of iterations (default = 3)
    num_iterations=$OPTARG
    if [ ${num_iterations} -le 1 ] 
    then
        echo "Number of iterations should be larger than 1, you passed ${num_iterations}. Exiting..."
        exit 1
    fi
   ;;
      o) # Output folder
   output_folder=$OPTARG
   if [ ! -e "${output_folder}" ]
   then
       mkdir ${output_folder} || {
           echo "Could not create output directory. Exiting..." 1>&2
           exit 1
       }
   fi
   ;;
      m) # Boolean for saving intermediate files (default = 0)
   save_intermediate=$OPTARG
   if [ ${save_intermediate} -ne 1 ] || [ ${save_intermediate} -ne 0 ]
   then
       echo "Boolean for saving intermediate files should be (1=save, 0=not save), you passed ${save_intermediate}. Exiting..."
       exit 1
   fi
   ;;
      \?) # getopts issues an error message
   echo "Unknown option found. Exiting..."
   exit 1
      ;;
  esac
done



if [ ! -e "${raw_t1}" ] || [ ! -e "${raw_dti}" ] || [ ! -e "${t1_template}" ] || [ ! -e "${dti_template}" ]
then
    echo  "Error in input files. Exiting...."
    exit 1 
fi






#################################Create working directory
filename=`basename ${raw_t1}`
filename=${filename%.nii*}

tmpdir=${output_folder}/registration_dir_${filename}
if [ -d ${tmpdir} ] ; then
    echo "${tmpdir} already exists. Exiting..."
    exit 1
fi

mkdir ${tmpdir} || {
    echo "Could not create temporary directory. Exiting..." 1>&2
    exit 1
}


################################Compute Trace map of DTI 
cp ${raw_dti} ${tmpdir}/
gunzip ${tmpdir}/*
filename=`basename ${raw_dti}`
filename=${filename%.nii*}
ComputeTRMap ${tmpdir}/${filename}.nii 
gzip ${tmpdir}/${filename}_TR.nii
mv ${tmpdir}/${filename}_TR.nii.gz ${tmpdir}/dti_scalar.nii.gz




################################Calculate initial alignment
if [ ${type_trans} -eq 1 ] 
then
    if [ ! -e "${initial_trans}" ]
    then
        echo  "Error in initial transformation provided (-t). Exiting...."
        exit 1
    fi


elif [ ${type_trans} -eq 2 ] 
then
    ANTS 3 -m MI[${raw_t1},${tmpdir}/dti_scalar.nii.gz,1,32] -o ${tmpdir}/Initial_dti_to_t1_ -i 0 --use-Histogram-Matching --number-of-affine-iterations 10000x10000x10000x10000x10000 —MI-option 64x64000 --do-rigid true
    mv ${tmpdir}/Initial_dti_to_t1_Affine.txt ${tmpdir}/Initial_dti_to_t1_Rigid.txt
    initial_trans=${tmpdir}/Initial_dti_to_t1_Rigid.txt


elif [ ${type_trans} -eq 3 ] 
then
    ANTS 3 -m MI[${raw_t1},${tmpdir}/dti_scalar.nii.gz,1,32] -o ${tmpdir}/Initial_dti_to_t1_ -i 0 --use-Histogram-Matching --number-of-affine-iterations 10000x10000x10000x10000x10000 —MI-option 64x64000 --rigid-affine true
    initial_trans=${tmpdir}/Initial_dti_to_t1_Affine.txt


else
    echo "Type of transformation should be (1=provided via -t, 2=rigid-body, 3=affine), you passed ${type_trans}. Exiting..."
    exit 1
fi





################################ Update
echo"------------------------------------------------------------"
echo Raw T1w image:                            ${raw_t1}
echo Raw DTI image:                            ${raw_dti}
echo Initial alignment from DTI to T1:         ${initial_trans}
echo T1w template:                             ${t1_template}
echo DTI template:                             ${dti_template}
echo Number of iterations:                     ${num_iterations}
echo Output folder:                            ${output_folder}
echo Boolean for saving intermediate:          ${save_intermediate}
echo"------------------------------------------------------------"



################################ Parameters in T1w registration using ANTs
RADIUS=4
TRANSFORMATION=SyN[0.25]
REGULARIZATION=Gauss[3,0.5] 
MAXITERATIONS=100x100x100
LINEARSET="--use-Histogram-Matching --number-of-affine-iterations 10000x10000x10000x10000x10000 --MI-option 64x64000 --affine-gradient-descent-option 0.1x0.5x1.e-4x1.e-4"





################################ Parameters in DTI registration using DR-TAMAS
mtrans="-t SyN[0.3,2,0.1]"
shrink="-f 6x4x2x1"
smooth="-s 1.0x1.0x0.5x0.05"
iter_num="-c 100x100x100x100"
metric="--metric DTITK --metric Trace"
opt_rot="--optimization_reorientation FS"
final_rot="--final_reorientation FS"
perform_affine="--perform_affine 0"
DFsmoothing="--adaptiveDFsmoothing 0"
combinefield="--output_combined_disp_field 0"







################################ Obtain DTI template
if [[ ${dti_template} == *.gz ]]
then
    cp ${dti_template} ${tmpdir}/dti_template.nii.gz
    dti_template=${tmpdir}/dti_template.nii.gz
elif [[ ${dti_template} == *.nii ]]; then
    cp ${dti_template} ${tmpdir}/dti_template.nii
    dti_template=${tmpdir}/dti_template.nii
else
    echo "Failed to obtain DTI template. Exiting..."
    exit 1
fi

################################ Create itentity matrix
cat > ${tmpdir}/identity_matrix.txt <<EOF
#Insight Transform File V1.0
#Transform 0
Transform: AffineTransform_double_3_3
Parameters: 1 0 0 0 1 0 0 0 1 0 0 0 
FixedParameters: 0 0 0

EOF

################################ Preparation for transformations
transform[0]="-t ${initial_trans}"

for i in {1..19..2}  
do
    let iter=(i+1)/2
    transform[$i]="-t ${tmpdir}/T1_trans_iter${iter}_Warp.nii.gz -t ${tmpdir}/T1_trans_iter${iter}_Affine.txt"
done

for i in {2..20..2}  
do 
    let iter=i/2
    transform[$i]="-t ${tmpdir}/DT_trans_iter${iter}_Warp.nii.gz"
done




################################ Registration starts here
count=1
steps=0
while [ ${count} -le ${num_iterations} ]
do

    let previous_count=count-1

    if [ ${count} -eq 1 ]
    then
        T1_update=${raw_t1}
    else
        T1_update=${tmpdir}/T1_iter${previous_count}_step2.nii.gz
    fi 

    ##### T1w registration
    echo"\n------------------------------------------------------------"
    echo "Iteration ${count}, Step 1:"
    echo"------------------------------------------------------------\n"

    let steps=steps+1
    def_prefix=${tmpdir}/T1_trans_iter${count}_
    echo "ANTS 3 -m CC[${t1_template},${T1_update},1,${RADIUS}] -t ${TRANSFORMATION} -r ${REGULARIZATION} -i ${MAXITERATIONS} ${LINEARSET} -o ${def_prefix}"
    ANTS 3 -m CC[${t1_template},${T1_update},1,${RADIUS}] -t ${TRANSFORMATION} -r ${REGULARIZATION} -i ${MAXITERATIONS} ${LINEARSET} -o ${def_prefix}


    ##### Apply T1w transformations to DTI
    transformations=""
    for ((i=${steps}; i>=0; i--))
    do 
        transformations="${transformations} ${transform[$i]}"
    done
    antsApplyTransforms -d 3 --verbose 1 -i ${tmpdir}/dti_scalar.nii.gz -o [${tmpdir}/dti_overalltrans_iter${count}_step1.nii,1] ${transformations} -r ${t1_template}
    ApplyTransformationToTensor ${raw_dti} ${tmpdir}/dti_overalltrans_iter${count}_step1.nii ${tmpdir}/dti_iter${count}_step1.nii FS ${dti_template}


    ##### DTI registration
    echo"\n------------------------------------------------------------"
    echo "Iteration ${count}, Step 2:"
    echo"------------------------------------------------------------\n"

    let steps=steps+1
    echo "DTIREG --fixed_tensor ${dti_template} --moving_tensor ${tmpdir}/dti_iter${count}_step1.nii --initial-affine-transform ${tmpdir}/identity_matrix.txt ${mtrans} ${shrink} ${smooth} ${iter_num} ${metric} ${opt_rot} ${final_rot} ${perform_affine} ${DFsmoothing} ${combinefield}"
    DTIREG --fixed_tensor ${dti_template} --moving_tensor ${tmpdir}/dti_iter${count}_step1.nii --initial-affine-transform ${tmpdir}/identity_matrix.txt ${mtrans} ${shrink} ${smooth} ${iter_num} ${metric} ${opt_rot} ${final_rot} ${perform_affine} ${DFsmoothing} ${combinefield} 
    mv ${tmpdir}/dti_iter${count}_step1_deffield_MINV.nii ${tmpdir}/DT_trans_iter${count}_Warp.nii.gz


    ##### Apply DTI transformations to T1w
    transformations=""
    for ((i=${steps}; i>=1; i--))
    do 
        transformations="${transformations} ${transform[$i]}"
    done
    antsApplyTransforms -d 3 --verbose 1 -i ${raw_t1} -o ${tmpdir}/T1_iter${count}_step2.nii.gz ${transformations} -r ${t1_template}



    let count=count+1
done

echo"\n------------------------------------------------------------"
echo "*** Registrations completed ***"
echo"------------------------------------------------------------\n"


################################ Combine all transformations and generate the final T1w
let final_step1=steps-1
transformations=""
for ((i=${final_step1}; i>=1; i--))
do 
    transformations="${transformations} ${transform[$i]}"
done

filename=`basename ${raw_t1}`
filename=${filename%.nii*}

antsApplyTransforms -d 3 --verbose 1 -i ${raw_t1} -o [${output_folder}/${filename}_T1w_overall_transformations.nii.gz,1] ${transformations} -r ${t1_template}
antsApplyTransforms -d 3 --verbose 1 -i ${raw_t1} -o ${output_folder}/${filename}_T1w_deformed.nii.gz ${transformations} -r ${t1_template}



################################ Combine all transformations and generate the final DTI
transformations=""
for ((i=${steps}; i>=0; i--))
do 
    transformations="${transformations} ${transform[$i]}"
done

filename=`basename ${raw_dti}`
filename=${filename%.nii*}

antsApplyTransforms -d 3 --verbose 1 -i ${tmpdir}/dti_scalar.nii.gz -o [${output_folder}/${filename}_DTI_overall_transformations.nii,1] ${transformations} -r ${t1_template}
ApplyTransformationToTensor ${raw_dti} ${output_folder}/${filename}_DTI_overall_transformations.nii ${output_folder}/${filename}_DTI_deformed.nii.gz FS ${dti_template}


################################ Cleanup


if [ ${save_intermediate} -eq 0 ] 
then
    rm -r ${tmpdir}
fi