#!/usr/bin/tcsh -f
# gtmseg

set VERSION = 'gtmseg mockbuild-local';

set outdir = ();
set subject = ();
set USF = 2;
set OutputUSF = ();
set SubSegWM = 0
set KeepCC = 0
set KeepHypo = 0;
set dmax = 5;
set outvol = gtmseg.mgz # relative to subject/mri
set headseg = apas+head.mgz # created by xcerebralseg
set ctxannot = ()
set wmannot = ()
set ctxlhbase=0;
set ctxrhbase=0;
set wmlhbase=0;
set wmrhbase=0;
set DoSegStats = 0;
set ctab = ();
set SubSegCblumWM = 0;
set AddPons = 1;
set AddVermis = 1;
set DoXCerSeg = ();
set UseSamseg = 1
set threads = 1
set xcthresh = ()
set nErodeWM = 0
set MergeSeg = ()
set MergeSegIds = ()
set MergeSegCtab = ()
set ForceUpdate = 0

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

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'`;
set year  = `date +%Y`
set month = `date +%m`
set day   = `date +%d`
set hour   = `date +%H`
set min    = `date +%M`

set outdir = $SUBJECTS_DIR/$subject/mri
if($#tmpdir == 0) then
  set tmpdir = `fs_temp_dir --scratch`
endif
mkdir -p $tmpdir

# Set up log file
mkdir -p $SUBJECTS_DIR/$subject/scripts/log/
if($#LF == 0) then
  set LF = $SUBJECTS_DIR/$subject/scripts/log/$outvol.log
  if(-e $LF) mv $LF $SUBJECTS_DIR/$subject/scripts/log/$outvol.$$.log
endif
if($LF != /dev/null) rm -f $LF
echo "Log file for gtmseg" >> $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($DoXCerSeg) then
  # This will create a seg of the whole head, including skull, xCSF,
  # etc (named apas+head.mgz by default) in native anat space. It is
  # actually a merging of extracerebral segs with the whole brain
  # segmentation to create a whole head. The xcerseg and the
  # subcortical GM seg will not be modified beyond this so subcortical
  # seg must be right. The subcortical WM may be modified if
  # subsegmenting WM. Note that any external segs must be merged with
  # this headseg and then passed with --head. The cortical GM can be
  # changed by specifying a different annot.  The aparc is ok here
  # because it can be replaced by another annot below, but it must be
  # a seg whose cortex has been fixed with the ribbon eg via
  # mri_surf2volseg because mri_gtmseg will not fix it.
  set cmd = (xcerebralseg --s $subject --threads $threads --o $headseg)
  if(! $cleanup) set cmd = ($cmd --tmpdir $tmpdir/xcerebralseg)
  if($AddPons == 0)   set cmd = ($cmd --no-pons)
  if($AddVermis == 0) set cmd = ($cmd --no-vermis)
  if(! $UseSamseg) set cmd = ($cmd --no-samseg)
  if($#xcthresh) set cmd = ($cmd --thresh $xcthresh)
  if($ForceUpdate) set cmd = ($cmd --force)
  echo $cmd | tee -a $LF
  $cmd | tee -a $LF
  if($status) goto error_exit;
endif

# Add cblumwmgyri; overwrites $heaseg. Should be done in xcerebralseg,
# but not really used now
if($SubSegCblumWM) then
  set base = `basename $headseg .mgz`
  set headseg2 = $base+cblumwmgyri.mgz
  set cmd = (cblumwmgyri --s $subject --seg $headseg --n 2 \
    --o $headseg2 --no-segstats --no-log)
  echo $cmd | tee -a $LF
  $cmd | tee -a $LF
  if($status) goto error_exit;
  # Rename it
  mv $SUBJECTS_DIR/$subject/mri/$headseg2 $SUBJECTS_DIR/$subject/mri/$headseg |& tee -a $LF
  if($status) exit 1;
endif

# Create a lobs annotation. Why? 
if($SubSegWM && $#wmannot == 0) then
  foreach hemi (lh rh)
    set cmd = (mri_annotation2label --subject $subject --hemi $hemi --lobesStrict lobes)
    echo $cmd | tee -a $LF
    $cmd | tee -a $LF
    if($status) goto error_exit;
  end
endif

set ud = `UpdateNeeded $SUBJECTS_DIR/$subject/mri/$outvol $ctxannotlist $wmannotlist $SUBJECTS_DIR/$subject/mri/$headseg $MergeSeg`
if($ud || $ForceUpdate) then
  set cmd = (mri_gtmseg --s $subject --usf $USF --o $outvol --threads $threads)
  set cmd = ($cmd --apas $headseg --wm-erode $nErodeWM)
  if($#ctxannot)  set cmd = ($cmd --ctx-annot $ctxannot $ctxlhbase $ctxrhbase)
  if($#wmannot)   set cmd = ($cmd --wm-annot  $wmannot  $wmlhbase  $wmrhbase)
  if($SubSegWM)   set cmd = ($cmd --subseg-wm --dmax $dmax)
  if(! $SubSegWM) set cmd = ($cmd --no-subseg-wm) # why needed?
  if($KeepCC)     set cmd = ($cmd --keep-cc)
  if(! $KeepCC)   set cmd = ($cmd --no-keep-cc)
  if($KeepHypo)     set cmd = ($cmd --keep-hypo)
  if(! $KeepHypo)   set cmd = ($cmd --no-keep-hypo)
  if($#OutputUSF)   set cmd = ($cmd --output-usf $OutputUSF)
  if($#ctab)        set cmd = ($cmd --ctab $ctab)
  if($#MergeSeg)    set cmd = ($cmd --merge $MergeSeg $MergeSegIds)
  if($#MergeSegCtab) set cmd = ($cmd --ctab $MergeSegCtab)
  echo $cmd | tee -a $LF
  $cmd | tee -a $LF
  if($status) goto error_exit;
endif

if($DoSegStats) then
  set gtmsegstem = `fname2stem $outvol`
  set gtmctab = $SUBJECTS_DIR/$subject/mri/$gtmsegstem.ctab
  set stat = $SUBJECTS_DIR/$subject/stats/$gtmsegstem.stats
  set gtmseg = $SUBJECTS_DIR/$subject/mri/$outvol
  set ud = `UpdateNeeded $stat $gtmseg`
  if($ud || $ForceUpdate) then
    set cmd = (mri_segstats --seg $gtmseg --ctab $gtmctab --sum $stat \
      --excludeid 0 --etiv --subject $subject)
    echo $cmd | tee -a $LF
    $cmd | tee -a $LF
    if($status) goto error_exit;
  endif
endif


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

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

set stem = `fname2stem $outvol`
set ctab = $stem.ctab

# 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 "Gtmseg-Run-Time-Sec $tSecRun" |& tee -a $LF
echo "Gtmseg-Run-Time-Hours $tRunHours" |& tee -a $LF
echo " " |& tee -a $LF
echo "To check run:" | tee -a $LF
echo "tkmeditfv $subject nu.mgz -seg $outvol $ctab -opacity 1" | tee -a $LF
echo ""| tee -a $LF

echo "gtmseg Done" |& tee -a $LF
exit 0

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

############--------------##################
error_exit:
echo "ERROR: $cmd" |& tee -a $LF
echo "gtmseg exited with errors" |& tee -a $LF

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

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

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

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

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

    case "--xcerseg":
      set headseg = apas+head.mgz
      set DoXCerSeg = 1;
      breaksw

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

    case "--no-xcerseg":
      set DoXCerSeg = 0;
      breaksw

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

    case "--no-pons":
      set AddPons = 0;
      breaksw

    case "--no-vermis":
      set AddVermis = 0;
      breaksw

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

    case "--merge":
      if($#argv < 2) goto arg2err;
      set MergeSeg = $argv[1]; shift;
      set MergeSegIds = ()
      foreach a ($argv)
        set b = `echo $a | cut -c1`
        if("$b" == "-") break;  
        set MergeSegIds = ($MergeSegIds $a)
        shift
      end
      if($#MergeSegIds == 0) then
        echo "ERROR: could not find any seg ids with --merge"
        echo "  proper usage is --merge segfilename segid1 ..."
        exit 1
      endif
      breaksw

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

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

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

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

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

    case "--ctx-annot":
      if($#argv < 3) goto arg3err;
      set ctxannot = $argv[1]; shift;
      set ctxlhbase = $argv[1]; shift;
      set ctxrhbase = $argv[1]; shift;
      breaksw

    case "--wm-annot":
      if($#argv < 3) goto arg3err;
      set wmannot = $argv[1]; shift;
      set wmlhbase = $argv[1]; shift;
      set wmrhbase = $argv[1]; shift;
      set SubSegWM = 1;
      breaksw

    case "--no-subsegwm":
      set SubSegWM = 0;
      set wmannot = ();
      breaksw
    case "--wmsubseg":
    case "--wm-subseg":
    case "--subsegwm":
    case "--subseg-wm":
      set SubSegWM = 1;
      breaksw

    case "--keep-hypo":
    case "--no-label-hypo":
      set KeepHypo = 1;
      breaksw
    case "--no-keep-hypo":
    case "--label-hypo":
      set KeepHypo = 0;
      breaksw

    case "--keep-cc":
    case "--no-label-cc":
      set KeepCC = 1;
      breaksw
    case "--label-cc":
      set KeepCC = 0;
      breaksw

    case "--subseg-cblum-wm":
      set SubSegCblumWM = 1;
      breaksw
    case "--no-subseg-cblum-wm":
      set SubSegCblumWM = 0;
      breaksw

    case "--seg-stats":
      set DoSegStats = 1;
      breaksw
    case "--no-seg-stats":
      set DoSegStats = 0;
      breaksw

    case "--samseg":
      set UseSamseg = 1;
      breaksw
    case "--no-samseg":
      set UseSamseg = 0;
      breaksw

    case "--threads":
    case "--nthreads":
      set threads = $argv[1]; shift;
      breaksw

    case "--force":
     set ForceUpdate = 1
     breaksw
    case "--no-force":
     set ForceUpdate = 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 "--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($#subject == 0) then
  echo "ERROR: must spec subject"
  exit 1;
endif
if(! -e $SUBJECTS_DIR/$subject) then
  echo "ERROR: cannot find $subject"
  exit 1;
endif
set ctxannotlist = ()
if($#ctxannot) then
  foreach hemi (lh rh)
    set f = $SUBJECTS_DIR/$subject/label/$hemi.$ctxannot
    if(-e $f) continue
    set f = $SUBJECTS_DIR/$subject/label/$hemi.$ctxannot.annot
    if(-e $f) then
      set ctxannot = $ctxannot.annot
      set ctxannotlist = ($ctxannotlist $ctxannot)
      continue
    endif
    echo "ERROR: cannot find $SUBJECTS_DIR/$subject/label/$hemi.$ctxannot"
    exit 1;
  end
endif
set wmannotlist = ()
if($#wmannot) then
  foreach hemi (lh rh)
    set f = $SUBJECTS_DIR/$subject/label/$hemi.$wmannot
    if(-e $f) continue
    set f = $SUBJECTS_DIR/$subject/label/$hemi.$wmannot.annot
    if(-e $f) then
      set wmannot = $wmannot.annot
      set wmannotlist = ($wmannotlist $wmannot)
      continue
    endif
    echo "ERROR: cannot find $SUBJECTS_DIR/$subject/label/$hemi.$wmannot"
    exit 1;
  end
endif

if($UseSamseg) then
  set headseg = apas+head.samseg.mgz # created by xcerebralseg
  if($#outvol == 0) set outvol = gtmseg.samseg.mgz
endif

if(! -e $SUBJECTS_DIR/$subject/mri/$headseg) then
  # Headseg is not there
  if($#DoXCerSeg == 0) then
    # No instruction as to whether create it or not, so create it by default
    set DoXCerSeg = 1;
  endif
else
  # Headseg is there
  if($ForceUpdate) set DoXCerSeg = 1;
  if($#DoXCerSeg == 0) then
    # No instruction as to whether to overwrite it or not
    echo "ERROR: $SUBJECTS_DIR/$subject/mri/$headseg exists. This is ok"
    echo "but you must indicate whether to use what is there (--no-xcerseg)"
    echo "or create a new one and overwrite what is there (--xcerseg)"
    echo "or specify your own headseg (--head)"
    echo ""
    exit 1;
  endif
endif

if($DoXCerSeg == 0 && ! -e $SUBJECTS_DIR/$subject/mri/$headseg) then
  echo "ERROR: $SUBJECTS_DIR/$subject/mri/$headseg does not exist"
  exit 1;
endif

if($#MergeSeg) then
  set f = $SUBJECTS_DIR/$subject/mri/$MergeSeg
  if(! -e $f) then
    echo "ERROR: cannot find $f"
    exit 1
  endif
endif

if($#MergeSegCtab && $#ctab) then
  echo "ERROR: cannot have both --ctab and merge-ctab"
  exit 1
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
############--------------##################
arg3err:
  echo "ERROR: flag $flag requires three arguments"
  exit 1
############--------------##################

############--------------##################
usage_exit:
  echo ""
  echo "gtmseg -- create an anatomical segmentation for the geometric transfer matrix (GTM)"
  echo "   run with --help for more help"
  echo ""
  echo " --s subject     : subject to analyze (required)"
  echo " --xcerseg : run xcerebralseg on this subject to create $headseg"
  echo " --samseg : use samseg to create whole head seg rather than GCA"
  echo " --threads nthreads : only applies when running xcerseg"
  echo ""
  echo " --o outvol      : relative to subject/mri (default is $outvol)" 
  echo " --usf USF       : upsampling factor (default is $USF)"
  echo " --subsegwm      : subsegment WM into lobes (default)"
  echo " --keep-hypo     : do not relabel hypointensities as WM when subsegmenting WM"
  echo " --keep-cc       : do not relabel corpus callosum as WM"
  echo " --dmax dmax     : distance threshold to use when subsegmenting WM (default is $dmax)"
  echo " --ctx-annot annot lhbase rhbase : annotation to use for cortical segmentation (default is aparc 1000 2000)"
  echo " --wm-annot  annot lhbase rhbase : annotation to use for WM segmentation (with --subsegwm, default is lobes 3200 4200)"
  echo " --output-usf OutputUSF : set output USF different than USF, mostly for debugging"
  echo " --head headseg : use headseg instead of $headseg"
  echo " --subseg-cblum-wm : subsegment cerebellum WM into core and gyri":
  echo " --no-pons :   do not add pons segmentation when doing ---xcerseg":
  echo " --no-vermis : do not add vermis segmentation when doing ---xcerseg":
  echo " --wm-erode N : erode WM (2,41) by N, replacing with 5001 and 5002 (Left/Right-Shell-Cerebral-White-Matter)":
  echo " --merge segname segid1 ... : merge given segids in subject/mri/segname"
  echo " --merge-ctab ctab : colortable to use with merge if needed (must have tissue type)"
  echo " --ctab colortable "
  echo " --seg-stats : compute segmentation stats"
  echo " --xcthresh xcthresh : mri_seghead thresh"
  echo " --force : force rerunning and overwriting results"
  echo ""
  echo " --debug "
  echo " --help "


  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

Creates a segmentation that can then be used with the geometric
transfer matrix for performing PET partial volume correction. It will
create a segmentation of the head (approximate extra-cerebral CSF,
approximate skull, air cavities, and the rest of the head) using the
xcerebralseg script; this also adds segments for pons and vermis by
default; With --subseg-cblum-wm, the cerebellum WM is subsegmented
into a core component and gyri. The core components keep the
{Left,Right}-Cerebellum-White-Matter names and indices. The gyri will
become CbmWM_Gyri_{Left,Right} with indices 690 and 691.

The WM can be subsegmented into lobes with --subsegwm. Hypointensities
can be kept with --keep-hypo. The corpus callosum can be kept with
--keep-cc. If these segments are not "kept" then they are relabelled
as WM. The upsampling factor (USF) controls the resolution of the
segmentation; this make a difference for cortical
segmentations. Higher resolution is more accurate but takes up more
disk space and will require more computations later on.

Example:

gtmseg --s subject

Creates a file called $SUBJECTS_DIR/subject/mri/gtmseg.mgz without
subsegmenting WM. Hypos are re-labeled as WM as is CC. The USF is 2
(meaning .5mm voxel size).

gtmseg --s subject --keep-hypo --subsegwm --o gtmseg.wmseg.hypo.mgz --usf 1

Creates a file called $SUBJECTS_DIR/subject/mri/gtmseg.wmseg.hypo.mgz
with subsegmenting of WM and keeping yypos. CC is re-labeled as
WM. The USF is 1 (meaning 1mm voxel size).

You can use your own segmentation or a modified FS segmentation. 
It will be easiest if you modify apas+head.mgz to insert your
segmentations. apas+head.mgz is created by gtmseg but you can
create it with 

xcerebralseg --s $subject --o $SUBJECTS_DIR/$subject/mri/apas+head.mgz

=====================================================================
Using a custom subcortical segementation: there are two methods:

------------------------------------------------------------------------
In the first method, create a segmentation to merge into the
aparc+aseg.mgz. It must be in the mri folder and the same size as
aparc+aseg.mgz. Eg, if it is called myseg.mgz, then run 

gtmseg --s subject --merge myseg.mgz 801 802 --merge-ctab myseg.ctab --o gtmseg+myseg ...

This will find places where myseg is 801 or 802 and transfer 801 and
802 to the gtmseg. The color table (ctab) must be one that includes a
tissue type schema (see below). You must include a ctab unless FS already
knows internally what the tissue type is (unlikely for custom segs). 

You must manage the merge segmentation numbers (eg, 801 and 802). They
should not collide with numbers in the aparc+aseg.mgz (there is no
error check for this). 

------------------------------------------------------------------------
In the second method, run gtmseg as default to create the
apas+head.mgz file.  Create a new/custom subcortical segmentation by
replacing voxels in apas+head.mgz with your segmentation: make sure
that the index you give your new segmentation(s) is not already
present in apas+head.mgz. Next create a color table text file that
looks something like this:

# TissueTypeSchema 
#ctTType   0  unknown                           0   0   0    0
#ctTType   1  cortex                          205  62  78    0
#ctTType   2  subcort_gm                      230 148  34    0
#ctTType   3  wm                                0 255   0    0
#ctTType   4  csf                             120  18 134    0
#ctTType   5  head                            150 150 200    0
265 MySeg            219 100 176 0 2

In the above case, it is assumed that the index that you used
was 265 and that your segmentation is for subcortical gray matter
(thus making the last number=2 in the MySeg row). If you have more
than one segmentation you are adding, then add rows into the 
color table. The first 5 entries codes for the different possible
tissue types.

You can then run

gtmseg --s subject --head apas+head+myseg.mgz \
   --ctab myseg.colortable.txt --o gtmseg+myseg.mgz  

where apas+head+myseg.mgz is apas+head.mgz with your new segmentation,
myseg.colortable.txt is the color table you created, and
gtmseg+myseg.mgz will be your output volume. Note that an output color
table gtmseg+myseg.ctab will be generated. You should check in this to
make sure that your segmentation appears. You can also run

tkmeditfv subject orig.mgz -seg gtmseg+myseg.mgz  myseg.colortable.txt

To use, eg, aparc.a2009s, then
  --ctx-annot aparc.a2009s.annot 11100 12101
  The 11100 and 12101 will make it agree with  $FREESURFER/FreeSurferColorLUT.txt


===========================================================
Using a custom cortical parcellation:

If you have your own cortical annotation, then just specify
  --ctx-annot your.annot 1000 2000 
Where your cortical parcellation is in subject/label/?h.your.annot

The values 1000 and 2000 are used to create the color table and
segmentation indices for the output volume. In this case, 1000 will be
added to the left hemi and 2000 will be added to the right hemi.  The
actual values (1000, 2000) are semi-arbitrary. There needs to be
enough of a difference so that the indices for lh will not overlap
with those of the rh. They need to be 1000 or over so that the
software recognizes them as cortex (needed in the colortable to
specify the tissue type).


