ppf_tsnr

#!/bin/bash

# this is a more complicated way of doing:
#  tsnr(){ 3dTstat -overwrite -prefix mean.nii.gz -mean $1; 3dDetrend -overwrite -polort 4 -prefix det.nii.gz $1; 3dTstat -overwrite -prefix std.nii.gz -stdev det.nii.gz; 3dcalc -overwrite -prefix tsnr.nii.gz -m mean.nii.gz -s std.nii.gz -exp 'm/s';}

# default files -- also can be set by command switches
DEFRES=2.3
gm_mask() { echo "/opt/ni_tools/standard/mni_icbm152_nlin_asym_09c/mni_icbm152_gm_tal_nlin_asym_09c_${1}mm.nii";}
[ -z $GM_MNI ]      && GM_MNI="$(gm_mask $DEFRES)"
[ -z "$TEMPLATE" ]  && TEMPLATE=template_brain.nii # what does mni look like
[ -z "$WARPCOEF" ]  && WARPCOEF=func_to_standard_warp_allv.nii.gz  # how to warp native to mni
[ -z "$STARTFILE" ] && STARTFILE="_*.nii.gz"
# insturction settingts
[ -z "$CLEANUP" ]   && CLEANUP=yes   # remove ~850MB of nifti files when we are done
[ -z "$CHECKFLAG" ] && CHECKFLAG=yes # should we check for complete file flag .preprocessfunctional_complete?
[ -z "$IS_WARP" ] && IS_WARP=1 # usually we have warped the data (20190604: pp w/ -no_warp)
[ -z "$DROPVOLS" ] && DROPVOLS=0 # default to not truncating 
ONLYONE=""
PDIR=""
VERBOSE=0

# dont give warnings about oblique datasets (3dDetrend, 3dTstats, maskave on native space data)
export AFNI_NO_OBLIQUE_WARNING=YES 

usage() {
    [ -n "$1" ] && echo "ERROR:  $@" && echo -n "USAGE:  "
    me=$(basename $0)
    echo -n "$me [options] [path/to/preprocessFunctional]"
    # if error message, dont print everything
    if [ -n "$1" ]; then
       echo "; see $me -h"
       exit 1
    fi
    echo
    echo " OPTIONS: "
    echo "   -g MASK      MNI greymatter mask ($GM_MNI) or native space mask if also -n"
    echo "   -t TEMPLATE  MNI template ($TEMPLATE)"
    echo "   -w WARPCOEF  native to mni warp ($WARPCOEF)"
    echo "   -s STARTFILE example (#tr, suffix) file to use ($STARTFILE)"
    echo "   -O NIIFILE   only calculate for this file"
    echo "   -x X         remove first X volumes (ideal usage is in preprocessFunctional, not here)"
    echo "   -f   do not look for .preprocessfunctional_complete 'flag' file"
    echo "   -c   do not clean up tsnr/ temporary files"
    echo "   -n   native/no warp preprocessing, bet mc_target.nii.gz for mask instead"
    echo "   -v   be verbose, repeat for more text; -v -v -v => set -x"
    echo "   -h   this help message"
    echo " OUTPUT:"
    echo "   tsnr directory with *_tsnr.{nii.gz,txt} for each preprocessing step"
    echo " EXAMPLE: start with warped func, specify warpfile and template"
    echo "   cd /Volumes/Hera/preproc/cog_task/rest_spikemin/10124_20070829/snip"
    echo "   o=\$(dirname \$(readlink -f _func.nii.gz ))"
    echo "   ppf_tsnr -s wdktm_func.nii.gz  -t  \$o/template_brain.nii -w \$o/func_to_standard_warp_allv.nii.gz -v -v"
    echo " EXAMPLE: run for only the input file with native space mask (mask mc_target.nii.gz), ignore missing .preprocessfunctional_complete"
    echo "   ppf_tsnr -O _func.nii.gz  -n -f "
    exit 0
}

find_final() {
   # list all sufixes with the same (or greater) tr
   # sort by lenght of prefix. print only the longests (lastest)
   suffix=$1;shift
   ntr=$1; shift
   3dinfo -iname -nt *_$suffix |
     $suffix// |awk "(\$2 >= $ntr){print length(\$1), \$1}"  |
     sort -nr |uniq |
     sed 1q
}

is_same_grid() {
  [ $(3dinfo -same_grid $@|sort|sed 1q ) -eq 1 ] && return 0
  return 1
}

