#!/usr/bin/env tcsh
# scgm-mask - sources
if(-e $FREESURFER_HOME/sources.csh) then
  source $FREESURFER_HOME/sources.csh
endif

set VERSION = '$Id$';
set scriptname = `basename $0`

set outdir = ();
set subjectlist = ();
set ForceUpdate = 0
set segvol = aseg.mgz

# These are the non-GM indices
set idx = ( 0 2 3 4 5 14 15 24 31 41 42 43 44 63 72 77 78 79 251 252 253 254 255 )
set asegctab = $FREESURFER/ASegStatsLUT.txt
set fsctab = $FREESURFER_HOME/FreeSurferColorLUT.txt

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
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`

mkdir -p $outdir/log $outdir/subjects $outdir/count
pushd $outdir > /dev/null
set outdir = `pwd`;
popd > /dev/null

if($#tmpdir == 0) then
  if(-dw /scratch)   set tmpdir = /scratch/tmpdir.scgm-mask.$$
  if(! -dw /scratch) set tmpdir = $outdir/tmpdir.scgm-mask.$$
endif
#mkdir -p $tmpdir

# Set up log file
if($#LF == 0) set LF = $outdir/log/scgm-mask.Y$year.M$month.D$day.H$hour.M$min.log
if($LF != /dev/null) rm -f $LF
echo "Log file for scgm-mask" >> $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
ls -l $0  | 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
echo "pid $$" | tee -a $LF
if($?PBS_JOBID) then
  echo "pbsjob $PBS_JOBID"  >> $LF
endif
if($?SLURM_JOB_ID) then
  echo SLURM_JOB_ID $SLURM_JOB_ID >> $LF
endif

#========================================================
# Resample the segs to mni152 space via the synthmorph nonlinear reg
set segmnilist = ()
@ n = 0
foreach subject ($subjectlist)
  @ n = $n + 1
  set seg  = $SUBJECTS_DIR/$subject/mri/$segvol
  set warp = $SUBJECTS_DIR/$subject/mri/transforms/synthmorph.1.0mm.1.0mm/warp.to.mni152.1.0mm.1.0mm.nii.gz
  set segmni = $outdir/subjects/$subject.seg.nii.gz
  set ud = `UpdateNeeded $segmni $seg $warp`
  if($ud || $ForceUpdate) then
    echo "$n/$#subjectlist $subject `date` =======================" | tee -a $LF
    set cmd = (mri_convert -rt nearest $seg -at $warp $segmni)
    echo $cmd | tee -a $LF
    $cmd |& tee -a $LF
    if($status) goto error_exit
  endif
  set segmnilist = ($segmnilist $segmni)
end

# Create a stack of the segmentations in mni152 space
set segstack = $outdir/all.seg.nii.gz
set ud = `UpdateNeeded $segstack $segmnilist`
if($ud || $ForceUpdate) then
  set cmd = (mri_concat $segmnilist --o $segstack)
  echo $cmd | tee -a $LF
  $cmd |& tee -a $LF
  if($status) goto error_exit
endif

# Binarize (0,1) the stack of segmentations.
set segstackbin = $outdir/all.seg.bin.nii.gz
set ud = `UpdateNeeded $segstackbin $segstack`
if($ud || $ForceUpdate) then
  set cmd = (mri_binarize --match-ctab $asegctab $idx --i $segstack --o $segstackbin)
  echo $cmd | tee -a $LF
  $cmd |& tee -a $LF
  if($status) goto error_exit
endif

# Create a probably map of the likelihood of a voxel being SCGM in the
# given subject list by computing the mean of the bins.
set pscgm = $outdir/p.scgm.nii.gz
set ud = `UpdateNeeded $pscgm $segstackbin`
if($ud || $ForceUpdate) then
  set cmd = (mri_concat $segstackbin --mean --o $pscgm)
  echo $cmd | tee -a $LF
  $cmd |& tee -a $LF
  if($status) goto error_exit
endif

# Create a list the number of voxels at several thresholds
# to help choose a threshold. Not sure how helpful it is.
set nvoxfile = $outdir/nvox.thresh.dat
set ud = `UpdateNeeded $nvoxfile $pscgm`
if($ud || $ForceUpdate) then
  rm -f $nvoxfile
  set th100list = (10 20 30 40 50 60 70 80 90)
  foreach th100 ($th100list)
    set th = `echo "$th100/100" | bc -l`
    set countfile = $outdir/count/count$th100.dat
    set ud = `UpdateNeeded $countfile $pscgm`
    if($ud || $ForceUpdate) then
      set cmd = (mri_binarize --i $pscgm --min $th --count $countfile)
      echo $cmd | tee -a $LF
      $cmd |& tee -a $LF
      if($status) goto error_exit
      echo "" | tee -a $LF
    endif
    set count = `cat $countfile | awk '{print $1}'`
    echo $th $count | tee -a $nvoxfile
  end
endif
echo "" | tee -a $LF
cat $nvoxfile | tee -a $LF
echo "" | tee -a $LF

# Create a mask at a 50% threshold, then "close" it by dilating by 1
# and then eroding by 1. This is a reasonable mask to use.
set mask = $outdir/scgm.mask.th50.close.nii.gz
set ud = `UpdateNeeded $mask $pscgm`
if($ud || $ForceUpdate) then
  set countfile = $outdir/pscgm.mask.th50.close.count.dat
  set cmd = (mri_binarize --i $pscgm --min 0.5 --dilate 1 --erode 1 --o $mask \
    --count $countfile --binval 559 --ctab $fsctab)
  echo $cmd | tee -a $LF
  $cmd |& tee -a $LF
  if($status) goto error_exit
endif

# This just creates a freeview command line to make it easy to view
set mni = $FREESURFER/average/mni_icbm152_nlin_asym_09c/mni_icbm152_t1_tal_nlin_asym_09c.nii.gz
set fvcmd = (fsvglrun freeview -neuro-view -rotate-around-cursor --hide-3d-slices \
  -v $mni ${pscgm}:colormap=heat:heatscale=.001,1 -v ${mask}:outline=1)
echo "" | tee -a $LF
echo $fvcmd | tee -a $LF
echo "" | tee -a $LF


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

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

# Done
echo " " |& tee -a $LF
set tSecEnd = `date '+%s'`;
@ tSecRun = $tSecEnd - $tSecStart;
set tRunMin = `echo $tSecRun/60|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 "Scgm-Mask-Run-Time-Sec $tSecRun" |& tee -a $LF
echo "Scgm-Mask-Run-Time-Min $tRunMin" |& tee -a $LF
echo "Scgm-Mask-Run-Time-Hours $tRunHours" |& tee -a $LF
echo " " |& tee -a $LF
echo "scgm-mask 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 "--o":
      if($#argv < 1) goto arg1err;
      set outdir = $argv[1]; shift;
      breaksw

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

    case "--f":
      if($#argv < 1) goto arg1err;
      set subjectlistfile = $argv[1]; shift;
      set subjectlist = ($subjectlist `cat $subjectlistfile`)
      breaksw

    case "--sd":
      if($#argv < 1) goto arg1err;
      setenv SUBJECTS_DIR $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($#outdir == 0) then
  echo "ERROR: must spec outdir"
  exit 1;
endif
if($#subjectlist == 0) then
  echo "ERROR: must spec subjects"
  exit 1;
endif
foreach subject ($subjectlist)
  set seg  = $SUBJECTS_DIR/$subject/mri/$segvol
  set warp = $SUBJECTS_DIR/$subject/mri/transforms/synthmorph.1.0mm.1.0mm/warp.to.mni152.1.0mm.1.0mm.nii.gz
  foreach f ($seg $warp)
    if(-e $f) continue
    echo "ERROR: cannot find $f"
    exit 1;
  end
end

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 "scgm-mask - create a subcortical gray matter mask in mni152 space"
  echo " --f slist"
  echo " --s subject <--s subject>"
  echo " --o outdir"
  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

This script creates a subcortical gray matter (SCGM) mask in mni152
space.  It loops through all the given subjects, resamples the
aseg.mgz to mni152 1mm space (mni_icbm152_t1_tal_nlin_asym_09c) using
the synthmorph non-linear warp created during recon-all with v8,
concatenates the asegs into a single file (a "stack"), then binarizes
the subcortical GM structures. These binary maps are then averaged to
gether to give a probability (0-1) of SCGM at each voxel (the
p.scgm.nii.gz volume). 

This script can be used to create a tighter mask that is specific to
the given subjects in the group. This can be compared to
$FREESURFER/average/mni_icbm152_nlin_asym_09c/synthseg.t1.subcortgm.mask.nii.gz,
which is the SCGMM created from segmenting the MNI152. The result can
be used in programs like vlrmerge-mni.

By default, a binary mask called scgm.mask.th50.close.nii.gz is
created by thresholding p.scgm.nii.gz volume at 0.5 (ie, 50% of the
subjects must be SCGM in a voxel for it to be in the mask) after which
a "closing" operation is applied by dilating by 1 voxel then eroding
by 1 voxel. The user can create different masks by binarizing
p.scgm.nii.gz.  This can be done in a number of ways. The differences
generally just affect what happens at the edge of the SCGM where the
registration is not perfect.  More agressive thresholding leaves
voxels where most/all subjects agree but eliminates voxels on the
edge. Each change of 0.2 of the threshold roughly changes the edge by
one voxel.  The threshold you actually use can depend somewhat on what
you are doing. Eg, if you are just making figures, then a more liberal
mask might work. If you are using it to only smooth within the mask,
then you have to decide on how important those extra voxels are vs
mixing SCGM signal with non-SCGM. This script creates a file called
nvox.thresh.dat that has a table of threshold vs number of voxels in
the mask.

Simple thresholding:
mri_binarize --i p.scgm.nii.gz --min 0.5 --o scgm.mask.nii.gz
This thresholds at 0.5, meaning that half of the subjects must have
SCGM at that voxel for it to be in the mask. Setting this threshold
higher means that more voxels are excluded. 

Simple thresholding with dilation:
mri_binarize --i p.scgm.nii.gz --min 0.5 --dilate 1 --o scgm.mask.nii.gz
This is the same as simple thresholding but then the mask is expanded 
(dilated) by one voxel.

Simple thresholding with closing:
mri_binarize --i p.scgm.nii.gz --min 0.5 --dilate 1 --erode 1 --o scgm.mask.nii.gz
This is the same as simple thresholding but then the mask is expanded 
(dilated) by one voxel and then eroded by one voxel, this tends to make the mask
have rounder, less jagged boundaries.