#! /bin/tcsh -f

#
# make_average_volume
#
# Creates average volumes from a set of subjects.
#
# --help option will show usage
#
# Original Author: Doug Greve
#
# Copyright © 2021 The General Hospital Corporation (Boston, MA) "MGH"
#
# Terms and conditions for use, reproduction, distribution and contribution
# are found in the 'FreeSurfer Software License Agreement' contained
# in the file 'LICENSE' found in the FreeSurfer distribution, and here:
#
# https://surfer.nmr.mgh.harvard.edu/fswiki/FreeSurferSoftwareLicense
#
# Reporting: freesurfer@nmr.mgh.harvard.edu
#
#

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

set VERSION = 'make_average_volume dev';

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 ConfFlag = "-conform"
set ctab = ()

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

source $FREESURFER_HOME/sources.csh

goto parse_args;
parse_args_return:

goto check_params;
check_params_return:

set StartTime = `date`;
set tSecStart = `date '+%s'`;

# 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)
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 (nu norm wm) #orig nu T1 brain wm # wm needed for surfaces?
  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
      echo "ERROR: cannot find $volid for $subject" | tee -a $LF
      exit 1;
    endif

    set thisXHemi = `basename $subject`
    if( "$thisXHemi" != "xhemi") then
      set xfmvol = $tmpdir/$volid-$subject.mgh
    else
      set tmp = `dirname $subject`
      set xfmvol = $tmpdir/$volid-$tmp-xhemi.mgh
    endif

    set cmd = (mri_convert $invol $xfmvol --apply_transform $xfm)
    if($transform_fname == talairach.xfm)  set cmd = ($cmd -oc 0 0 0 ); # sets output c_ras=0
    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 (this will conform as well unless -noconform)
  set outvol = $outdir/mri/$volid.mgz
  set cmd = (mri_average $ConfFlag $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

pushd $outdir/mri/
ln -s nu.mgz rawavg.mgz
ln -s nu.mgz orig.mgz
ln -s nu.mgz T1.mgz
ln -s norm.mgz brain.mgz
ln -s norm.mgz brainmask.mgz
popd

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
      echo "xfm $xfm" | tee -a $LF
      set xfmvol = $tmpdir/seg-$subject.mgh
      if($transform_fname != talairach.xfm) then
        set cmd = (mri_convert $invol $xfmvol --apply_transform $xfm \
           --resample_type 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
        if("$ConfFlag" == "-conform") then
	  set cmd = (mri_convert $xfmvol --conform -rt nearest -ns 1 -odt int $xfmvol)
          $cmd
          if($status) then
            pwd |& tee -a $LF
            echo "ERROR: $cmd" |& tee -a $LF
            exit 1;
          endif
        endif
      else
        set templatevol = $FREESURFER_HOME/average/mni305.cor.mgz
        set cmd = (mri_vol2vol --mov $invol --targ $templatevol --o $xfmvol --interp nearest)
        set bn0 = `basename $xfm`
        set bn = `basename $xfm .xfm`
        echo "bn = $bn"  | tee -a $LF
        if($bn0 == $bn.xfm) set cmd = ($cmd --xfm $xfm)
        set bn = `basename $xfm .m3z`
        echo "bn = $bn"  | tee -a $LF
        if($bn0 == $bn.m3z) set cmd = ($cmd --m3z $bn0 --s $subject)
        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
      endif

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

    # Concat and vote. 
    # Make sure seg is same space as brainmask
    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)
    if($#ctab) set cmd = ($cmd --ctab $ctab)
    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;

    # persurf is needed for the creation of ribbon and aparc+aseg
    # these steps are done with a call to recon-all in make_average_volume
    if($seg == aseg) then
      pushd $outdir/mri
      ln -s aseg.mgz aseg.presurf.mgz
      ln -s aseg.mgz aseg.presurf.hypos.mgz
      popd
    endif

 end # Loop over segs

endif #DoASeg

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

if($cleanup) rm -r $tmpdir

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 "make_average_volume-Run-Time-Sec $tSecRun" |& tee -a $LF
echo "make_average_volume-Run-Time-Hours $tRunHours" |& tee -a $LF
echo " " |& tee -a $LF

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 "--s":
    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 ( $getting_subjects ) then
        # got em all, from --subjects variable arg loop
        set getting_subjects = 0;
        shift;
      endif
      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 "--f":
      if ( $getting_subjects ) then
        # got em all, from --subjects variable arg loop
        set getting_subjects = 0;
        shift;
      endif
      if ( $#argv == 0) goto arg1moreerr;
      set fname = $argv[1]; shift;
      if(! -e $fname) then
        echo "ERROR: cannot find $fname"
        exit 1;
      endif
      if($?SUBJECTS) then
        set SUBJECTS = ($SUBJECTS `cat $fname`)
      else
        set SUBJECTS = (`cat $fname`)
      endif
      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":
    case "--v-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 "--ctab":
      if ( $getting_subjects ) then
        # got em all, from --subjects variable arg loop
        set getting_subjects = 0;
        shift;
      endif
      set ctab = $argv[1]; shift;
      breaksw

    case "--ctab-default":
      if ( $getting_subjects ) then
        # got em all, from --subjects variable arg loop
        set getting_subjects = 0;
        shift;
      endif
      set ctab = $FREESURFER_HOME/FreeSurferColorLUT.txt
      breaksw

    case "--threads"
      if ( $getting_subjects ) then
        # got em all, from --subjects variable arg loop
        set getting_subjects = 0;
        shift;
      endif
      if($#argv < 1) goto arg1err;
      set threads = $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":
    case "--template-only":
    case "--no-template-only":
      breaksw

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

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

    # These 1-arg flags passed to make_average_subject, but dont apply here.
    # Need to eat the argument.
    case "--surfreg":
    case "--surf-reg":
    case "--surf_reg":
    case "--ico":
    case "--annot"
    case "--meas"
    case "--rca-threads":
    case "--s-xform":
    case "--s-dest-lta"
      if ( $getting_subjects ) then
        set getting_subjects = 0;
        # got em all, from --subjects variable arg loop
        shift;
      endif
      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 "--conform":
    case "-conform":
      set ConfFlag = "-conform"
      if ( $getting_subjects ) then
        set getting_subjects = 0;
        # got em all, from --subjects variable arg loop
      endif
      breaksw
    case "--no-conform":
    case "-no-conform":
    case "-noconform":
      set ConfFlag = "-noconform"
      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;
  set outdir = $sdout/${average_subject}
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 "   --ctab colortable : embed into segmentations"
  echo "   --ctab-default : embed $FREESURFER_HOME/FreeSurferColorLUT.txt into segmentations"
  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