#!/bin/tcsh -f

# seg2filled - creates a filled.mgz from an aseg-style
# segmentation. The original intent is to use a SAMSEG segmentation to
# created the filled so don't have to do it using the recon-all volume
# stream

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

set VERSION = 'seg2filled dev';

set seg = ();
set ndil = 1
set hemilist = (lh rh)
set norm = ()
set SimulateCavity = 0; # for testing filling of cavities
set tmpdir = ();
set cleanup = 1;
set LF = ();
set filled = ()
set surfname = ()
set surfdir = ();
set subject = ();

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
pushd $outdir > /dev/null
set outdir = `pwd`;
popd > /dev/null
if($#surfname) mkdir -p $surfdir

if($#tmpdir == 0) then
  # This can generate a lot of output and may fill up /scratch
  set tmpdir = $outdir/tmpdir.seg2filled.$$
endif
mkdir -p $tmpdir

# Set up log file
set base = `fname2stem $filled`
if($#LF == 0) set LF = $base.log
if($LF != /dev/null) rm -f $LF
echo "Log file for seg2filled" >> $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

#========================================================
# This is a list of structures to be included in the subcortical mass
# Exclude InvLatVent, Hip, Amyg, Cblum and any unlateralized, except for WMH
# Exclude CC (251-255). 
set lhseglist = ( 2  4  9 10 11 12 13 26 27 28 29 30 31 78 81)
set rhseglist = (41 43 48 49 50 51 52 58 59 60 61 62 63 79 82)
if($SimulateCavity) then
  # Remove #13 or #52 (pallidum) to simulate a cavity. 
  echo "Removing 13 and 52 (pallidum) to simulate a cavity." | tee -a $LF
  set lhseglist = ( 2  4  9 10 11 12    26 27 28 29 30 31 78 81)
  set rhseglist = (41 43 48 49 50 51    58 59 60 61 62 63 79 82)
endif
# Add WM hypointensities. Hypos are not lateralized, but it should not
# matter unless one is on the boundary of left and right, and even
# then maybe not.
set lhseglist = ($lhseglist 77) # WMH
set rhseglist = ($rhseglist 77) # WMH

set c1list = ()
foreach hemi ($hemilist)
  if($hemi == lh) then
    set match = ($lhseglist)
    set hemivalue = 255;
  endif
  if($hemi == rh) then
    set match = ($rhseglist)
    set hemivalue = 127;
  endif

  # Binarize the seg to create the subcoritcal mass
  set subcortmass0 = $tmpdir/subcortmass0.$hemi.mgh
  set cmd = (mri_binarize --i $seg --match $match --o $subcortmass0)
  echo $cmd |& tee -a $LF
  $cmd |& tee -a $LF
  if($status) goto error_exit;

  # Extact all connected components (clusters).  This will eliminate
  # edge and corner defects if they are on the edge of the mass. Still
  # need to run pretess.
  set sum = $tmpdir/subcortmass.$hemi.ocn.sum
  set ocn = $tmpdir/subcortmass.$hemi.ocn.mgh
  set cmd = (mri_volcluster --in $subcortmass0 --match 1 --sum $sum --ocn $ocn)
  echo $cmd |& tee -a $LF
  $cmd |& tee -a $LF
  if($status) goto error_exit;
  if($cleanup) rm $subcortmass0

  # Extract the main component (--match 1)
  set c1 = $tmpdir/subcortmass.$hemi.c1.unfilled.mgh
  set cmd = (mri_binarize --i $ocn --match 1 --o $c1 --binval $hemivalue)
  echo $cmd |& tee -a $LF
  $cmd |& tee -a $LF
  if($status) goto error_exit;

  # Look for cavities in the main component 
  # Invert the main component
  set c1inv = $tmpdir/subcortmass.$hemi.c1.inv.mgh
  set cmd = (mri_binarize --i $ocn --match 1 --inv --o $c1inv)
  echo $cmd |& tee -a $LF
  $cmd |& tee -a $LF
  if($status) goto error_exit;
  if($cleanup) rm $ocn
  # For speed, create a mask before extracting connected
  # components. Otherwise, there will be a huge concomp outside the
  # brain which will take a long time to extract. To create the mask,
  # binarize and dilate the seg. The mask must cover any cavities.
  # Could remove extracerebral segs to make faster. 
  set segmaskdil = $tmpdir/subcortmass.$hemi.c1.mgh
  set cmd = (mri_binarize --i $seg --min 0.5 --dilate $ndil --o $segmaskdil)
  echo $cmd |& tee -a $LF
  $cmd |& tee -a $LF
  if($status) goto error_exit;
  # Get clusters on the inverted. Use dilatated as a mask
  set invsum = $tmpdir/subcortmass.$hemi.c1.inv.ocn.sum
  set invocn = $tmpdir/subcortmass.$hemi.c1.inv.ocn.mgh
  set cmd = (mri_volcluster --in $c1inv --match 1 --sum $invsum --ocn $invocn --mask $segmaskdil)
  echo $cmd |& tee -a $LF
  $cmd |& tee -a $LF
  if($status) goto error_exit;
  if($cleanup) rm $c1inv
  # If there is only one cluster, then there are no cavities
  set NClusters = `grep NClusters $invsum | awk '{print $3}'`
  echo "NClusters $NClusters" |& tee -a $LF
  if($NClusters > 1) then
    # Fill the holes
    # Remove the big cluster 
    set invocn2 = $tmpdir/subcortmass.$hemi.c1.inv.ocn2.mgh
    set cmd = (mri_binarize --i $invocn --o $invocn2 --replaceonly 1 0)
    echo $cmd |& tee -a $LF
    $cmd |& tee -a $LF
    if($status) goto error_exit;
    # Binarize what is left
    set invocn3 = $tmpdir/subcortmass.$hemi.c1.inv.ocn3.mgh
    set cmd = (mri_binarize --i $invocn2 --o $invocn3 --min 0.5 --binval $hemivalue)
    echo $cmd |& tee -a $LF
    $cmd |& tee -a $LF
    if($status) goto error_exit;
    # add to the main component
    set c1filled = $tmpdir/subcortmass.$hemi.c1.filled.mgh
    set cmd = (mri_concat --sum $c1 $invocn3 --o $c1filled)
    echo $cmd |& tee -a $LF
    $cmd |& tee -a $LF
    if($status) goto error_exit;
    if($cleanup) rm $c1 $invocn $invocn2 $invocn3 
    set c1 = $c1filled
  endif

  # Run pretess
  set c1pretess = $tmpdir/subcortmass.$hemi.c1.pretess.mgh
  set cmd = (mri_pretess $c1 wm $norm $c1pretess)
  echo $cmd |& tee -a $LF
  $cmd |& tee -a $LF
  if($status) goto error_exit;
  if($cleanup) rm $c1
  # Have to binarize it again because pretess changes intensities
  set cmd = (mri_binarize --i $c1pretess --min 0.5 --o $c1pretess --binval $hemivalue)
  echo $cmd |& tee -a $LF
  $cmd |& tee -a $LF
  if($status) goto error_exit;

  set c1list = ($c1list $c1pretess)
end # hemi

if($#hemilist == 2) then
  # Create the filled.mgz
  # Add lh and rh together
  set filled0 = $tmpdir/filled0.mgz
  set cmd = (mri_concat $c1list --sum --o $filled0)
  echo $cmd |& tee -a $LF
  $cmd |& tee -a $LF
  if($status) goto error_exit;
  if($cleanup) rm $c1list
  # Remove any voxels present in both lh and rh. This can happen
  # because pretess might add a voxel. Which begs the question
  # as to whether it should be removed. Below keeps only voxels
  # that are 127 or 255 (a little hacky). 
  set fillmask = $tmpdir/fillmask.mgh
  set cmd = (mri_binarize --i $filled0 --match 127 255 --o $fillmask)
  echo $cmd |& tee -a $LF
  $cmd |& tee -a $LF
  if($status) goto error_exit;
  set filledf = $tmpdir/filled.float.mgz
  set cmd = (mri_mask $filled0 $fillmask $filledf)
  echo $cmd |& tee -a $LF
  $cmd |& tee -a $LF
  if($status) goto error_exit;
  # Convert to uchar, just in case
  if($#filled == 0) set filled = $outdir/filled.mgz
  set cmd = (mri_convert $filledf -odt uchar --no_scale 1 $filled)
  echo $cmd |& tee -a $LF
  $cmd |& tee -a $LF
  if($status) goto error_exit;
else
  set cmd = (mri_convert $c1list[1] -odt uchar --no_scale 1 $filled)
endif

if($#surfname) then
  foreach hemi ($hemilist)
    # Create the surface based on the (filled) main component
    if($hemi == lh) set hemivalue = 255;
    if($hemi == rh) set hemivalue = 127;
    set surfout = $surfdir/$hemi.$surfname
    set cmd = (mri_binarize --i $filled --match $hemivalue --surf $surfout)
    echo $cmd |& tee -a $LF
    $cmd |& tee -a $LF
    if($status) goto error_exit;
  end
endif

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

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

# Done
echo " " |& tee -a $LF
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 "Seg2filled-Run-Time-Sec $tSecRun" |& tee -a $LF
echo "Seg2filled-Run-Time-Min $tRunMin" |& tee -a $LF
echo "Seg2filled-Run-Time-Hours $tRunHours" |& tee -a $LF
echo " " |& tee -a $LF
echo "seg2filled 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 filled = $argv[1]; shift;
      breaksw

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

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

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

    case "--lh":
      set hemilist = (lh)
      breaksw

    case "--rh":
      set hemilist = (rh)
      breaksw

    case "--cavity":
      set SimulateCavity = 1
      breaksw

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

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

    case "--ndil":
      if($#argv < 1) goto arg1err;
      set ndil = $argv[1]; shift;
      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) then
  if(! -e $SUBJECTS_DIR/$subject) then
    echo "ERROR: cannot find $subject"
    exit 1;
  endif
  set seg = $SUBJECTS_DIR/$subject/mri/aseg.presurf.mgz
  set norm = $SUBJECTS_DIR/$subject/mri/norm.mgz
  set filled = $SUBJECTS_DIR/$subject/mri/aseg.filled.mgz
  set surfname = orig.aseg.nofix
  set surfdir  = $SUBJECTS_DIR/$subject/surf
endif

if($#seg == 0) then
  echo "ERROR: must spec seg"
  exit 1;
endif
if($#norm == 0) then
  echo "ERROR: must spec norm"
  exit 1;
endif
if($#filled == 0) then
  echo "ERROR: must spec output"
  exit 1;
endif
set outdir = `dirname $filled`
if($#surfdir == 0) set surfdir = $outdir

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 "seg2filled --seg seg.mgz --norm norm.mgz --o filled.mgz"
  echo "  --ndil ndil : used to speed cavity detection"
  echo "  --cavity : simulate a cavity to test the filling operation"
  echo "  --surf surfname : create ?h.surfname"
  echo "  --surfdir surfdir : dir to put surf (default is same as filled)"
  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

Creates a filled.mgz from an aseg-style segmentation. The original
intent is to use a SAMSEG segmentation to created the filled so do not
have to do it using the recon-all volume stream. It should not be
necessary to run mri_pretess on the output.