#!/usr/bin/env tcsh

@global_parse `basename $0` "$*" ; if ($status) exit 0

#set version   = "0.0";  set rev_dat   = "June 3, 2022"
# [PT]  creation of this script, which (at present) is a useful
#       precursor for EPI-anatomical alignment.
#
#set version   = "1.0";  set rev_dat   = "June 3, 2022"
# [PT]  add help file
#
#set version   = "1.1";  set rev_dat   = "July 4, 2022"
# [PT]  change how masking is turned off
#
#set version   = "1.2";  set rev_dat   = "Jan 29, 2024"
# [PT]  change how local_mask is parsed, so '-mask DSET' is 
#       properly handled with an initial dset copy to the wdir
#
set version   = "1.3";  set rev_dat   = "May 11, 2025"
# [PT]  add -sharpen_with_phi, to possibly insert 3dSharpen into
#       the processing stream
#
set version   = "1.4";  set rev_dat   = "July 2, 2025"
# [PT]  add -automask_opts and -automask_tool options, to be able
#       to be able to control automasking (which now also happens 
#       as its own separate step) with more refinement, if 
#       necessary
#
# ----------------------------------------------------------------

set abbrev = "LocalUni"

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

set input     = ""
set opref     = ""
set odir      = "."

set pppp      = "`3dnewid -fun11`"
set wdir_base = "__wdir_${abbrev}"
set wdir_name = "${wdir_base}_${pppp}"

set local_perc = 50
set local_rad  = -3
set local_mask = "-automask"
set sharp_phi  = ""

set amask_opts = ""
set amask_tool = ""

set filt_thr   = 1.5

set overwrite  = ""

set DO_CLEAN  = 1                       # default: keep working dir

