#!/bin/bash
# Author: Michael Hanke <michael.hanke@gmail.com>
# License: MIT

set -e
set -u

if [ $# -lt 3 ]; then
  cat << EOT
This script implements MCFLIRT-based motion correction using an dedicated
target image. The difference to using MCFLIRT's --reffile option is that
an initial alignment to the target image can be performed with more
flexibility. The actual motion correction is performed using the mean
volume as a reference and a constant subsequent transformation to the
final target is applied in addition. Each volume is transformed separately
with a single interpolation/reslicing step that applies the combined
transformations.

Call this script with at least three arguments

 1. Input bold time series image
 2. Skull-stripped reference image
 3. Output filename for the motion-corrected time series image.

(4) Options to FLIRT for the alignment to the target image. Default: "-dof 6"

Examples:

 $0 bold.nii.gz subjtmpl.nii.gz bold_mc.nii.gz "-dof 6 -cost normcorr"

 $0 bold.nii.gz subjtmpl.nii.gz bold_mc.nii.gz "-dof 6 -cost normcorr"

EOT
  exit 1
fi
# fake test for FSL
ignore=$FSLDIR

inbold=$(readlink -f $1) || true
ref=$(readlink -f $2) || true
outbold=$(readlink -m $3)
refsmooth=${4:-0}
initflirtopts=${5:--dof 6}

if [ -z "$inbold" ]; then
  echo "File not found '$1'"
  exit 1
fi
if [ -z "$ref" ]; then
  echo "File not found '$2'"
  exit 1
fi

outbolddir=$(dirname "${outbold}")
mkdir -p "$outbolddir"

wdir=$(mktemp -d --suffix=2stagemc)
trap "rm -rf $wdir" SIGINT SIGTERM
cd "$wdir"


# motion correction reference
$FSLDIR/bin/fslmaths "$inbold" -Tmean bold -odt input
# perform motion correction
$FSLDIR/bin/mcflirt -in "$inbold" -reffile bold -mats -plots -out mc
# mean bold volume after MC
$FSLDIR/bin/fslmaths mc -Tmean meanbold_mc -odt input
$FSLDIR/bin/bet meanbold_mc bold_brain
# smooth reference volume somewhat
$FSLDIR/bin/fslmaths "$ref" -s $(echo "scale=4; $refsmooth / 2.355" | bc) refsmooth
$FSLDIR/bin/flirt ${initflirtopts} -in bold_brain -ref refsmooth -omat bold2brain.mat
cp mc.par "$($FSLDIR/bin/remove_ext "$outbold")_mcparams.txt"
for i in mc.mat/MAT*; do convert_xfm -omat mc.mat/tmpl_$(basename $i) -concat bold2brain.mat $i; done
mkdir bold
$FSLDIR/bin/fslsplit "$inbold" bold/
mkdir mc
for i in bold/*.nii*; do
  bname=$(basename $i)
  idx=${bname%.nii*}
  $FSLDIR/bin/flirt -in "$i" -ref "$ref" -applyxfm -init mc.mat/tmpl_MAT_${idx} -out mc/${idx} -interp sinc; done
$FSLDIR/bin/fslmerge -tr "$outbold" mc/* $(fslval "$inbold" pixdim4)

cd -
rm -rf "$wdir"