#!/bin/tcsh -e

# !!!!!?   Why use f00*, f01*, f02*, etc.? is iidx instead?

# written by PA Taylor (NIMH, NIH, USA)
# started March, 2017

# --------------------- revision history -------------------------
#
#set version   = "1.0";    set rev_dat   = "March 10, 2017"
#   + start: align t1 to t2
#
#set version   = "1.1";    set rev_dat   = "March 11, 2017"
#   + change name...
#   + help text added
#   + no-clean option added
#
#set version   = "1.2";    set rev_dat   = "Apr 12, 2017"
#   + standardize I/O, helpfile, option names
#
#set version   = "1.3";    set rev_dat   = "Apr 26, 2017"
#   + newer wdir convention
#
#set version   = "1.4";    set rev_dat   = "May 1, 2017"
#   + Align t2w -> t1w and reverse (better masking/weighting)
#
#set version   = "1.5";    set rev_dat   = "June 20, 2017"
#   + new opt: can have mask for t2w vol
#
#set version   = "1.6";    set rev_dat   = "July 25, 2017"
#   + new opt: can have more than warp = shift+rot only; default is
#   just that, still
#
#set version   = "1.7";    set rev_dat   = "Aug 11, 2017"
#   + switch a ">>" to ">" for general system compatability
#
#set version   = "1.8";    set rev_dat   = "Sep 04, 2017"
#   + update to work with newer @chauffeur*, -prefix only
#
set version   = "1.9";    set rev_dat   = "Jan 12, 2018"
#   + wt vol into working dir
#
# ---------------------------------------------------------------

set this_prog = "fat_proc_align_anat_pair"
set tpname    = "${this_prog:gas/fat_proc_//}"
set here      = $PWD

# ----------------- find AFNI and set viewer ---------------------

# find AFNI binaries directory and viewer location
set adir      = ""
set my_viewer = ""
which afni >& /dev/null
if ( $status ) then
    echo "** Cannot find 'afni' (?!)."
    goto BAD_EXIT
else
    set aa   = `which afni`
    set adir = $aa:h
endif

# default location of viewer: user could modify!
set my_viewer = "$adir/@chauffeur_afni"

# ----------------------- set defaults --------------------------

set it1w       = ""
set it2w       = ""
set imtrx      = ""              # user can input 12 DOF made elsewhere
set it2w_mask  = ""

# def: do prep t1w for freesurfing -> matrix has even numbered dims ->
# 1mm iso voxels, or other iso
set DO_PREP_FS = 1               

set RES        = ""              # def: 1mm output res, or smallest vox dim
set OUT_T2W_GRID = 0

set odir       = ""
set opref      = ""

set DO_VIEWER  = 1               # def: do viewing
set DO_CLEAN   = "1"
set wdir       = "__WORKING_$tpname"
set WDIR_EX  = "1"               # put opref on wdir (unless user names)
set output_cmd = 1               # def: output copy of this command
set cmd_file   = ""              # def: same name as viewer
set qc_prefix  = ""              # def: autoname; user can enter
set postfix    = "_"             # stick into name
set warp_opt   = "-warp shift_rotate" # default, but can be changed

# ------------------- process options, a la rr ----------------------

