#! /bin/tcsh -f

#
# make_average_volume
#
# Creates average volumes from a set of subjects.
#
# --help option will show usage
#
# Original Author: Doug Greve
# CVS Revision Info:
#    $Author: greve $
#    $Date: 2009/09/14 23:29:09 $
#    $Revision: 1.29 $
#
# Copyright (C) 2002-2007,
# The General Hospital Corporation (Boston, MA).
# All rights reserved.
#
# Distribution, usage and copying of this software is covered under the
# terms found in the License Agreement file named 'COPYING' found in the
# FreeSurfer source code root directory, and duplicated here:
# https://surfer.nmr.mgh.harvard.edu/fswiki/FreeSurferOpenSourceLicense
#
# General inquiries: freesurfer@nmr.mgh.harvard.edu
# Bug reports: analysis-bugs@nmr.mgh.harvard.edu
#

# uncomment this to increase number of allowable open files to maximum:
#limit descriptors unlimited

set VERSION = '$Id: make_average_volume,v 1.29 2009/09/14 23:29:09 greve Exp $';

set PrintHelp = 0;
set transform_fname = talairach.xfm
set average_subject = average
set KeepAllOrig = 0;
set sdout = ();
set DoASeg = 1;
set cleanup = 1;
set DoXHemi = 0;