# ------------------- 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

    if ( "$argv[$ac]" == "-echo" ) then
        set echo

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

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

    # ----------- 3dLocalStat opts

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

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

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

    else if ( "$argv[$ac]" == "-local_mask" ) then
        if ( $ac >= $#argv ) goto FAIL_MISSING_ARG
        @ ac += 1
        set local_mask = "$argv[$ac]"
        # special case to turn off masking
        if ( "${local_mask}" == "None" ) then
            set local_mask = ""
        endif

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

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

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

        set check     = `echo "${wdir_name}" | grep '/'`
        if ( "${check}" != "" ) then
            echo "\n\n** ERROR: -wdir_name arg must not have '/':"
            echo "   '${wdir_name}'\n\n"
            goto BAD_EXIT
        endif

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

    else if ( "$argv[$ac]" == "-overwrite" ) then
        set overwrite = "-overwrite"

    else if ( "$argv[$ac]" == "-no_clean" ) then
        set DO_CLEAN = 0
        
    else
        echo "\n\n** ERROR: unexpected option #$ac = '$argv[$ac]'\n\n"
        goto BAD_EXIT
        
    endif
    @ ac += 1
end

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

if ( "${input}" == "" ) then
    echo "** ERROR: missing input file. Use '-input ..'"
    goto BAD_EXIT
endif

if ( "${opref}" == "" ) then
    echo "** ERROR: missing output prefix. Use '-prefix ..'"
    goto BAD_EXIT
endif

if ( "${sharp_phi}" != "" ) then
    set is_ok = `ccalc -i -eval "within(${sharp_phi},0.1,0.9)"`
    if ( ! ${is_ok} ) then
        echo "** ERROR: Sharpen phi must be in range [0.1, 0.9], not: ${sharp_phi}"
        echo "          Please fix your input to '-sharpen_with_phi ..'"
        goto BAD_EXIT
    endif
endif

# ===================== output dir + wdir =======================
 
echo "++ Start 3dLocalUnifize work"

# check output directory, use input one if nothing given

if ( ! -e "${odir}" ) then
    echo "++ Making new output directory: ${odir}"
    \mkdir -p "${odir}"
endif

set wdir = "${odir}/${wdir_name}"

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

# check about whether local_mask opt uses a mask dset; if so, copy
# that mask to wdir and change var to use that local copy
set nlm = `python -c "print(len('${local_mask}'.split()))"`
if ( "${nlm}" == "2" ) then
    # does local_mask start with '-mask'?
    set lll = `echo "${local_mask}" | awk '{print $1}'`
    if ( "${lll}" == "-mask" ) then
        # if here, we have a mask dset input

        # get name of mask; careful of subbrick selectors
        set mmm = `echo "${local_mask}" | awk '{print $2}'`

        # copy the dset
        3dTcat ${overwrite}                                         \
            -prefix   "${wdir}/dset_00_mask.nii.gz"                 \
            "${mmm}"

        if ( $status ) then
            echo "** ERROR when processing local_mask opt: '${local_mask}'"
            goto BAD_EXIT
        endif
            
        # edit the local_mask opt to use the local wdir mask name
        set local_mask = "-mask dset_00_mask.nii.gz"
    endif
endif


# ======================= main program work ===========================

# Copy original image in

set idset    = "${input}"
set odset_cp = dset_00_cp.nii.gz

3dTcat ${overwrite}                                                 \
    -prefix   "${wdir}/${odset_cp}"                                 \
    "${idset}"

# move into working dir
cd "${wdir}"

# if automasking, calculate separately first
if ( "${local_mask}" == "-automask" ) then

    3dAutomask ${overwrite}                                         \
        ${amask_opts}                                               \
        -prefix  "dset_00_mask.nii.gz"                              \
        "${odset_cp}"

    if ( $status ) then
        echo "** ERROR when automasking"
        goto BAD_EXIT
    endif
        
    # edit the local_mask opt to use the local wdir mask name
    set local_mask = "-mask dset_00_mask.nii.gz"

    if ( "${amask_tool}" != "" ) then
        3dmask_tool ${overwrite}                                  \
            ${amask_tool}                                         \
            -input     dset_00_mask.nii.gz                        \
            -prefix    dset_00_maskB.nii.gz 

        if ( $status ) then
            echo "** ERROR when running 3dmask_tool on the automask"
            goto BAD_EXIT
        endif

        # edit the local_mask opt to use the local wdir mask name
        set local_mask = "-mask dset_00_maskB.nii.gz"
    endif
endif


# Calc the smooth/median image

set idset    = "${odset_cp}"
set odset_lu = dset_01_lu.nii.gz

3dLocalstat ${overwrite}                                            \
    ${local_mask}                                                   \
    -nbhd     "SPHERE(${local_rad})"                                \
    -stat     "perc:${local_perc}:${local_perc}:1"                  \
    -prefix   "${odset_lu}"                                         \
    "${idset}"

# Calc the ratio of original to this

set idset     = "${odset_lu}"
set jdset     = "${odset_cp}"
set odset_rat = dset_02_rat.nii.gz

3dcalc ${overwrite}                                                \
    -a       "${idset}"                                            \
    -b       "${jdset}"                                            \
    -expr    "b/a"                                                 \
    -prefix  "${odset_rat}"

set dset_out = "${odset_rat}"

# (opt) Sharpen

if ( "${sharp_phi}" != "" ) then
    echo "++ Apply sharpening to image with phi: ${sharp_phi}"

    set idset       = "${odset_rat}"
    set odset_sharp = dset_03_sharp.nii.gz

    if ( "${overwrite}" != "" && -f ${odset_sharp} ) then
        \rm ${odset_sharp}
    endif

    3dSharpen                                                      \
        -phi     ${sharp_phi}                                      \
        -input   "${idset}"                                        \
        -prefix  "${odset_sharp}"

    set dset_out = "${odset_sharp}"
endif

# (opt) Filter possible outlier ratio values

if ( `echo "${filt_thr} > 0 " | bc` ) then
    echo "++ Apply filter threshold to scaled image: ${filt_thr}"

    set idset      = "${dset_out}"
    set odset_filt = dset_04_filt.nii.gz

    3dcalc ${overwrite}                                            \
        -a       "${idset}"                                        \
        -expr    "maxbelow(${filt_thr},a)"                         \
        -prefix  "${odset_filt}"

    set dset_out = "${odset_filt}"
else
    echo "++ NO filter threshold will be applied to the scaled image"

endif

# ----------- finish

3dcopy ${overwrite}                                               \
    "${dset_out}"                                                 \
    "../${opref}"

# ------------------------- wrap up ----------------------------

# move out of wdir to the odir
cd ..
set whereout = "$PWD"

if ( "$DO_CLEAN" == "1" ) then
    echo "\n+* Removing the working dir: '${wdir_name}'\n"
    \rm -rf "${wdir_name}"
else
    echo "\n++ NOT removing the working dir: '${wdir_name}'\n"
endif

echo ""
echo "++ DONE! View the finished, locally-unifized product:"
echo "     ${whereout}/${opref}"
echo ""


goto GOOD_EXIT

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

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

OVERVIEW ~1~

This program takes an input and generates a simple "unifized" output
volume.  It estimates the median in the local neighborhood of each
voxel, and uses that to scale each voxel's brightness.  The result is
a new dataset of brightness of order 1, which still has the
interesting structure(s) present in the original.

This program's output looks very useful to help with dataset alignment
(esp. EPI-to-anatomical) in a wide array of cases.

ver   : ${version}
date  : ${rev_dat}
auth  : PA Taylor (SSCC, NIMH, NIH)


USAGE ~1~

This program is generally run as:

    3dLocalUnifize [options] -prefix DSET_OUT -input DSET_IN

where the following options exist:

  -input     DSET_IN  :(req) input dataset
  
  -prefix   DSET_OUT  :(req) output dataset name, including path

  -wdir_name WD       :name of temporary working directory, which 
                       should not contain any path information---it will be
                       created in the same directory as the final dataset
                       is created
                       (def: ${wdir_base}_, plus a random alphanumeric str)

  -echo               :run this program very verbosely (def: don't do so)

  -no_clean           :do not remove the working directory (def: remove it)

 ... and the following are 'tinkering' options, likely not needed in
     most cases:

  -local_rad LR       :the spherical neighborhood's radius for the 
                       3dLocalStat step (def: ${local_rad})

  -local_perc LP      :the percentile used in the 3dLocalStat step,
                       generating the scaling volume 
                       (def: ${local_perc})

  -local_mask LM      :provide the masking option to be used in the
                       3dLocalStat step, which should be enclosed in
                       quotes for passing along to the internal
                       program call.  So, to use a pre-existing mask,
                       you might call this option like:
                         -local_mask "-mask my_mask.nii.gz"
                       To remove any masking, put the special keyword
                       "None" as the option value.
                       (def: "${local_mask}")

  -automask_opts 'AO ...' :when automasking is applied ('-local_mask -automask' 
                       or, since that is the default, not adding any masking 
                       opt), users can use this option to provide one or more
                       options to 3dAutomask when it is run 
                       (def: "${amask_opts}").

  -automask_tool 'AT ...' :when automasking is applied ('-local_mask -automask' 
                       or, since that is the default, not adding any masking 
                       opt), users can use this option to run 3dmask_tool on that
                       initial automask, providing one or more options to 
                       3dmask_tool when it is run (def: "${amask_tool}").

  -sharpen_with_phi PHI :before possible filtering and output, apply 3dSharpen
                       to sharpen the locally unifized result.  The value of
                       PHI must be in range [0.1, 0.9]; this sets the 
                       value of the '-phi ..' opt in 3dSharpen---see that 
                       program's help for details (def: no sharpening applied)

  -filter_thr FT      :put a ceiling on values in the final, scaled dataset,
                       whose values should be of order 1; setting FT to be a
                       value <=0 turns off this final filtering
                       (def: ${filt_thr})


NOTES ~1~

  This program is designed to not need a lot of tinkering with
  options, such as the '-local_* ..' ones.  In most cases, the default
  scaling will be useful.


EXAMPLES ~1~

  1. Basic local unifizing:
     3dLocalUnifize                                                  \\
        -prefix  vr_base_LU                                          \\
        -input   vr_base_min_outlier+orig.HEAD

  2. Same as above, without masking:
     3dLocalUnifize                                                  \\
        -prefix      vr_base_LU_FOV                                  \\
        -input       vr_base_min_outlier+orig.HEAD                   \\
        -local_mask  None

  3. Same as Ex. 1, but adding sharpening:
     3dLocalUnifize                                                  \\
        -prefix            vr_base_LU_SHARP.nii.gz                   \\
        -input             vr_base_min_outlier+orig.HEAD             \\
        -sharpen_with_phi  0.5

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

GOOD_EXIT:
    exit 0