#!/bin/tcsh -f
# fsr-coreg 

#  In theory, it would be better to resample all modes into an
#    in-between space rather than to a reference.

#  Need a way to better specify args for internal commands (eg, rob
#    reg, and coreg)
#

if(-e $FREESURFER_HOME/sources.csh) then
  source $FREESURFER_HOME/sources.csh
endif

set VERSION = 'fsr-coreg 8.2.0';
set selfname = `basename $0`

set importdir = ();
set outdir = ()
set refmode = ();
set threads = 1
set cleanup = 1;
set LF = ();
set tmpdir = ()
set XOptsFile = ()
set GlobXOptsFile = ()
set ForceUpdate = 0
set rtsat = 4.685

if($?FS_V8_XOPTS == 0) setenv FS_V8_XOPTS 0
set UseV8 = $FS_V8_XOPTS;

set inputargs = ($argv);
set PrintHelp = 0;
if($#argv == 0) goto usage_exit;
set n = `echo $argv | grep -e -help | wc -l` 
if($n != 0) then
  set PrintHelp = 1;
  goto usage_exit;
endif
set n = `echo $argv | grep -e -version | wc -l` 
if($n != 0) then
  echo $VERSION
  exit 0;
endif
goto parse_args;
parse_args_return:
goto check_params;
check_params_return:

set V8XoptsFile = ()
if($UseV8) set V8XoptsFile = $FREESURFER/etc/global-expert-options.v8.txt

set StartTime = `date`;
set tSecStart = `date '+%s'`;
set year  = `date +%Y`
set month = `date +%m`
set day   = `date +%d`
set hour   = `date +%H`
set min    = `date +%M`

mkdir -p $outdir/log

# Set up log file
if($#LF == 0) set LF = $outdir/log/fsr-coreg.Y$year.M$month.D$day.H$hour.M$min.log
if($LF != /dev/null) rm -f $LF
echo "Log file for fsr-coreg" >> $LF
date  | tee -a $LF
echo "" | tee -a $LF
echo "setenv SUBJECTS_DIR $SUBJECTS_DIR" | tee -a $LF
echo "cd `pwd`"  | tee -a $LF
echo $0 $inputargs | tee -a $LF
echo "" | tee -a $LF
cat $FREESURFER_HOME/build-stamp.txt | tee -a $LF
echo $VERSION | tee -a $LF
uname -a  | tee -a $LF
if($?PBS_JOBID) then
  echo "pbsjob $PBS_JOBID"  >> $LF
endif

#========================================================
date | tee -a $LF

# Create the files to keep track of the inputs and params
if(! -e $parfile) then
  echo refmode $refmode >> $parfile
  if($#XOptsFile) cp $XOptsFile $outdir/log/fsr-coreg.expert.txt
endif

# Go through each mode and do within mode registration
echo "\nWithin-mode registration" | tee -a $LF
foreach modename ($modenames)
  echo "$modename `date` -------------------------"
  set modedir = $outdir/$modename
  mkdir -p $modedir

  set nruns = `cat $importdir/$modename/nruns.txt`
  # If only one run, then just create a symlink
  if($nruns == 1) then
    if("$outdir" == "$importdir") then
      pushd $modedir >& /dev/null
      rm -f runavg.mgz
      ln -s run001.mgz runavg.mgz
      popd >& /dev/null
    else
      rm -f $modedir/runavg.mgz
      ln -s $importdir/$modename/run001.mgz $modedir/runavg.mgz
    endif
    echo "$modename, only one run, creating symlink" | tee -a $LF
    continue
  endif

  # Multiple runs, register and average together
  set srclist = ()
  set ltalist = ()
  @ nthrun = 0;
  while ($nthrun < $nruns)
    @ nthrun = $nthrun + 1
    set nthrunstr = `printf %03d $nthrun`
    set srcfile = $importdir/$modename/run$nthrunstr.mgz
    set ltafile = $outdir/$modename/reg$nthrunstr.lta
    if(! -e $srcfile) then
      echo "ERROR: cannot find $srcfile"
      exit 1;
    endif
    set srclist = ($srclist $srcfile)
    set ltalist = ($ltalist $ltafile)
  end

  set average = $importdir/$modename/runavg.mgz

  # Check whether anything needs to be done
  set UN = `UpdateNeeded $average $srclist`
  if($UN == 0 && $ForceUpdate == 0) then
    echo "$modename, within mode registration update not needed" | tee -a $LF
    continue
  endif

  # If it gets here, then register refmode runs together
  set cmd = (mri_robust_template --template $average --mov $srclist --lta $ltalist --sat $rtsat)
  set xopts = `fsr-getxopts $selfname-mri_robust_template $V8XoptsFile $GlobXOptsFile $XOptsFile `;
  set cmd = ($cmd $xopts)
  date | tee -a $LF
  echo "cd `pwd`" | tee -a $LF
  echo $cmd |& tee -a $LF
  fs_time $cmd |& tee -a $LF
  if($status) goto error_exit

end

# Start a file for collecting tkregister commands
if($#modenames > 1) then
  set rc = $outdir/log/registration-check-commands
  rm -f $rc
  echo "# These are commands that can be used to check the cross-modal registration" | tee -a $rc
  echo "" | tee -a $rc
endif

#---------------------------------------------------------------------------
echo "\nBetween-mode registration" | tee -a $LF
set refvol = $outdir/$refmode/runavg.mgz
foreach modename ($modenames)
  echo "$modename `date` -------------------------" | tee -a $LF

  if($modename == $refmode) then
    # This is the reference mode, just create a link
    pushd $outdir/$modename >& /dev/null
    rm -f runavg-refmodespace.mgz
    ln -s runavg.mgz runavg-refmodespace.mgz
    # Create an LTA too?
    popd >& /dev/null
    continue
  endif

  # Before generating a new LTA, check whether the LTA exists
  # and whether it is different than the auto LTA. If the LTA
  # does not exist or is the same as the auto LTA, then it is
  # safe to copy any newly generated auto LTA to the LTA
  set ltaauto = $outdir/$modename/reg.avg-to-refmodespace.auto.lta
  set lta = $outdir/$modename/reg.avg-to-refmodespace.lta
  set LTAsAreDiff = 1;
  if(! -e $lta) set LTAsAreDiff = 0;
  if(-e $lta && -e $ltaauto) then
    # Check whether there is a difference
    set LTAsAreDiff = `diff $lta $ltaauto | wc -l`
  endif
  # LTAsAreDiff will == 1 if lta does not exist or is the same as ltaauto

  # Run mri_coreg to create the auto lta
  set movvol = $outdir/$modename/runavg.mgz
  set UN = `UpdateNeeded $ltaauto $movvol`
  if($UN == 1 || $ForceUpdate == 1) then
    set cmd = (mri_coreg --mov $movvol --ref $refvol --reg $ltaauto --threads $threads)
    set xopts = `fsr-getxopts $selfname-mri_coreg $V8XoptsFile $GlobXOptsFile $XOptsFile `;
    set cmd = ($cmd $xopts)
    date | tee -a $LF
    echo "cd `pwd`" | tee -a $LF
    echo $cmd |& tee -a $LF
    fs_time $cmd |& tee -a $LF
    if($status) goto error_exit
    # Create the final LTA. If it does not exist or is identical, then just copy 
    if($LTAsAreDiff == 0) then
      set cmd = (cp -p $ltaauto $lta)
      echo $cmd |& tee -a $LF
      $cmd |& tee -a $LF
      if($status) goto error_exit
    endif
  else
    echo "$modename, between mode registration update not needed" | tee -a $LF
  endif

  # Add a command to check the reg to the registration-check-commands file
  set cmd = (tkregisterfv --mov $movvol --targ $refvol --reg $lta)
  echo $cmd | tee -a $rc
  echo "" | tee -a $rc

  set outvollist = ()
  @ nthrun = 0;
  while ($nthrun < $nruns)
    @ nthrun = $nthrun + 1
    set nthrunstr = `printf %03d $nthrun`
    if($nruns > 1) then
      set runreflta = $outdir/$modename/reg$nthrunstr.refmodespace.lta
      set runlta = $outdir/$modename/reg$nthrunstr.lta
      set UN = `UpdateNeeded $runreflta $runlta $lta`
      if($UN == 1 || $ForceUpdate) then
        set cmd = (mri_concatenate_lta $runlta $lta $runreflta)
        set xopts = `fsr-getxopts $selfname-mri_concatenate_lta $V8XoptsFile $GlobXOptsFile $XOptsFile `;
        set cmd = ($cmd $xopts)
        date | tee -a $LF
        echo "cd `pwd`" | tee -a $LF
        echo $cmd |& tee -a $LF
        fs_time $cmd |& tee -a $LF
        if($status) goto error_exit
      endif
    else
      set runreflta = $outdir/$modename/reg.avg-to-refmodespace.lta
    endif

    # Resample run to reference mode space (use cubic)
    set movvol = $outdir/$modename/run$nthrunstr.mgz
    set outvol = $outdir/$modename/run$nthrunstr-refmodespace.mgz
    set UN = `UpdateNeeded $outvol $runreflta $movvol`
    if($UN == 1 || $ForceUpdate == 1) then
      set cmd = (mri_vol2vol --mov $movvol --reg $runreflta --o $outvol --interp cubic)
      set xopts = `fsr-getxopts $selfname-mri_vol2vol $V8XoptsFile $GlobXOptsFile $XOptsFile `;
      set cmd = ($cmd $xopts)
      date | tee -a $LF
      echo "cd `pwd`" | tee -a $LF
      echo $cmd |& tee -a $LF
      fs_time $cmd |& tee -a $LF
      if($status) goto error_exit
    endif
    set outvollist = ($outvollist $outvol)
  end # run

  # Now average together. Should not have to worry about scale here
  set outvol = $outdir/$modename/runavg-refmodespace.mgz
  set UN = `UpdateNeeded $outvol $outvollist`
  if($UN == 1 || $ForceUpdate) then
    set cmd = (mri_concat $outvollist --mean --o $outvol)
    date | tee -a $LF
    echo "cd `pwd`" | tee -a $LF
    echo $cmd |& tee -a $LF
    fs_time $cmd |& tee -a $LF
    if($status) goto error_exit
  endif

end #mode

date | tee -a $LF
echo "" | tee -a $LF

if($#modenames > 1) then
  # Create one freeview command to look at all the resampled modes together
  set cmd = (freeview)
  foreach modename ($modenames)
    set outvol = $outdir/$modename/$modename.mgz
    set cmd = ($cmd $outvol)
  end
  echo $cmd >> $outdir/log/registration-check-commands
  echo " ----------------------------------------- " | tee -a $LF
  cat $outdir/log/registration-check-commands | tee -a $LF
  echo " ----------------------------------------- " | tee -a $LF
  echo ""  | tee -a $LF
endif

# For fun, create a fsr command line (but don't run it)
# Make sure that refmode is the first input because fsr
# uses that for registration to the atlas (or not)
set outvol = $outdir/$refmode/$refmode.mgz
set cmd = (fsr --i $outvol)
foreach modename ($modenames)
  if($modename == $refmode) continue;
  set outvol = $outdir/$modename/$modename.mgz
  set cmd = ($cmd --i $outvol)
end
set cmd = ($cmd --bin --threads $threads)
echo "To run fsr (add --o outputdir): "| tee -a $LF
echo $cmd | tee -a $LF
echo ""| tee -a $LF

#========================================================

# Cleanup
# if($cleanup) rm -rf $tmpdir

# Done
set tSecEnd = `date '+%s'`;
@ tSecRun = $tSecEnd - $tSecStart;
set tRunMin = `echo $tSecRun/50|bc -l`
set tRunMin = `printf %5.2f $tRunMin`
set tRunHours = `echo $tSecRun/3600|bc -l`
set tRunHours = `printf %5.2f $tRunHours`
echo "Started at $StartTime " |& tee -a $LF
echo "Ended   at `date`" |& tee -a $LF
echo "Fsr-coreg-Run-Time-Sec $tSecRun" |& tee -a $LF
echo "Fsr-coreg-Run-Time-Min $tRunMin" |& tee -a $LF
echo "Fsr-coreg-Run-Time-Hours $tRunHours" |& tee -a $LF
echo " " |& tee -a $LF
echo "fsr-coreg Done" |& tee -a $LF
exit 0

###############################################

############--------------##################
error_exit:
echo "ERROR:"

exit 1;
###############################################

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

  set flag = $argv[1]; shift;
  
  switch($flag)

    case "-v8": 
    case "--v8": 
      set UseV8 = 1
      breaksw;
    case "-no-v8": 
    case "--no-v8": 
      set UseV8 = 0
      breaksw;

    case "--importdir":
    case "--id":
      if($#argv < 1) goto arg1err;
      set importdir = $argv[1]; shift;
      if(! -e $importdir) then
        echo "ERROR: $importdir does not exist, run fsr-import"
        exit 1
      endif
      breaksw

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

    case "--ref":
    case "--refmode":
      if($#argv < 1) goto arg1err;
      set refmode = $argv[1]; shift;
      breaksw

    case "--nthreads":
    case "--threads":
      if($#argv < 1) goto arg1err;
      set threads = $argv[1]; shift;
      breaksw

    case "--force-update":
    case "--force":
      set ForceUpdate = 1;
      breaksw

    case "--expert":
      if( $#argv < 1) goto arg1err;
      set XOptsFile = $argv[1]; shift;
      fsr-checkxopts $XOptsFile
      if($status) goto error_exit;
      set XOptsFile = `getfullpath $XOptsFile`
      breaksw

    case "--log":
      if($#argv < 1) goto arg1err;
      set LF = $argv[1]; shift;
      breaksw

    case "--nolog":
    case "--no-log":
      set LF = /dev/null
      breaksw

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

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

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

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

    default:
      echo ERROR: Flag $flag unrecognized. 
      echo $cmdline
      exit 1
      breaksw
  endsw

end

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

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

if($#importdir == 0) then
  echo "ERROR: must spec importdir"
  exit 1;
endif
if(! -e $importdir) then
  echo "ERROR: $importdir does not exist"
  exit 1
endif

if($#outdir == 0) set outdir = $importdir
set importdir = `getfullpath $importdir`
set outdir    = `getfullpath $outdir`

set modenamefile  = $importdir/log/fsr-import.modenames.txt
set umodenamefile = $importdir/log/fsr-import.unique.modenames.txt
foreach f ($modenamefile $umodenamefile)
  if(! -e $f) then
    echo "ERROR: cannot find $f"
    exit 1;
  endif
end
set modenames = (`cat $modenamefile`)
set nmodes = $#modenames;
set umodenames = (`cat $umodenamefile`)

set parfile = $outdir/log/fsr-coreg.par.txt
if(-e $parfile) then
  if($#refmode || $#XOptsFile) then
    echo "ERROR: output already exists, don't spec input parameters"
    exit 1
  endif
  set refmode = `cat $parfile | awk '{if($1=="refmode") print $2}'`
  set xo = $outdir/log/fsr-coreg.expert.txt
  if(-e $xo) set XOptsFile = $xo # wont be there if not using it
else
  if($#umodenames == 1 && $#refmode == 0) set refmode = $umodenames
  if($#refmode == 0) then
    echo "ERROR: must spec a reference mode with --ref"
    echo " Choices for this input are $umodenames"
    exit 1;
  endif
  # Make sure reference mode is in the list of references
  set ok = 0;
  foreach modename ($modenames)
    if("$refmode" == "$modename") set ok = 1;
  end
  if(! $ok) then
    echo "ERROR: reference mode $refmode is not a mode"
    echo "Available modes: $modenames"
    exit 1;
  endif
endif


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

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

############--------------##################
usage_exit:
  echo ""
  echo "fsr-coreg : co-registers input data in prep for FreeSurfer analysis"
  echo "  --i importdir : data dir created by fsr-import"
  echo "  --ref modename : mode to use as a reference (all modes register to this mode)"
  echo "  --threads nthreads"
  echo "  --force-update : force update of files regardless of time stamp"
  echo "  --o outdir : set the output dir (default is importdir)"
  echo "  --expert xoptsfile"
  echo "  --v8/--no-v8 ($UseV8)"
  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

fsr-coreg : co-registers input data in prep for FreeSurfer analysis

The input folder is that as created by fsr-import, which creates a folder of modalities

For example:

fsr-import --t1w run1.T1.slice0.dicom  --t1w run2.T1.slice0.dicom --t2w T2.nii.gz \
  --mode pd myPDimage.mgz --flair FLAIR-weighted.dcm --o importdir

will import 5 volumes, two of which are T1-weighted, one is T2
weighted, one is FLAIR weighted, the other is a custom mode called "pd".

The output will be a directory structure stored in import, eg,
import/pd/r001.mgz
import/t1w/r001.mgz and r002.mgz
import/t2w/r001.mgz

When fsr-coreg is run, it registers all the with-in modality runs
together (eg, r001.mgz and r002.mgz) using mri_robust_template to
create a file called mode.native.mgz. If there is only one run, then
mode.native.mgz is a symbolic link to r001.mgz. Next, each mode is
registered (with mri_coreg) to the reference mode to create an LTA
file (mode.reg-to-ref.lta). The native volume is then resampled with
cubic interpolation into the space of the reference mode to create a
file called mode.mgz (for the reference mode, this is just a link to
mode.native.mgz). In this way, a volume is produced for each mode in
alignment with the reference mode. These can be input to samseg (which
is automatically done when running samseg with
--mode/--t1w/--t2w/--flair).

t1w, t2w, and flair are given special consideration because they
can all be used as input to recon-all.