set cmdargs = ($argv);
if($#argv == 0) then
  # zero args is allowed only if SUBJECTS env var is declared
  if ( ! $?SUBJECTS) then
      goto usage_exit;
  endif
endif

goto parse_args;
parse_args_return:

goto check_params;
check_params_return:

# Include cross-hemi registration
if($DoXHemi) then
  set tmp = ();
  foreach s ($SUBJECTS)
    if(! -e $SUBJECTS_DIR/$s/xhemi) then
      echo "ERROR: $s/xhemi does not exist"
      exit 1;
    endif
    set tmp = ($tmp $s $s/xhemi)
  end
  set SUBJECTS = ($tmp);
endif

# This will put the data under $sdout (SUBJECTS_DIR by default)
set outdir = $sdout/${average_subject}
mkdir -p $outdir
if($status) then
  echo "ERROR: could not make $outdir"
  exit 1;
endif
mkdir -p $outdir/scripts
mkdir -p $outdir/mri
mkdir -p $outdir/mri/transforms

set LF = $outdir/scripts/make_average_volume.log
if(-e $LF) mv $LF $LF.bak

echo Log file is $LF
echo ""
echo $0  | tee -a $LF
echo $cmdargs  | tee -a $LF
echo ""
date | tee -a $LF
pwd  | tee -a $LF
echo "setenv SUBJECTS_DIR $SUBJECTS_DIR"  | tee -a $LF
echo "sdout $sdout" | tee -a $LF
echo $SUBJECTS >> $LF

set tmpdir = $outdir/tmp/make_average_vol-tmp-$$
mkdir -p $tmpdir

#----------------------------------------------------------
foreach volid (orig T1 brain)
  set invollist = ();

  foreach subject ($SUBJECTS)

    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 invol = $SUBJECTS_DIR/$subject/mri/$volid.mgz
    if(! -e $invol) then
      set invol = $SUBJECTS_DIR/$subject/mri/$volid
      if(! -e $invol/COR-.info) then
        echo "ERROR: cannot find $volid for $subject" | tee -a $LF
        exit 1;
      endif
    endif

    # Use mgh to keep from having to compress and decompress
    set xfmvol = $tmpdir/$volid-$subject.mgh

    set cmd = (mri_convert $invol $xfmvol  --apply_transform $xfm)
    set cmd = ($cmd -oc 0 0 0 ); # sets output c_ras=0
    #if($transform_fname == talairach.xfm) set cmd = ($cmd --devolvexfm $subject);
    echo $cmd | tee -a $LF
    $cmd  |& tee -a $LF
    if($status) then
      pwd |& tee -a $LF
      echo "ERROR: mri_convert failed." |& tee -a $LF
      exit 1;
    endif

    set invollist = ($invollist $xfmvol);
  end

  # Average the volumes together
  set outvol = $outdir/mri/$volid.mgz
  set cmd = (mri_average $invollist $outvol)
  echo $cmd | tee -a $LF
  $cmd  | tee -a $LF
  if($status) exit 1;

  # remove reference to individual subject's talairach.xfm
  set cmd = (mri_add_xform_to_header -c auto $outvol $outvol)
  echo $cmd | tee -a $LF
  $cmd  | tee -a $LF
  if($status) exit 1;

  # When keeping all orig, concat
  if($KeepAllOrig && $volid == orig) then
    echo "Concatenating all subjects' orig volumes"
    set cmd = (mri_concat $invollist --o $outdir/mri/orig.all.mgz)
    echo $cmd | tee -a $LF
    $cmd  | tee -a $LF
    if($status) exit 1;
  endif

end

cp $outdir/mri/brain.mgz $outdir/mri/brainmask.mgz

set outtalxfm = $outdir/mri/transforms/talairach.xfm
rm -f $outtalxfm
echo "MNI Transform File" >> $outtalxfm
echo "" >> $outtalxfm
echo "Transform_Type = Linear;" >> $outtalxfm
echo "Linear_Transform =" >> $outtalxfm
echo "1.0 0.0 0.0 0.0 " >> $outtalxfm
echo "0.0 1.0 0.0 0.0 " >> $outtalxfm
echo "0.0 0.0 1.0 0.0;" >> $outtalxfm

#----------------------------------------------------------
if($DoASeg) then
  foreach seg (aseg)
    # aparc+aseg and wmparc dont work so well
    set invollist = ()
    @ nth = 0;
    foreach subject ($SUBJECTS)
      @ nth = $nth + 1;
      echo "--------------------------------------" | tee -a $LF
      echo "$nth/$#SUBJECTS $subject $seg `date`" | tee -a $LF

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

      # ASeg
      set invol = $SUBJECTS_DIR/$subject/mri/$seg.mgz
      if(! -e $invol) then
        echo "ERROR: cannot find $volid for $subject" | tee -a $LF
        exit 1;
      endif

      # Reslice, use nearest neighbor, use mgh format
      set xfmvol = $tmpdir/seg-$subject.mgh
      set cmd = (mri_vol2vol --mov $invol \
           --targ $FREESURFER_HOME/average/mni305.cor.mgz \
           --xfm $xfm  --o $xfmvol --interp nearest)
      echo $cmd | tee -a $LF
      $cmd  |& tee -a $LF
      if($status) then
        pwd |& tee -a $LF
        echo "ERROR: mri_vol2vol failed." |& tee -a $LF
        exit 1;
      endif

      set invollist = ($invollist $xfmvol);
    end
    echo "--------------------------------------" | tee -a $LF

    # Concat and vote. 
    set cmd = (mri_concat $invollist --o $tmpdir/vote.mgz --vote \
      --mask $outdir/mri/brainmask.mgz --debug)
    echo $cmd | tee -a $LF
    $cmd  | tee -a $LF
    if($status) then
      echo "WARNING: could not create average segmentation for $seg ..." | tee -a $LF
      echo "  ... but continuing" | tee -a $LF
      continue;
    endif
    # Extract Seg
    set aseg = $outdir/mri/$seg.mgz
    set cmd = (mri_convert $tmpdir/vote.mgz $aseg --frame 0)
    echo $cmd | tee -a $LF
    $cmd  | tee -a $LF
    if($status) exit 1;
    # Extract Probability of ASeg
    set pseg = $outdir/mri/p.$seg.mgz
    set cmd = (mri_convert $tmpdir/vote.mgz $pseg --frame 1)
    echo $cmd | tee -a $LF
    $cmd  | tee -a $LF
    if($status) exit 1;

 end # Loop over segs

endif #DoASeg

#----------------------------------------------------------

if($cleanup) rm -r $tmpdir

date | tee -a $LF
echo "make_average_volume 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 SUBJECTS = $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 SUBJECTS = ($sl);
      breaksw

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

    case "--sd-out":
      if ( $getting_subjects ) then
        # got em all, from --subjects variable arg loop
        set getting_subjects = 0;
        shift;
      endif
      if ( $#argv == 0) goto arg1err;
      set sdout = $argv[1]; shift;
      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 "--sd":
    case "--sdir":
      if ( $getting_subjects ) then
        # got em all, from --subjects variable arg loop
        set getting_subjects = 0;
        shift;
      endif
      if ( $#argv == 0) goto arg1err;
      set SUBJECTS_DIR = $argv[1]; shift;
      setenv SUBJECTS_DIR $SUBJECTS_DIR
      breaksw

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

    # Symlink is passed to make_average_surface thru make_average_subject
    case "--symlink":
    case "--no-symlink":
    case "--no-link":
      breaksw

    # These are flags passed to make_average_subject, but dont apply here
    case "--no-link":
    case "--link":
    case "--no-surf":
    case "--no-vol":
    case "--force":
    case "--xhemi":
      if ( $getting_subjects ) then
        set getting_subjects = 0;
        # got em all, from --subjects variable arg loop
      endif
      breaksw

    # These are flags passed to make_average_subject, but dont apply here.
    # Need to eat the arguments.
    case "--surfreg":
    case "--surf-reg":
    case "--surf_reg":
    case "--ico":
      if ( $getting_subjects ) then
        set getting_subjects = 0;
        # got em all, from --subjects variable arg loop
        shift;
      endif
      set SurfReg = $argv[1]; shift;
      breaksw

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

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

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

    case "--xhemi":
      set DoXHemi = 1;
      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 SUBJECTS = ( $SUBJECTS $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 (! $?SUBJECTS) then
    echo "ERROR: no subjects declared!"
    echo "Either declare subjects in SUBJECTS variable,"
    echo "or declare using --subjects argument."
    exit 1
  endif
  if (! $?SUBJECTS_DIR) then
    echo "ERROR: SUBJECTS_DIR is not declared!"
    echo "Either set the SUBJECTS_DIR environment variable,"
    echo "or declare using --sdir argument, the root directory"
    echo "for subject data files."
    exit 1
  endif
  if(! -e $SUBJECTS_DIR ) then
    echo "ERROR: SUBJECTS_DIR $SUBJECTS_DIR does not exist."
    exit 1;
  endif
  if(! $?FREESURFER_HOME ) then
    echo "ERROR: environment variable FREESURFER_HOME not set."
    exit 1;
  endif
  if(! -e $FREESURFER_HOME ) then
    echo "ERROR: FREESURFER_HOME $FREESURFER_HOME does not exist."
    exit 1;
  endif
  if($#sdout == 0) set sdout = $SUBJECTS_DIR;
goto check_params_return;
############--------------##################

############--------------##################
arg1err:
  echo "ERROR: flag $flag requires one argument"
  exit 1
############--------------##################

############--------------##################
arg1moreerr:
  echo "ERROR: flag $flag requires one or more arguments"
  exit 1
############--------------##################

############--------------##################
usage_exit:
  echo ""
  echo "USAGE: make_average_volume"
  echo ""
  echo "Required Arguments"
  echo "   --subjects <subj1> <subj2> ... <subjN>"
  echo "             : or declare subjects in SUBJECTS env var"
  echo "   --fsgd fsgdfile : get subject list from fsgd"
  echo ""
  echo "Optional Arguments"
  echo "   --out <average subject name>    : default name is 'average'"
  echo "   --topdir topdir : put data here and link to SUBJECTS_DIR"
  echo "   --xform xformname : use mri/transforms/xformname (def is talairach.xfm)"
  echo "   --sdir <SUBJECTS_DIR to use instead of the one in the env>"
  echo "   --sd      : same as --sdir"
  echo "   --force   : overwrite existing average subject data"
  echo "   --keep-all-orig : concatenate all orig vols into mri/orig.all.mgz"
  echo "   --no-aseg : do not create 'average' aseg"
  echo ""
  echo "   --help    : short descriptive help"
  echo "   --version : script version info"
  echo "   --echo    : enable command echo, for debug"
  echo "   --debug   : same as --echo"
  echo "   --nocleanup : do not delete temporary files"
  echo ""
  echo "See also: recon-all, make_final_surfaces, morph_subject"
  echo ""

  if(! $PrintHelp) exit 1;

  echo Version: $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 average volumes from a set of subjects.

EXAMPLE

make_average_volume --out avgsubject --subjects subj1 subj2 subj3 subj4

will create $SUBJECTS_DIR/avgsubject with orig.mgz, brain.mgz, and T1.mgz
which will be averages of subjects 1-4.

SEE ALSO

make_average_subject make_average_surface

