#!/usr/bin/tcsh -f
# gca-apply

set VERSION = 'gca-apply mockbuild-local';

set gca = ();
set subject = ();
set norm = ();
set outdir = ()
set input = ();
set gcabase = ()
set GCARegIterations = ();
set DiceSeg = ();
set DiceFile = ();
set SrcLTA = ();
set SrcNorm = ();
set UseCoreg = 0;
set UseV6LabOpts = 1;
set m3zSpec = ()
set brainmask = brainmask.mgz
set EMRegAseg = 1;
set ctab = $FREESURFER_HOME/ASegStatsLUT.txt
set ForceUpdate = 0

set Overwrite = 0;
set tmpdir = ();
set cleanup = 1;
set LF = ();

if($?FS_OMP_NUM_THREADS) then
  setenv OMP_NUM_THREADS $FS_OMP_NUM_THREADS 
else
  setenv OMP_NUM_THREADS 1
endif

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

source $FREESURFER_HOME/sources.csh

goto parse_args;
parse_args_return:
goto check_params;
check_params_return:

set StartTime = `date`;
set tSecStart = `date '+%s'`;

mkdir -p $outdir/stats $outdir/mri/transforms $outdir/scripts
# outdir must exist before getfullpath
set outdir = `getfullpath $outdir`

