#! /bin/tcsh -f

#
# mri_motion_correct2
#
# Performs motion correction of anatomicals. This should produce
# the same output as mri_motion_correct
#
# Original Author: Doug Greve
#
# Copyright © 2021 The General Hospital Corporation (Boston, MA) "MGH"
#
# Terms and conditions for use, reproduction, distribution and contribution
# are found in the 'FreeSurfer Software License Agreement' contained
# in the file 'LICENSE' found in the FreeSurfer distribution, and here:
#
# https://surfer.nmr.mgh.harvard.edu/fswiki/FreeSurferSoftwareLicense
#
# Reporting: freesurfer@nmr.mgh.harvard.edu
#
#


set VERSION = 'mri_motion_correct2 dev';

set tmpdir      = ();
set inputlist   = ();
set outspec     = ();
set PrintHelp   = 0;
set LF          = ();
set CleanUp     = 1;
set CleanUpList = ();
set SaveXFM     = 1;
set target = ();
set PWD = pwd;
if ( -e /bin/pwd ) set PWD = /bin/pwd
set HiRes = '';
set outfmt = ();
set AllowWild = 0;

# If no args, print usage and exit #
if($#argv == 0) then
  goto usage_exit;
  exit 1;
endif

# Look for version #
set n = `echo $argv | grep version | wc -l`
if($n != 0) then
  echo $VERSION
  ##### gather version info used in this script
  mri_convert --version
  set prog = `which rawtominc`
  strings $prog | grep rawtominc.c
  set prog = `which minctracc`
  strings $prog | grep minctracc.c
  set prog = `which mincresample`
  strings $prog | grep mincresample.c
  set prog = `which mincaverage`
  strings $prog | grep mincaverage.c
  exit 0;
endif

source $FREESURFER_HOME/sources.csh

goto parse_args;
parse_args_return:

goto check_params;
check_params_return:

##### Create a log file ######
if($#LF == 0) then
  if("$outfmt" == "COR") then
    set LF = $outdir/mri_motion_correct2.log
  else
    set LF = $outspec.mri_motion_correct2.log
  endif
  if(-e $LF) mv $LF $LF.old
endif

echo "--------------------------------------------------------------"
echo "mri_motion_correct2 logfile is $LF"
echo "--------------------------------------------------------------"

echo "mri_motion_correct2 log file" >> $LF
echo $VERSION      >> $LF
echo $0            >> $LF
echo $argv         >> $LF
pwd               >> $LF
uname -a           >> $LF
date               >> $LF
which mri_convert  >> $LF
echo SaveXFM $SaveXFM

set StartTime = `date`;

# Get TR, TI, TE, and FlipAngle, check consist
foreach input ($inputlist)
  set TR = `mri_info $input --tr`;
  if($status) then
    echo "ERROR: $status"
    exit 1;
  endif
  set TE = `mri_info $input --te`;
  if($status) then
    echo "ERROR: $status"
    exit 1;
  endif
  set TI = `mri_info $input --ti`;
  if($status) then
    echo "ERROR: $status"
    exit 1;
  endif
  set FlipAngle = `mri_info $input --flip_angle`;
  if($status) then
    echo "ERROR: $status"
    exit 1;
  endif

  if($input == $inputlist[1]) then
    set TR0 = $TR;
    set TE0 = $TE;
    set TI0 = $TI;
    set FlipAngle0 = $FlipAngle
  endif

  if($TR != $TR0) then
    echo "WARNING: TRs are inconsistent $TR0 $TR"
  endif
  if($TE != $TE0) then
    echo "WARNING: TEs are inconsistent $TE0 $TE"
  endif
  if($TI != $TI0) then
    echo "WARNING: TIs are inconsistent $TI0 $TI"
  endif
  if($FlipAngle != $FlipAngle0) then
    echo "WARNING: FlipAngles are inconsistent $FlipAngle0 $FlipAngle"
  endif

end


#------------- Convert each input to MINC ------------#
set CleanUpList = ($CleanUpList );
set mncinputlist = ();
@ run = 0;
foreach input ($inputlist)
  @ run = $run + 1;

  echo "-----------------------------------------" |& tee -a $LF
  echo "Converting $input " |& tee -a $LF
  date |& tee -a $LF

  set mncinput = $tmpdir/cor-$run.mnc
  set mncinputlist = ($mncinputlist $mncinput);
  rm -f $mncinput

  set cmd = (mri_convert $input $mncinput)
  #set cmd = (mri_convert $input $mncinput -odt float) # should do this instead?
  pwd >> $LF
  echo $cmd >> $LF
  $cmd |& tee -a $LF
  if($status) then
    echo "ERROR: mri_convert failed" |& tee -a $LF
    exit 1;
  endif

  set CleanUpList = ($CleanUpList $mncinput);

end

#----------- Motion Correct ---------------#
set mncresampled = $tmpdir/resampled.mnc
if($#target == 0) then
  set mnctarget = $mncinputlist[1];
else
  set mnctarget = $tmpdir/target.mnc
  set cmd = (mri_convert $target $mnctarget)
  echo $cmd | tee -a $LF
  $cmd | tee -a $LF
endif
@ run = 0;
foreach mncinput ($mncinputlist)
  @ run = $run + 1;

  if($mncinput == $mnctarget) continue;

  if($SaveXFM) then
    set xfm = $outstem-mc$run.xfm
  else
    set xfm = $tmpdir/run$run.xfm
    set CleanUpList = ($CleanUpList $xfm);
  endif
  rm -f $xfm

  echo "-----------------------------------------" |& tee -a $LF
  echo "Motion Correcting $mncinput" |& tee -a $LF
  date |& tee -a $LF
  echo "xfm is $xfm" |& tee -a $LF
  # minctracc is MNI software without version option, but cvs'ed.
  set cmd = (minctracc -lsq6 $mncinput $mnctarget $xfm)
  pwd | tee -a $LF
  echo $cmd | tee -a $LF
  $cmd |& tee -a $LF
  if($status) then
    echo "ERROR: minctracc failed" |& tee -a $LF
    exit 1;
  endif

  echo "-----------------------------------------" |& tee -a $LF
  echo "Resampling $mncinput" |& tee -a $LF
  date |& tee -a $LF

  rm -f $mncresampled
  # mincresample is MNI software without version option, but cvs'ed
  set cmd = (mincresample -like $mnctarget -transform $xfm)
  set cmd = ($cmd $mncinput $mncresampled)
  pwd >> $LF
  echo $cmd >> $LF
  $cmd |& tee -a $LF
  if($status) then
    echo "ERROR: mincresample failed" |& tee -a $LF
    exit 1;
  endif

  mv $mncresampled $mncinput |& tee -a $LF

end

#----------- Average Volumes Together ---------------#
if($#mncinputlist > 1) then
  echo "-----------------------------------------" |& tee -a $LF
  echo "Averaging " |& tee -a $LF
  date |& tee -a $LF
  set avgvol = $tmpdir/avgvol.mnc
  set CleanUpList = ($CleanUpList $avgvol);
  rm -f $avgvol
  set cmd = (mincaverage $mncinputlist $avgvol);
  pwd | tee -a $LF
  echo $cmd | tee -a $LF
  $cmd |& tee -a $LF
  if($status) then
    echo "ERROR: mincaverage failed" |& tee -a $LF
    exit 1;
  endif
else
  echo "Only one input, so not averaging" |& tee -a $LF
  set avgvol = $mncinputlist 
endif

#----------- Convert Average to output ---------------#
echo "-----------------------------------------" |& tee -a $LF
echo "Converting Average to output " |& tee -a $LF
date |& tee -a $LF
set cmd = (mri_convert ${HiRes} $avgvol $outspec);
set cmd = ($cmd -tr $TR -te $TE -TI $TI -flip_angle $FlipAngle)
pwd >> $LF
echo $cmd >> $LF
$cmd |& tee -a $LF
if($status) then
  echo "ERROR: mri_convert failed" |& tee -a $LF
  exit 1;
endif

goto cleanup;
cleanup_return:

#---------------- end processing--------------------#
echo "Started at: $StartTime"    |& tee -a $LF
echo "Ended at:   `date`"        |& tee -a $LF
echo "mri_motion_correct2: done" |& tee -a $LF

exit 0;

################# subroutines #######################
#################             #######################

############--------------##################
parse_args:
set cmdline = ($argv);
while( $#argv != 0 )

  set flag = $argv[1]; shift;

  switch($flag)

    case "-help"
    case "--help"
      set PrintHelp = 1;
      goto usage_exit;
      breaksw

    case "-o":
      if ( $#argv == 0) goto arg1err;
      set outspec = $argv[1]; shift;
      breaksw

    case "-i":
      if ( $#argv == 0) goto arg1err;
      set inputlist = ($inputlist $argv[1]); shift;
      breaksw

    case "-t":
      if ( $#argv == 0) goto arg1err;
      set target = ($argv[1]); shift;
      breaksw

    case "-tmp":
    case "-tmpdir":
      if ( $#argv == 0) goto arg1err;
      set tmpdir = $argv[1]; shift;
      set CleanUp = 0;
      breaksw

    case "-umask":
      if ( $#argv == 0) goto arg1err;
      umask $argv[1]; shift;
      breaksw

    case "-nocleanup":
      set CleanUp = 0;
      breaksw

    case "-verbose":
      set verbose = 1;
      breaksw

    case "-echo":
      set echo = 1;
      breaksw

    case "-debug":
      set verbose = 1;
      set echo = 1;
      breaksw

    case "-wild":
      set AllowWild = 1;
      breaksw

    case "-cm":
      set HiRes = ('-cm')
      breaksw

    default:
      if($AllowWild) then
        set inputlist = ($inputlist $flag);
      else
        echo ERROR: Flag $flag unrecognized.
        echo $cmdline
        exit 1
      endif
      breaksw
  endsw

end

goto parse_args_return;
############--------------##################

############--------------##################
check_params:

  if($#inputlist == 0) then
    echo "ERROR: no inputs specified"
    exit 1;
  endif

  if($#inputlist == 1 && $#target == 0) then
    echo "ERROR: only one input specified"
    exit 1;
  endif

  if($#outspec == 0) then
    echo "ERROR: no output specified"
    exit 1;
  endif

  set outfmt = `mri_info $outspec --format`
  if("$outfmt" == "unknown") set outfmt = "COR";

  if("$outfmt" == "COR") then
    set outdir = $outspec
    set outstem = $outspec/COR
  else
    set outdir = `dirname $outspec`;
    set outstem = `fname2stem $outspec`
    if($status) then
      echo "$outstem"
      exit 1;
    endif
  endif
  mkdir -p $outdir
  if($status) then
    echo "ERROR: cannot create $outdir"
    exit 1;
  endif

  # Go through each input, make sure it exists
  foreach i ($inputlist)
    mri_info $i > /dev/null
    if($status) then
      echo "ERROR: with $i"
      exit 1;
    endif
  end

  # Create the tmp directory and get full path
  if($#tmpdir == 0) then
    set tmpdir = `fs_temp_dir --scratch --base $outdir`
  endif
  echo tmpdir is $tmpdir
  if($tmpdir != /tmp) then
    mkdir -p $tmpdir
    if($status) then
      echo "ERROR: cannot create $tmpdir"
      exit 1;
    endif
    pushd $tmpdir > /dev/null
    set tmpdir = `pwd`;
    popd > /dev/null
  endif

goto check_params_return;
############--------------##################

############--------------##################
arg1err:
  echo "ERROR: flag $flag requires one argument"
  exit 1
############--------------##################


############--------------##################
cleanup:

  if($CleanUp) then
    foreach f ($CleanUpList)
      rm -f $f
    end
    rm -r $tmpdir
  endif

goto cleanup_return;
############--------------##################




############--------------##################
usage_exit:
  echo "USAGE: mri_motion_correct2"
  echo ""
  echo "Required Arguments:"
  echo ""
  echo "  -o  output spec : output file or directory (for COR)"
  echo "  -i  input1 <-i input2 <-i input3>>"
  echo ""
  echo "Optional Arguments:";
  echo ""
  echo "  -t target      : use target instead of first file"
  echo "  -wild          : assume unmatched args are input files"
  echo "  -tmpdir tmpdir : directory for temporary files"
  echo "  -nocleanup     : do not delete temporary files"
  echo "  -umask umask   : set unix file permission mask"
  echo "  -cm            : COR volumes conform to min voxel size"
  echo "  -version       : print version and exit"
  echo "  -debug         : print lots of stuff to screen"
  echo ""

  if(! $PrintHelp) exit 1;

  echo $VERSION

  cat $0 | awk 'BEGIN{prt=0}{if(prt) print $0; if($1 == "BEGINHELP") prt = 1 }'

exit 1;

#---- Everything below here is printed out as part of help -----#
BEGINHELP

Aligns and averages two or more volumes. Uses minctracc with -lsq6 to
perform the alignment. Then uses mincresample and mincaverage.

The inputs and output can be any format accepted by mri_convert.

Example 1:

mri_motion_correct2 -i 002.mgz  -i 003 -o mc4.img

  002.mgz is a volume in compressed MGH format.
  003 is in COR format.
  mc4.img (the output) is in analyze format.

Example 2:

Say you have many input volumes, eg, 001.mgh ... 010.mgh, and
you do not want to list all of them on the command-line with
a -i. Then you can:

mri_motion_correct2 -o mc.mgh  -wild *.mgh

Note that -wild must appear BEFORE the wildcard.