TSNR_onerun(){
   [ -z "$PDIR" -o ! -d "$PDIR" ] && echo "$FUNCNAME: something went badly!" && exit 1
   cd $PDIR
   pdir=$(pwd)

   # some checks
   firstfile=$(3dinfo -iname -nt $STARTFILE |awk '($2>10){print $1;exit}')
   [ -z "$firstfile" -o ! -r "$firstfile" ] && usage "cannot find start file ($(pwd)/$STARTFILE w/more than 10 trs); use '-s STARTFILE' option" 
   [ "$CHECKFLAG" == "yes" -a ! -r "$pdir/.preprocessfunctional_complete" ] && usage "$(pwd) does not contain a preprocessfunctional_complete flag file (use '-f' to ignore)"

   # only need to warp if we are starting before warp file 
   # and we are not doing native space thing
   NEEDWARP=0
   [ $IS_WARP -eq 1 ] && ! [[ $(basename $STARTFILE) =~ w[a-z]*_ ]] && NEEDWARP=1
   [ $VERBOSE -gt 1 ] && echo "# NEEDWARP? $NEEDWARP"

   # make sure we have files
   [ $NEEDWARP -eq 1 ]  && 
    for v in TEMPLATE GM_MNI WARPCOEF; do
       # if start file is already warped, we dont need to warp anything else
       [ ! -r "${!v}" ] && usage "cannot find $v (${!v}), consider changing with command options"
    done

   [ ! -d tsnr ] && mkdir -p tsnr

   # find pipeline's first imaging file (nii.gz starts with _, has more than 10 trs) 

   # suffix is all but the first _, includes .nii.gz: _func.nii.gz -> func.nii.gz; brnswdktm_func_5.nii.gz -> func_5.nii.gz
   suffix=${firstfile#*_} #suffix=${firstfile:1:${#firstfile}}  # does not work if firstfile is not _
   ntr=$(3dinfo -nt "$firstfile")
   [ $VERBOSE -gt 0 ] && echo "first file $firstfile, pipeline suffix $suffix, $ntr trs"

   # make sure we are using the correct mask
   if [ $IS_WARP -eq 1 ] && ! is_same_grid $TEMPLATE $GM_MNI; then
      # if we are using the default mask, try one that is the same
      local tmpl_res=$(3dinfo -adi $TEMPLATE | sed 's/.\?0*$//')
      # if we have the default, try at current template res res
      [ "$GM_MNI" = "$(gm_mask $DEFRES)" ] && GM_MNI=$(gm_mask $tmpl_res)
      [ ! -r "$GM_MNI" ] ||  ! is_same_grid $TEMPLATE $GM_MNI &&
        usage "$TEMPLATE does not match grid of mask ($GM_MNI)! Does voxel size match? use -g to specify mask"
   fi
   #finalfile=$(find_final $suffix $ntr)
   #[ $(3dinfo -same_grid $finalfile $GM_MNI ) -ne 1 ] &&
   #   usage "final file ($finalfile) does not match grid of mask ($GM_MNI)"


   # grab all files matching the suffix with the same number of trs as first file

   cnt=0
   if [ -n "$ONLYONE" ]; then
      filelist="$STARTFILE"
      [  $VERBOSE -gt 0 ] && echo "using only $STARTFILE" 
   else
      filelist="$(3dinfo -iname -nt *_${suffix/.nii.gz/}*.nii.gz | awk "(\$2 == $ntr){print \$1}")"
      [ $VERBOSE -gt 1 ] && echo "looking for '*${suffix/.nii.gz/}*' w/$ntr trs:" && \
        3dinfo -nt -space -iname  *_${suffix/.nii.gz/}*.nii.gz |
        sort -n |sed 's/^/\t/'
   fi
   for file in $filelist; do
      [ $VERBOSE -gt 1 ] && echo "# looking at $file"
      [ ! -r $file ] && echo "UTOH! bad expected file $file" >&2 && continue

      let ++cnt # count files we expect to have

      #header="${header}$conditions $conditions" #will this even be right? #producing text file...
      fname=$(basename $file) #the file name of 'file'
      prefix=${fname//.nii.gz/} #the file name of 'file' without the extension
      outname=${fname//_${suffix}/}_tsnr.nii.gz
      [ $DROPVOLS -ne 0 ] && outname=${fname//_${suffix}/}_tsnr-$DROPVOLS.nii.gz

      # remove nii.gz from txt filename
      prefix_=${prefix%%_*}
      stepnum=$(printf "%02d" ${#prefix_})
      txtout=$pdir/tsnr/$stepnum-$prefix.txt
      [ $DROPVOLS -ne 0 ] && txtout=${txtout/.txt}-$DROPVOLS.txt && echo "# making $txtout from $outname"
      # dont do anything if we already have txt file
      [ -s "$txtout" ] && [ $(cat "$txtout" | wc -l  ) -gt 0 ] && continue
       
      #check to see what space you are in
      space=$(3dinfo -space ${prefix}.nii.gz)
      if [ $space == "MNI" ]; then
         gm_mask=$GM_MNI
      # 20190604 - allow nowarp
      elif [ $IS_WARP -eq 0 ]; then
         # if user didn't specify a mask (-g) then GM_MNI space = "MNI" (still default)
         # change default to mprage_bet.nii.gz
         if [ $(3dinfo -space "$GM_MNI") == "MNI" ]; then
            local nw_input="mc_target.nii.gz"
            local nw_mask="$(pwd)/tsnr/func_bet.nii.gz"
            # make the "gm" mask (actually just brain here) the output of bet
            gm_mask=${nw_mask/.nii.gz/_mask.nii.gz} 
            if [ ! -r $nw_input ]; then
              echo "missing '$nw_input' to make mask. set your own mask with -g"
              exit 1
           fi
           [ $VERBOSE -gt 0 ] && echo "# using mask '$nw_input' ($nw_mask)"
           [ ! -r "$gm_mask" ] && bet "$nw_input" "$nw_mask" -F -m
         else
            # GM_MNI is what was specified by -g, hopefully not MNI and probalby not GM
            gm_mask=$GM_MNI
            [ $VERBOSE -gt 0 ] && echo "# using specifed mask $gm_mask without warping"
         fi
      else
         gm_mask=$pdir/tsnr/gm_mask_native.nii.gz
         # we need native space gm mask
         if [ ! -r $gm_mask ]; then
            [ $VERBOSE -gt 0 ] && echo "# making native space grey matter mask using $file (as resp for all)"
            WARPEXAMPLE=$file
            #gray matter mask in native space -need this once for each subject
            #make the warp
            invwarp -w $WARPCOEF -o standard_to_func_warp.nii.gz -r $file
            #apply the warp
            applywarp --in="$GM_MNI" \
               --out=$gm_mask \
               --warp=standard_to_func_warp.nii.gz \
               --ref=${file} \
               --rel \
               --interp=spline
         fi
         # check mask
         #  initially used '$firstfile' for warp
         #  but that file could have not been in native space!
         if ! is_same_grid $file $gm_mask; then 
             echo "BAD WARP mask/file mismatch!!"
             echo "  $gm_mask does not match $file (used '$WARPEXAMPLE' as template [empty means not created here])" 
             echo "   consider 'rm $gm_mask'"
             continue
         fi
      fi
      [ ! -r "$gm_mask" ] && echo "cannot read mask: '$gm_mask', set with -g" && exit 1

      tsnrout=$pdir/tsnr/$outname 
      #echo "checking $tsnrout"
      if [ ! -r "$tsnrout" ] ; then
         echo "creating tsnr image $outname w/mask $gm_mask"
         # need to put mean back in for tsnr calc
         if [ ${file:0:1} == "b" ]; then
            #remove drift, this artificially inflate your SD, make SD look bigger than what it is
            [ ! -r $pdir/tsnr/$prefix.det.nii.gz ] && 3dDetrend -prefix $pdir/tsnr/$prefix.det.nii.gz -polort 4 $pdir/$prefix.nii.gz"[$DROPVOLS..$]"
            #calcualte SD on detrended data
            [ ! -r $pdir/tsnr/$prefix.det.stdev.nii.gz ] && 3dTstat -stdev -prefix $pdir/tsnr/$prefix.det.stdev.nii.gz $pdir/tsnr/$prefix.det.nii.gz

            # need to put back mean, use norm tmean -- this may have been calculated earlier
            normprefix=n${prefix#*n}     # e.g. nswudktm_func_5 from brgnswudktm_func_5
            [ ! -r $pdir/tsnr/$normprefix.tmean.nii.gz ] && 3dTstat -mean -prefix $pdir/tsnr/$normprefix.tmean.nii.gz $pdir/${normprefix}.nii.gz"[$DROPVOLS..$]" 

            [ ! -r $pdir/tsnr/$outname ] && 3dcalc -a $pdir/tsnr/${normprefix}.tmean.nii.gz -b $pdir/tsnr/$prefix.det.stdev.nii.gz -expr 'a/b' -float -prefix $pdir/tsnr/$outname
         else

            [ ! -r $pdir/tsnr/$prefix.tmean.nii.gz ] && 3dTstat -mean -prefix $pdir/tsnr/$prefix.tmean.nii.gz $pdir/$prefix.nii.gz"[$DROPVOLS..$]"
            [ ! -r $pdir/tsnr/$prefix.det.nii.gz ] && 3dDetrend -prefix $pdir/tsnr/$prefix.det.nii.gz -polort 4 $pdir/$prefix.nii.gz"[$DROPVOLS..$]"
            [ ! -r $pdir/tsnr/$prefix.det.stdev.nii.gz ] && 3dTstat -stdev -prefix $pdir/tsnr/$prefix.det.stdev.nii.gz $pdir/tsnr/$prefix.det.nii.gz
            [ ! -r $pdir/tsnr/$outname ] && 3dcalc -a $pdir/tsnr/$prefix.tmean.nii.gz -b $pdir/tsnr/$prefix.det.stdev.nii.gz -expr 'a / b' -float -prefix $pdir/tsnr/$outname
         fi
      fi
      echo "creating txt output '$txtout'"
      3dmaskave -quiet -mask "$gm_mask" "$tsnrout" |tee "$txtout"

      if [ ! -s "$txtout" ]; then 
         echo -e "FAILED "$txtout" is empty!\n\t3dmaskave -quiet -mask '$gm_mask' '$tsnrout'" >&2
         [ -n "$ONLYONE" ] && return 1
      fi

   done

   # no need to run count check if we ran for only one and we have it
   [ -n "$ONLYONE" ] && return 0 

   # return success if we have as many txt files as expected
   txtcnt=$(ls $pdir/tsnr/*.txt|wc -l)
   echo "# expect $cnt tsnr/*.txt files have $txtcnt"
   [ $cnt -eq $txtcnt ] && return 0
   return 1
}
parse_args(){
   while [ -n "$1" ]; do
      case "$1" in
         # files
         -t |-temp )   TEMPLATE="$2"; shift 2;;
         -g |-gm_mni ) GM_MNI="$2"; shift 2;;
         -s |-start )  STARTFILE="$2"; shift 2;;
         -O |-only)    STARTFILE="$2"; ONLYONE=1; shift 2;;
         -w |-warp )   WARPCOEF="$2"; shift 2;;
         # settings
         -x)   DROPVOLS="$2"; shift 2;;
         # flags
         -c |-no-clean)  CLEANUP="no"; shift 1;;
         -f |-no-flag) CHECKFLAG="no"; shift 1;;
         -n ) IS_WARP=0; shift;; 
         -v ) let ++VERBOSE; shift;;
         # help
         -h ) usage "";;
         -*) usage "unkown option: $1";;
         # non switch arg is preprocess directory
         *) PDIR="$1"; shift; echo setting PDIR to $PDIR;;
      esac
   done

   # if we have no input, input is the current directory
   [ -z "$PDIR" ] && PDIR="$(pwd)"
   #[ -f "$PDIR" ] && PDIR=$(dirname "$INPUT") && STARTFILE="$(basename $PDIR)"
   [ ! -d "$PDIR" ] && usage "bad preprocdir: $PDIR"

   for v in TEMPLATE GM_MNI WARPCOEF; do
      f="${!v}"
      [ ! -r "$f" ] && continue
      fd="$(cd $(dirname $f);pwd)"
      fb="$(basename $f)"
      printf -v $v "$fd/$fb"
   done

   # files can be relative to rundir so don't panic if we cannot find them
   # but make absolute path if they are relative to exec dir
   if [ $VERBOSE -gt 1 ]; then
      for v in PDIR TEMPLATE GM_MNI WARPCOEF STARTFILE CHECKFLAG CLEANUP VERBOSE; do
         echo "$v: ${!v}"
      done
   fi
}


parse_args $@
[ $VERBOSE -ge 3 ] && set -x
# rundir is dir of input if input is a file, otherwise rundir is input
# setup other options (cleanup, flag, etc)

# run and cleanup if needed
TSNR_onerun && 
 [ "$CLEANUP" = yes ] &&
 find $PDIR/tsnr/ -maxdepth 1 -type f -iname "*.nii.gz" -not -iname "*_tsnr.nii.gz" -delete