if ( $#argv == 0 ) goto SHOW_HELP

set ac = 1
while ( $ac <= $#argv )
    # terminal options
    if ( ("$argv[$ac]" == "-h" ) || ("$argv[$ac]" == "-help" )) then
        goto SHOW_HELP
    endif
    if ( "$argv[$ac]" == "-ver" ) then
        goto SHOW_VERSION
    endif

    # --------------- input dset(s) ----------------

    if ( "$argv[$ac]" == "-in_t1w" ) then
        if ( $ac >= $#argv ) goto FAIL_MISSING_ARG
        @ ac += 1
        set it1w = "$argv[$ac]"

    else if ( "$argv[$ac]" == "-in_t2w" ) then
        if ( $ac >= $#argv ) goto FAIL_MISSING_ARG
        @ ac += 1
        set it2w = "$argv[$ac]"

    else if ( "$argv[$ac]" == "-in_t2w_mask" ) then
        if ( $ac >= $#argv ) goto FAIL_MISSING_ARG
        @ ac += 1
        set it2w_mask = "$argv[$ac]"

    # [PT: July 25, 2017] works like opt of the same name in
    # 3dAllineate; takes shift_only, shift_rotate, etc.
    else if ( "$argv[$ac]" == "-warp" ) then
        if ( $ac >= $#argv ) goto FAIL_MISSING_ARG
        @ ac += 1
        set warp_opt = "-warp $argv[$ac]"

    # optional input: premade 3dAllineate matrix
    else if ( "$argv[$ac]" == "-matrix" ) then
        if ( $ac >= $#argv ) goto FAIL_MISSING_ARG
        @ ac += 1
        set imtrx = "$argv[$ac]"

    # ----------------- outputs ---------------------

    else if ( "$argv[$ac]" == "-prefix" ) then
        if ( $ac >= $#argv ) goto FAIL_MISSING_ARG
        @ ac += 1
        set opref = "$argv[$ac]"

    # ------------------- other opts ---------------

    else if ( "$argv[$ac]" == "-workdir" ) then
        if ( $ac >= $#argv ) goto FAIL_MISSING_ARG
        @ ac += 1
        set wdir = "$argv[$ac]"
        set WDIR_EX  = "0"

    else if ( "$argv[$ac]" == "-no_cmd_out" ) then
        set output_cmd = 0

    # set output resolution
    else if ( "$argv[$ac]" == "-newgrid" ) then
        if ( $ac >= $#argv ) goto FAIL_MISSING_ARG
        @ ac += 1
        set RES = "$argv[$ac]"

    # set output resolution by using t2w grid
    else if ( "$argv[$ac]" == "-out_t2w_grid" ) then
        set OUT_T2W_GRID = 1

    # *don't* prep for FS. Why?
    else if ( "$argv[$ac]" == "-no_fs_prep" ) then
        set DO_PREP_FS= 0

    else if ( "$argv[$ac]" == "-no_clean" ) then
        set DO_CLEAN = "0"

    else
        echo "** unexpected option #$ac = '$argv[$ac]'"
        goto BAD_EXIT

    endif
    @ ac += 1
end

# =======================================================================
# ============================ ** SETUP ** ==============================
# =======================================================================

# ============================ input files ==============================

echo "++ Start script version: $version"

# NEED these two inputs
if ( "$it1w" == "" ) then
    echo "** ERROR: no t1w file input?!"
    goto BAD_EXIT
else if ( "$it2w" == "" ) then
    echo "** ERROR: no t2w file input?!"
    goto BAD_EXIT
endif

# make sure we can read volumes OK
foreach ff ( "$it1w" "$it2w" )
    set check = `3dinfo "$ff"`
    if ( "$#check" == "0" ) then
        echo "** ERROR: can't find input file:  $ff !"
        goto BAD_EXIT
    else
        echo "++ Found input file:   $ff"
    endif
end

# ===================== output and working dirs =========================

if ( "$opref" == "" ) then
    echo "** ERROR: need '-prefix ...' option provided!"
    echo "   See the helpfile for more information."
    goto BAD_EXIT
else
    set odir = `dirname $opref`
    set opref = `basename $opref`
    echo ""
    echo "++ Based on prefix, the output directory will be:"
    echo "     $odir"
    echo "++ Based on prefix, the output prefix will be:"
    echo "     $opref"
    echo ""
endif

# check output directory, use input one if nothing given

# default output dir, if nothing input.
if ( ! -e "$odir" ) then
    echo "+* Output directory didn't exist.  Trying to make it now."
    mkdir "$odir"
endif

# and put working directory as subdirectory.
if ( "$WDIR_EX" == "1" ) then
    set wdir = $odir/${wdir}_$opref
else
    set wdir = $odir/$wdir
endif

# make the working directory
if ( ! -e $wdir ) then
    echo "++ Making working directory: $wdir"
    mkdir $wdir
else
    echo "+* WARNING: Somehow found a premade working directory (?):"
    echo "      $wdir"

    # don't clean preexisting directories-- could be user mistake.
    echo "   NB: will *not* clean it afterwards."
    set DO_CLEAN = "0"
endif

# --------------------- output res ---------------------------

# If the user doesn't input a voxel resolution for output, then the
# spatial resolution will be 1mm iso, or the smallest dimension of an
# input voxel (so the user doesn't lose resolution by default).

if ( "$RES" == "" ) then

    set RES = "1.0"
    set dim = `3dinfo -ad3 "$it1w"`

    foreach dd ( $dim )
        set comp = `echo "$RES > $dd" | bc`
        if ( "$comp" != "0" ) then
            echo "++ Found smaller vox dim: $dd < ${RES} !"
            set RES = "$dd"
        endif
    end

    echo "++ Note: final volume grid resolution will be $RES mm isotropic."

endif

# ========================= output fnames ==========================

set ocmd   = "${opref}_cmd.txt"      # name for output command
set ot1w   = "${opref}.nii.gz"       # name for output t1w
set omtrx  = "${opref}_map_anat.aff12.1D" # save matrix with other files

# =======================================================================
# =========================== ** PROCESS ** =============================
# =======================================================================

echo "\n-----> STARTING $this_prog ---->"

# ---------------------------- CMD ---------------------------------

echo "\n\nThis command:"
echo "$this_prog $argv\n\n"

if ( "$cmd_file" == "" ) then
    set cmd_file = "$odir/$ocmd"
endif

# copy original command:
# dump copy of command into workdir/..
if ( $output_cmd == 1 ) then
    echo "++ Echoing the command to: $cmd_file"

    set rec_afni_ver = `afni -ver`
    echo "### AFNI version:"  > $cmd_file
    echo "# $rec_afni_ver\n"            >> $cmd_file

    echo "### Executed from the directory location:"  >> $cmd_file
    echo "# $here\n"            >> $cmd_file
    echo "### The command was:" >> $cmd_file
    echo "# $this_prog $argv"   >> $cmd_file
    echo "\n"                   >> $cmd_file
endif

# --------------------- start proc ---------------------------

# copy the t1w file to wdir via resampling: match t2w-t1w orients
set orient_t2w = `3dinfo -orient "$it2w"`
set t1res      = $wdir/f00_t1res.nii
set T1TOMAP    = $t1res                 # we will apply map to this!
3dresample                          \
    -overwrite                      \
    -prefix $t1res                  \
    -orient $orient_t2w             \
    -inset "$it1w"

# prep this side for aligning with a bit of unifizing and brightness
# thresholding
echo "++ Unifize T1w proc vol."
set t1uni = $wdir/f01_t1uni.nii
3dUnifize                            \
    -GM                              \
    -prefix  $t1uni                  \
    -input   "$t1res"                \
    -overwrite 

echo "++ Threshold T1w proc vol."
set t1thr = $wdir/f02_t1thr.nii
set v1 = `3dBrickStat -non-zero -automask   \
            -percentile 95 1 95 ${t1uni}`
echo ${v1[2]}
3dcalc -a  ${t1uni}                         \
       -expr "maxbelow(${v1[2]},a)"         \
       -prefix $t1thr                       \
       -overwrite

set t1align   = $wdir/f05_t1align.nii
if ( "$imtrx" == "" ) then

    echo "++ Calculating alignment (t2w -> t1w) matrix."

    set mt_t2_t1 = $wdir/map_t2w_to_t1w.aff12.1D
    set mt_t1_t2 = $wdir/map_t1w_to_t2w.aff12.1D

    set imtrx = $wdir/map_to_t2w.aff12.1D

    if ( "$it2w_mask" == "" ) then
        set it2w_mask = "-source_automask"
    else
        set it2w_mask = "-source_mask $it2w_mask"
    endif

    3dAllineate   -echo_edu               \
        -1Dmatrix_save  $mt_t2_t1         \
        -prefix         "$t1align"        \
        -source         "$it2w"           \
        -base           "$t1thr"          \
        -twopass                          \
        -cost lpc                         \
        -cmass                            \
        $warp_opt                         \
        $it2w_mask                        \
        -wtprefix      $wdir/wtvol.nii.gz \
        -autoweight                       \
        -final wsinc5                     \
        -overwrite

    # get inverse
    echo "++ Inverting alignment matrix (to get t1w -> t2w)."
    cat_matvec -ONELINE                   \
        $mt_t2_t1 -I                      \
        > $mt_t1_t2
    
    set imtrx = $mt_t1_t2
endif

echo "++ Apply alignment matrix."

set t1mapped = $wdir/f06_t1mapped.nii
if ( $OUT_T2W_GRID == 1 ) then
    echo "++ Outputting final image to t2w space"
    3dAllineate   -echo_edu               \
        -1Dmatrix_apply   $imtrx          \
        -source           $T1TOMAP        \
        -master           "$it2w"         \
        -prefix           $t1mapped       \
        -float                            \
        -final wsinc5

else
    echo "++ Outputting final image to $RES mm iso voxels"
    3dAllineate   -echo_edu               \
        -1Dmatrix_apply   $imtrx          \
        -source           $T1TOMAP        \
        -prefix           $t1mapped       \
        -base             "$it2w"         \
        -float                            \
        -newgrid         "$RES"           \
        -final wsinc5
endif

set T1TOCOPY = $t1mapped
\cp $imtrx $odir/$omtrx

# zeropad if needed
if ( $DO_PREP_FS == 1) then

    set t1padded = $wdir/f08_t1padded.nii

    echo "++ Perhaps for prepping for FreeSurfing:"
    echo "    -> making sure that there are even numbers of voxels"
    echo "       along each spatial dimension."

    set dims = `3dinfo -n4 ${t1mapped}`
    echo "++ The dimensions are:  $dims[1] $dims[2] $dims[3]"

    set odds = ( 0 0 0 )

    # silly stuff to deal with orientation
    set ori  = `3dinfo -orient ${t1mapped}`
    set orilet = ( "" "" "" )
    set orilet[1] = `echo $ori | awk '{print substr($0,1,1)}'`
    set orilet[2] = `echo $ori | awk '{print substr($0,2,1)}'`
    set orilet[3] = `echo $ori | awk '{print substr($0,3,1)}'`

    set build_cmd = "" 

    foreach i ( `seq 1 1 3` ) 
        set odds[$i] = `echo "$dims[$i] % 2" | bc`
        if ( $odds[$i] != 0 ) then
            echo "+* Found an odd dimension in the [${i}]th axis!"
            echo "   --> will hopefully fix."

            set build_cmd = "$build_cmd -${orilet[$i]} 1 "
        endif
    end

    if ( "$build_cmd" != "" ) then
        
        echo "++ Zeropadding:  $build_cmd"
        3dZeropad     -echo_edu                 \
            -overwrite                          \
            -prefix $t1padded                   \
            $build_cmd                          \
            $t1mapped

        set T1TOCOPY = $t1padded
    endif

endif

set fdims = `3dinfo -n4 ${T1TOCOPY}`
set fres  = `3dinfo -ad3 ${T1TOCOPY}`

echo "++ Done workin'.  Just need to copy the final dset."

3dcalc                         \
    -overwrite                 \
    -a $T1TOCOPY               \
    -expr 'a'                  \
    -prefix $odir/$ot1w

if ( $DO_VIEWER == 1 ) then

    set vedge = $wdir/f10_edges.nii

    echo "++ More QC images: b0 on initial ref."
    3dedge3                         \
        -overwrite                  \
        -prefix $vedge              \
        -input  $odir/$ot1w

    if( $qc_prefix == "" ) then
        set vpref0 = ${opref}${postfix}_qc00_t2w_u_et1w
        set vpref1 = ${opref}${postfix}_qc01_t2w_u_t1w
    else
        set vpref0 = ${qc_prefix}${postfix}_qc00_t2w_u_et1w
        set vpref1 = ${qc_prefix}${postfix}_qc01_t2w_u_t1w
    endif

    echo "\n\n"
    echo "++ QC image 00 ($odir/$ot1w edges on $it2w): $vpref0"
    echo "\n\n"
    # need to put '[0]' on $iref?
    $my_viewer                            \
        -ulay "$it2w"                     \
        -ulay_range "2%" "98%"            \
        -olay "$vedge"                    \
        -func_range_perc 50               \
        -pbar_posonly                     \
        -cbar "red_monochrome"            \
        -opacity 6                        \
        -prefix "$odir/$vpref0"           \
        -montx 5 -monty 3                 \
        -set_xhairs OFF                   \
        -label_mode 1 -label_size 3       \
        -do_clean 

    echo "\n\n"
    echo "++ QC image 01 ($odir/$ot1w olay on $it2w): $vpref1"
    echo "\n\n"

    $my_viewer                            \
        -ulay "$it2w"                     \
        -ulay_range "2%" "98%"            \
        -olay "$odir/$ot1w"               \
        -pbar_posonly                     \
        -opacity 4                        \
        -prefix "$odir/$vpref1"           \
        -montx 5 -monty 3                 \
        -set_xhairs OFF                   \
        -label_mode 1 -label_size 3       \
        -do_clean 

endif

# clean, by default
if ( "$DO_CLEAN" == "1" ) then
    echo "\n++ Cleaning working directory!\n"
    \rm -rf $wdir
else
    echo "\n++ NOT removing working directory '$wdir'.\n"
endif

# final messages
echo ""
echo "++ The final data set is here:  $odir/$ot1w"
echo "++ The final spatial dims are:  $fdims[1] $fdims[2] $fdims[3]"
echo "++ The spatial resolution is:   $fres[1] $fres[2] $fres[3]"
echo ""

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

goto GOOD_EXIT

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

SHOW_HELP:
cat << EOF
# -----------------------------------------------------------------------

 This program is for aligning a T1w anatomical to a T2w anatomical
 using solid body parameters (i.e., only translation and rotation);
 this program does not clean or alter the T1w's brightness values
 (beyond minor smoothing from regridding).  If one is going to be
 entering the T1w volume into the realms of FreeSurfer (FS), one might
 want to do this just *before* that step.  If one wants axialized (or
 AC-PC-ized) anatomical volumes, one could perform that step on the
 T2w volume *before* using this function.

 This program mainly assumes that the T1w and T2w volume come from the
 same subject, and have similar contrasts expected for standard
 sequences and healthy adult brains.  This might still work for other
 applications, but caveat emptor (even more than usual!).  This would
 *not* be recommended for aligning brains that aren't from the same
 subject.

 As part of this alignment, the T1w volume will end up with the same
 orientation and a similar FOV as the T2w volume.  Additionally, by
 default, the anatomical will be prepped a bit with an eye toward
 using FS, to have properties favorable to using it: 

   + the T1w volume is resampled to isotropic spatial resolution of
     either 1 mm voxel edges or, if the input volume has any edge
     length smaller than this, to that value (i.e., resampled to 1 mm
     or the minimum input voxel edge length, whichever is less).  The
     user can adjust this with the '-newgrid ...' option, or decide to 
     match the grid of the T2w volume via '-out_t2w_grid'. 

   + the T1w will have a FOV matching or quite similar to the T2w
     volume (as well as matching orientation).

   + the output T1w will be checked to ensure having an even number of
     rows along each axis; this is because FS somehow seems to want
     this for best behavioring. A slice of zeros will be appended in
     any direction with with an odd number of slices; turned off
     by using '-no_fs_prep'.

 Note that, if you are preparing to use FS afterwards, then make sure
 to use their current help files, archives, etc. for all options and
 settings.  For example, while at present (March, 2017) FS does seem
 to prefer isotropic voxels with 1 mm edge length by default, one can
 use high resolution options for data acquired at higher resolution.
 Anyways, you can read more about that there.

  Ver. $version (PA Taylor, ${rev_dat})

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

  OUTPUT:

     + NIFTI file: aligned T1w volume

     + QC snapshots of the T1w volume overlaying the T2w volume, and
       also the T1w edges overlaying the T2w volume.

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

  RUNNING:

    $this_prog \
      -in_t1w  T1W                          \
      -in_t2w  T2W                          \
      -prefix  PPP                          \
      {-newgrid RES}                        \
      {-out_t2w_grid}                       \
      {-in_t2w_mask MASK_T2W}               \
      {-matrix MMM}                         \
      {-workdir WWW}                        \
      {-no_fs_prep}                         \
      {-warp WAR}                           \
      {-no_cmd_out}                         \
      {-no_clean} 

  where:

   -in_t1w  T1W   :T1w volume (required).
   -in_t2w  T2W   :T2w volume (required; preferably from same subject as
                   T1W).

   -prefix  PPP   :output prefix for files and snapshots (required).

   -newgrid RES   :specify output T1w volume's final resolution; will be
                   isotropic in this value (default: 1 mm, or smallest voxel
                   edge length of input T1W if that value is < 1 mm).
   -out_t2w_grid  :final T1w volume is on the T2W volume's grid (with 
                   possible addition of padding with a slice of zeros if
                   prepping for FS).

   -no_fs_prep    :check output T1w volume to have an even number of slices
                   in all directions, because FS appears to like/require this
                   (at least at this slicetime of the universe, March, 2017).

   -in_t2w_mask MASK_T2W
                  :can input a mask to apply to the t2w volume for
                   alignment purposes; might help in times of aligning 
                   hardship.

   -warp WAR      :can choose which of the possible affine degrees of freedom
                   are employed in the warping, selecting them in the same
                   manner described in 3dAllineate's help;  that is, WAR can
                   be any of shift_only, shift_rotate, shift_rotate_scale, or
                   affine_general.  Default: WAR = shift_rotate.

   -matrix MMM    :one can apply a pre-made matrix that has been made by
                   3dAllineate previously.  With this option.  If you want.

   -workdir WWW   :specify a working directory, which can be removed;
                   (default name = '$wdir')

   -no_cmd_out    :don't save the command line call of this program
                   and the location where it was run (otherwise, it is
                   saved by default in the ODIR/).                     
   -no_clean      :no not delete temporary working directory (default is to 
                   remove it to save disk space).

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

  EXAMPLE

    # have isotropic 1x1x1 mm final anat:
    $this_prog  \
        -in_t1w    MPRAGE.nii.gz        \
        -in_t2w    T2w_anat.nii.gz      \
        -newgrid   1.0                  \
        -prefix    t1w_alnd

    # match the final anat resolution to that of the t2w dset:
    $this_prog  \
        -in_t1w    MPRAGE.nii.gz        \
        -in_t2w    T2w_anat.nii.gz      \
        -out_t2w_grid                   \
        -prefix    t1w_alndb


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

EOF

    goto GOOD_EXIT

SHOW_VERSION:
    echo "version  $version (${rev_dat})"
    goto GOOD_EXIT

FAIL_MISSING_ARG:
    echo "** ERROR! Missing an argument after option flag: '$argv[$ac]'"
    goto BAD_EXIT

BAD_EXIT:
    exit 1

# send everyone here, in case there is any cleanup to do
GOOD_EXIT:
    exit 0
