#!/bin/tcsh -f
# thalseg

set VERSION = 'make_average_subcort dev';

set outvol = ();
set subjlist = ();
set tmpdir = ();
set cleanup = 1;
set LF = ();
set transform_fname = talairach.xfm

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 outdir = `dirname $outvol`;
mkdir -p $outdir
pushd $outdir > /dev/null
set outdir = `pwd`;
popd > /dev/null

if($#tmpdir == 0) then
  set tmpdir = `fs_temp_dir --scratch`
endif
mkdir -p $tmpdir

set outstem = `fname2stem $outvol`
if($#LF == 0) set LF = $outstem.log
if($LF != /dev/null) rm -f $LF

echo "Log file for make_average_subcort" >> $LF
date  | tee -a $LF
echo "" | 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
echo "" | tee -a $LF
cat $FREESURFER_HOME/build-stamp.txt | tee -a $LF
uname -a  | tee -a $LF

set ctxvollist = ()
set subctxvollist = ()
set cblumvollist = ()
@ nth = 0
foreach subject ($subjlist)
  @ nth = $nth + 1
  echo "#@# $nth/$#subjlist $subject `date`" | tee -a $LF

  set xfm   = $SUBJECTS_DIR/$subject/mri/transforms/$transform_fname
  if(! -e $xfm) then
    echo "ERROR: cannot find $xfm" | tee -a $LF
    exit 1;
  endif

  set aseg = $SUBJECTS_DIR/$subject/mri/aseg.mgz

  # Make a mask of cortex, can use aseg with v6
  set ctx = $tmpdir/ctx.mgh
  set cmd = (mri_binarize --i $aseg --o $ctx --match 3 --match 42)
  echo $cmd | tee -a $LF
  $cmd | tee -a $LF
  if($status) exit 1;
  set xfmctx = $tmpdir/ctx.mask.$subject.tal.mgh
  set cmd = (mri_convert $ctx $xfmctx --apply_transform $xfm \
    --resample_type nearest -oc 0 0 0 ); # sets output c_ras=0
  echo $cmd | tee -a $LF
  $cmd  |& tee -a $LF
  if($status) exit 1;

  # Make a mask of all subcortical GM, specifying
  # the things to exclude
  set subctx = $tmpdir/subctx.mgh
  set cmd = (mri_binarize --i $aseg --inv --o $subctx \
    --match 0 --match 3 \
    --match 42 --match 4 --match 14 --match 15 --match 43 \
    --match 44 --match 72 --match 75 --match 76 --match 2 \
    --match 41 --match 7 --match 46 --match 251 --match 252 \
    --match 253 --match 254 --match 255 \
    --match 24 --match 31 --match 63  \
    --match 77 --match 78 --match 79 --match 80 \
    --match 81 --match 82)
  echo $cmd | tee -a $LF
  $cmd | tee -a $LF
  if($status) exit 1;
  set xfmsubctx = $tmpdir/subctx.$subject.tal.mgh
  set cmd = (mri_convert $subctx $xfmsubctx --apply_transform $xfm \
    --resample_type nearest -oc 0 0 0 ); # sets output c_ras=0
  echo $cmd | tee -a $LF
  $cmd  |& tee -a $LF
  if($status) exit 1;

  # Now make a mask of the cerebellum
  set cblum = $tmpdir/cblum.mgh
  set cmd = (mri_binarize --i $aseg --o $cblum --match 8 --match 47)
  echo $cmd | tee -a $LF
  $cmd | tee -a $LF
  if($status) exit 1;
  set xfmcblum = $tmpdir/cblum.$subject.tal.mgh
  set cmd = (mri_convert $cblum $xfmcblum --apply_transform $xfm \
    --resample_type nearest -oc 0 0 0 ); 
  echo $cmd | tee -a $LF
  $cmd  |& tee -a $LF
  if($status) exit 1;

  set ctxvollist = ($ctxvollist $xfmctx)
  set subctxvollist = ($subctxvollist $xfmsubctx)
  set cblumvollist = ($cblumvollist $xfmcblum)

end

# Average the volumes together
set pctx = $tmpdir/p.ctx.mgh
set cmd = (mri_concat $ctxvollist --mean --o $pctx)
echo $cmd | tee -a $LF
$cmd  | tee -a $LF
if($status) exit 1;

set psubctx = $tmpdir/p.subctx.mgh
set cmd = (mri_concat $subctxvollist --mean --o $psubctx)
echo $cmd | tee -a $LF
$cmd  | tee -a $LF
if($status) exit 1;

set pcblum = $tmpdir/p.cblum.mgh
set cmd = (mri_concat $cblumvollist --mean --o $pcblum)
echo $cmd | tee -a $LF
$cmd  | tee -a $LF
if($status) exit 1;

# Create mask of voxels with p<0.8 of cortex
set pctx80 = $tmpdir/p80.ctx.mgh
set cmd = (mri_binarize --i $pctx --min 0.8 --inv --o $pctx80)
echo $cmd | tee -a $LF
$cmd  | tee -a $LF
if($status) exit 1;

# Create mask of subcortex voxels that exclude any voxel
# with 80% or more cortex
set subcortmask0 = $tmpdir/subcort0.mgh
set cmd = (mri_binarize --i $psubctx --min 0.00001 \
  --mask $pctx80 --o $subcortmask0)
echo $cmd | tee -a $LF
$cmd  | tee -a $LF
if($status) exit 1;

# Create mask of voxels with p>0.2 of cortex
set pctx20 = $tmpdir/p20.ctx.mgh
set cmd = (mri_binarize --i $pctx --min 0.2  --o $pctx20)
echo $cmd | tee -a $LF
$cmd  | tee -a $LF
if($status) exit 1;

# Create mask of cblum voxels that exclude any voxel
# with 20% or more cortex; 
set cblummask = $tmpdir/cblum+ctx.mask.mgh
set cmd = (mri_binarize --i $pcblum --min 0.00001 \
  --mask $pctx20 --o $cblummask)
echo $cmd | tee -a $LF
$cmd  | tee -a $LF
if($status) exit 1;

# Invert the mask to leave out ctx voxels near cblum
set cblummaskex = $tmpdir/cblum+ctx.mask.ex.mgh
set cmd = (mri_binarize --i $cblummask --min 0.5 --inv --o $cblummaskex)
echo $cmd | tee -a $LF
$cmd  | tee -a $LF
if($status) exit 1;

# Create another refinement of the subcort mask with ctx voxels
# near cblum removed
set subcortmask1 = $tmpdir/subcort1.mgh
set cmd = (mri_binarize --i $subcortmask0 --min 0.5 --mask $cblummaskex --o $subcortmask1)
#set cmd = ($cmd --dilate 1 --erode 1)
echo $cmd | tee -a $LF
$cmd  | tee -a $LF
if($status) exit 1;

# Connected components, remove things less than 1000 vox
set subcortmask2 = $tmpdir/subcort2.mgh
set cmd = (mri_volcluster --in $subcortmask1 --minsizevox 1000 --ocn $subcortmask2 --thmin 0.5)
echo $cmd | tee -a $LF
$cmd  | tee -a $LF
if($status) exit 1;

# Final binarization with dilation and erosion. 4 was selected because
# it eliminated holes from the mask
set subcortmask3 = $tmpdir/subcort3.mgh
set cmd = (mri_binarize --i $subcortmask2 --min 0.5 --dilate 4 --erode 4 --o $outvol)
echo $cmd | tee -a $LF
$cmd  | tee -a $LF
if($status) exit 1;

if($cleanup) rm -r $tmpdir

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

exit 0

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

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

  set flag = $argv[1];
  if (! $getting_subjects) then
    shift;
  endif

  switch($flag)

    case "--help":
      set PrintHelp = 1;
      goto usage_exit;
      exit 0;

    case "--version":
      echo $VERSION
      exit 0;

    case "--subjects":
      if ( $#argv == 0) goto arg1moreerr;
      set subjlist = $argv[1]; shift;
      # loop on getting variable number of subject names
      set getting_subjects = 1; # see 'default:' case
      breaksw

    case "--fsgd":
      if ( $#argv == 0) goto arg1err;
      set fsgdf = $argv[1]; shift;
      if(! -e $fsgdf) then
        echo "ERROR: cannot find $fsgdf";
        exit 1;
      endif
      set sl = `cat $fsgdf | awk '{if($1 == "Input") print $2}'`;
      set subjlist = ($sl);
      breaksw

    case "--o":
      if ( $getting_subjects ) then
        # got em all, from --subjects variable arg loop
        set getting_subjects = 0;
        shift;
      endif
      if ( $#argv == 0) goto arg1err;
      set outvol = $argv[1]; shift;
      breaksw

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

    case "--xform":
      if ( $getting_subjects ) then
        # got em all, from --subjects variable arg loop
        set getting_subjects = 0;
        shift;
      endif
      if ( $#argv == 0) goto arg1err;
      set transform_fname = $argv[1]; shift;
      breaksw

    case "--nocleanup":
      set cleanup = 0;
      if ( $getting_subjects ) then
        set getting_subjects = 0;
        # got em all, from --subjects variable arg loop
      endif
      breaksw

    case "--debug":
    case "--echo":
      set echo = 1;
      set verbose = 1
      if ( $getting_subjects ) then
        set getting_subjects = 0;
        # got em all, from --subjects variable arg loop
      endif
      breaksw

    default:
      if ( $getting_subjects ) then
        # loop on getting variable number of subjects,
        # until a new flag is found, or no more args
        set subjlist = ( $subjlist $argv[1] ); shift;
        set getting_subjects = 1;
      else
        echo ERROR: Flag $flag unrecognized.
        echo $cmdline
        exit 1
      endif
      breaksw
  endsw

end

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

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

if($#outvol == 0) then
  echo "ERROR: no output specified"
  exit 1;
endif
if($#subjlist == 0) then
  echo "ERROR: no subjects specified"
  exit 1;
endif
foreach s ($subjlist)
  if(! -e $SUBJECTS_DIR/$s) then
    echo "ERROR: cannot find $s in $SUBJECTS_DIR"
    exit 1;
  endif
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 "make_average_subcort"
  echo "  --subjects s1 s2 ..."
  echo "  --o outvol"
  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 creates an average subcortical mask for the given inputs
subjects.  The final mask is in fsaverage/mni305 space. It is designed
to be used for the subcortical analysis in FSFAST where one needs to
exclude cortical voxels from the subcortical analysis.

Example Usage:

make_average_subcort --o subcort.mgz \
  --subjects 004 008 017 021 032 039 040 045 049 067 073 074 \
   080 084 091 092 093 095 097 099 102 103 106 108 111 114 123 \
   124 128 129 130 131 133 136 138 140 141 144 145 149