if($#LF == 0) set LF = $outdir/scripts/gca-apply.$gcabase.log
if($LF != /dev/null && -e $LF) mv $LF $LF.bak
echo "Log file for gca-apply" >> $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

#========================================================
# Have to cd because mri_em_register writes a bunch of stuff to 
# current dir
cd $outdir/mri 

date | tee -a $LF
echo "cd `pwd`"| tee -a $LF

# mri_em_register
set LTA = $outdir/mri/transforms/$gcabase.lta
if($#SrcLTA == 0) then
  # LTA NOT passed on command line, create using mri_em_register
  set opts = (-mask $mdir/$brainmask)
  set opts = ($opts -uns 3) # used by recon-all, but not with skull
  # but docs imply it is only with skull that it should be used
  set cmd = (mri_em_register $opts $mdir/$input $gca $LTA)
  set ud = `UpdateNeeded $LTA $mdir/$brainmask $mdir/$input $gca`
else
  # LTA passed on command line
  set cmd = (cp $SrcLTA $LTA)
  set ud = `UpdateNeeded $LTA $SrcLTA`
endif
if($ud || $ForceUpdate) then
  echo "cd `pwd`" |& tee -a $LF
  echo $cmd |& tee -a $LF
  fs_time $cmd |& tee -a $LF
  if($status) exit 1;
else
  echo "$LTA update not needed" |& tee -a $LF
endif

date | tee -a $LF

# mri_ca_normlize
set norm = $outdir/mri/norm.$gcabase.mgz
if($#SrcNorm == 0) then
  set cmd = (mri_ca_normalize -mask $mdir/$brainmask $mdir/$input $gca $LTA $norm)
  set ud = `UpdateNeeded $norm $mdir/$brainmask $mdir/$input $gca $LTA `
else
  set cmd = (cp $SrcNorm $norm)
  set ud = `UpdateNeeded $norm $SrcNorm`
endif
if($ud || $ForceUpdate) then
  echo "cd `pwd`" |& tee -a $LF
  echo $cmd |& tee -a $LF
  fs_time $cmd |& tee -a $LF
  if($status) exit 1;
else
  echo "$norm update not needed" |& tee -a $LF
endif

date | tee -a $LF

# mri_ca_register
if($#m3zSpec == 0) then
  set m3z = $outdir/mri/transforms/$gcabase.m3z
  set ud = `UpdateNeeded $m3z $mdir/$brainmask $LTA $norm $gca`
  if($ud || $ForceUpdate) then
    set opts = (-align-after -mask $mdir/$brainmask -T $LTA)
    set opts = ($opts -nobigventricles) # Used by recon-all
    # These are not used by recon-all but used in rebuild_gca_atlas.csh
    #set opts = ($opts -smooth 1.0 -levels 2) 
    set cmd = (mri_ca_register $opts $norm $gca $m3z)
    if($#GCARegIterations) set cmd = ($cmd -gcareg-iters $GCARegIterations)
    echo $cmd |& tee -a $LF
    fs_time $cmd |& tee -a $LF
    if($status) exit 1;
  else
    echo "$m3z update not needed" |& tee -a $LF
  endif
endif

date | tee -a $LF

# mri_ca_label
set aseg = $outdir/mri/$gcabase.aseg.mgz
set ud = `UpdateNeeded $aseg $m3z $norm $gca`
if($ud || $ForceUpdate) then
  set opts = (-align)
  if($UseV6LabOpts) then
    set opts = ($opts -relabel_unlikely 9 .3)
    set opts = ($opts -prior 0.5)
  endif
  #set opts = ($opts -nobigventricles) # not used in recon-all, used in rebuild_gca_atlas.csh
  set cmd = (mri_ca_label $opts $norm $m3z $gca $aseg)
  echo $cmd |& tee -a $LF
  fs_time $cmd |& tee -a $LF
  if($status) exit 1;
else
  echo "$aseg update not needed" |& tee -a $LF
endif

date | tee -a $LF

# Compute segstats
set statsfile = $outdir/stats/$gcabase.stats
set ud = `UpdateNeeded $statsfile $aseg $norm`
if($ud || $ForceUpdate) then
  set cmd = (mri_segstats --seg $aseg --sum $statsfile \
    --in $norm --pv $norm --empty --excludeid 0 \
    --ctab $ctab --subject $subject \
    --brainmask $mdir/brainmask.mgz --brain-vol-from-seg \
    --in-intensity-name norm.$gcabase --in-intensity-units MR )
  set talxfm = $sdir/mri/transforms/talairach.xfm
  if(-e $talxfm ) set cmd = ($cmd --etiv --talxfm $talxfm);
  echo $cmd |& tee -a $LF
  $cmd |& tee -a $LF
  if($status) exit 1;
else
  echo "$statsfile update not needed" |& tee -a $LF
endif

date | tee -a $LF

if($#DiceSeg) then
  set cmd = (mri_compute_seg_overlap -cortex 0 -L $outdir/mri/$DiceFile \
    -table $DiceFile.table $aseg $mdir/$DiceSeg)
  echo $cmd |& tee -a $LF
  $cmd |& tee -a $LF
  if($status) exit 1;
endif

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

# Done
echo " " |& tee -a $LF
set tSecEnd = `date '+%s'`;
@ tSecRun = $tSecEnd - $tSecStart;
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 "gca-apply-Run-Time-Sec $tSecRun" |& tee -a $LF
echo "gca-apply-Run-Time-Hours $tRunHours" |& tee -a $LF
echo " " |& tee -a $LF
echo "gca-apply Done" |& tee -a $LF
exit 0

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

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

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

    case "--gca":
      if($#argv < 1) goto arg1err;
      set gca = $argv[1]; shift;
      if(! -e $gca) then
        echo "ERROR: cannot find $gca"
        exit 1;
      endif
      breaksw

    case "--gca-rb-2016":
      set gca = $FREESURFER_HOME/average/RB_all_2016-05-10.vc700.gca
      if(! -e $gca) then
        echo "ERROR: cannot find $gca"
        exit 1;
      endif
      breaksw

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

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

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

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

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

    case "--ctab":
      if($#argv < 1) goto arg1err;
      set ctab = $argv[1]; shift;
      if(! -e $ctab) then
        echo "ERROR: cannot find $ctab"
        exit 1;
      endif
      breaksw

    case "--hypothalamus":
      # For convenience, for now
      set gca = /autofs/cluster/fsm/users/greve/subjects/gca.limbic6.fsm010/gca/gca.i02.gca
      set ctab = /autofs/cluster/fsm/users/greve/subjects/gca.limbic6.fsm010/fsm.jca.2017c+aseg.lut
      set gcabase = gca.limbic6
      breaksw

    case "--hypothalamus-sym":
      # For convenience, for now
      set gca = /autofs/cluster/fsm/users/greve/subjects/gca.limbic7.fsm010.sym/gca/gca.i02.gca
      set ctab = /autofs/cluster/fsm/users/greve/subjects/gca.limbic7.fsm010.sym/fsm.jca.2017c+aseg.lut
      set gcabase = gca.limbic7.sym
      breaksw

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

    case "--no-segstats":
      set DoSegStats = 0;
      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 "--gcareg-iters":
      # For testing, to make ca_reg run faster
      if($#argv < 1) goto arg1err;
      set GCARegIterations = $argv[1]; shift;
      breaksw

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

    case "-sd":
    case "--sd":
      if ( $#argv < 1) goto arg1err;
      setenv SUBJECTS_DIR $argv[1]; shift;
      if(! -e $SUBJECTS_DIR) then
        echo "ERROR: cannot find SUBJECTS_DIR $SUBJECTS_DIR"
        exit 1;
      endif
      setenv SUBJECTS_DIR `getfullpath $SUBJECTS_DIR`
      breaksw

    case "--lta":
      if($#argv < 1) goto arg1err;
      set SrcLTA = $argv[1]; shift;
      if(! -e $SrcLTA) then
        echo "ERROR: cannot find $SrcLTA"
        exit 1;
      endif
      set SrcLTA = `getfullpath $SrcLTA`
      breaksw

    case "--norm":
      if($#argv < 1) goto arg1err;
      set SrcNorm = $argv[1]; shift;
      if(! -e $SrcNorm) then
        echo "ERROR: cannot find $SrcNorm"
        exit 1;
      endif
      set SrcNorm = `getfullpath $SrcNorm`
      breaksw

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

    case "--no-v6labopts":
      set UseV6LabOpts = 0;
      breaksw

    case "--v6labopts":
      set UseV6LabOpts = 1;
      breaksw

    case "--dice":
      if($#argv < 2) goto arg2err;
      set DiceSeg = $argv[1]; shift;
      set DiceFile = $argv[1]; shift;
      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($#gca == 0) then
  echo "ERROR: must spec gca"
  exit 1;
endif
if($#subject == 0) then
  echo "ERROR: must spec subject"
  exit 1;
endif
set sdir = $SUBJECTS_DIR/$subject
set mdir = $sdir/mri
if(! -e $sdir) then
  echo "ERROR: cannot find $subject"
  exit 1;
endif
if($#input == 0) set input = nu.mgz
foreach f ($mdir/$brainmask $mdir/$input)
  if(! -e $f) then
    echo "ERROR: cannot find $f"
    exit 1;
  endif
end
if($#outdir == 0) set outdir = $sdir

if($#DiceSeg && ! -e $mdir/$DiceSeg) then
  echo "ERROR: cannot find $mdir/$DiceSeg"
  exit 1;
endif

if($#gcabase == 0) set gcabase = `basename $gca .gca`
set lta   = $outdir/mri/transforms/$gcabase.lta
set norm  = $outdir/mri/norm.$gcabase.mgz
set aseg  = $outdir/mri/$gcabase.aseg.mgz
set stats = $outdir/stats/$gcabase.stats
if($#m3zSpec == 0) then
  set m3z = $outdir/mri/transforms/$gcabase.m3z
else
  set m3z = $outdir/transforms/$m3zSpec
  if(! -e $m3z) then
    echo "ERROR: cannot find $m3z"
    exit 1
  endif
endif

# Check whether the outputs already exist
if($ForceUpdate) then
  if(-e $lta && ! $Overwrite) then
    echo "ERROR: $lta already exists"
    exit 1;
  endif
  if(-e $norm && ! $Overwrite) then
    echo "ERROR: $norm already exists"
    exit 1;
  endif
  if(-e $aseg && ! $Overwrite) then
    echo "ERROR: $aseg already exists"
    exit 1;
  endif
  if(-e $stats && ! $Overwrite) then
    echo "ERROR: $stats already exists"
    exit 1;
  endif
  if($#m3zSpec == 0) then
    if(-e $m3z && ! $Overwrite) then
      echo "ERROR: $m3z already exists"
      exit 1;
    endif
  endif
endif

set gca = `getfullpath $gca`

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 "gca-apply : apply a gca, including em_reg, ca_norm, ca_reg, and ca_label"
  echo "  --gca gcafile"
  echo "  --s subject"
  echo ""
  echo " Other options"
  echo "  --nthreads nthreads : control number of OMP threads"
  echo "  --base gcabase : use gcabase when naming output files (default is basename gcafile)"
  echo "  --no-segstats : do not compute segstats "
  echo "  --sd SUBJECTS_DIR, or -sd "
  echo "  --dice DiceSeg DiceFile "
  echo "  --lta SrcLTA : use SrcLTA instead of running mri_em_register "
  echo "  --norm SrcNorm : use SrcNorm instead of running mri_ca_normalize "
  echo "  --input input.mgz : defalt is nu.mgz "
  echo "  --brainmask newbrainmask.mgz : defalt is brainmask.mgz "
  echo "  --o outdir : use this for output instead of SUBJECTS_DIR/subject "
  echo "  --no-v6labopts : do not include -relabel_unlikely 9 .3 -prior 0.5 with mri_ca_label":
  echo "  --m3z m3zfile":
  echo "  --gca-rb-2016 : use RB_all_2016-05-10.vc700.gca"
  echo "  --force-update : force recreation of a file even if it is younger then its parents"
  echo ""
  echo "  --gcareg-iters : set to 1, only for testing"
  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

Applies a GCA, performing the steps of mri_em_register,
mri_ca_normalize, mri_ca_register, and mri_ca_label (and possibly
mri_segstats). This script is designed to replicate the stages in
recon-all but does not overwrite any of those files. Single threaded,
this script may take 8 hours or so because mri_ca_register is done.
By default, an existing file will not be recreated if its inputs
are older than it (use --force-update to create regardless).

Creates files with the names:

gcabase.lta
gcabase.m3z
norm.gcabase.mgz
gcabase.aseg.mgz
gcabase.stats

Where gcabase is either `basename gcafile` or whatever is passed with --gcabase

When computing segstats, it uses partial volume correction with norm.gcabase.mgz.

Requires that nu.mgz (or --input instead of nu.mgz) and brainmask.mgz
(or --brainmask instead of brainmask.mgz) be present and ideally
created in the same way as when the atlas was trained on the training
subjects. When using gcatrain, the brainmask.mgz and nu.mgz are
created with:

recon-all -s subject -autorecon1 -no-talcheck

If --dice is specified, then the dice coefficient is computed using
mri_compute_seg_overlap using $SUBJECTS_DIR/$subject/mri/$DiceSeg
and the output will be put in $SUBJECTS_DIR/$subject/$DiceFile

--no-v6labopts 

do not include "-relabel_unlikely 9 .3 -prior 0.5" as part of the
mri_ca_label command options.  This are v6 options that are not part
of the 5.3 command line.

--o outputdir

If this is specified, then the output will be created in this folder.
This can be useful when you do not have permissions to the source
data or do not want to create additional files in the source tree.
outputdir will look like the subjects recon-all folder, so make sure
that the subject name is in the outputdir.



