#!/bin/bash
# Copyright(C) 1999-2020, 2023 National Technology & Engineering Solutions
# of Sandia, LLC (NTESS).  Under the terms of Contract DE-NA0003525 with
# NTESS, the U.S. Government retains certain rights in this software.
#
# See packages/seacas/LICENSE for details
########################################################################
function usage {
  if [ $OLDGETOPT -eq 1 ] ; then
    usage_old
  else
    usage_new
  fi
}

function usage_new {
cat <<DECOMP_USAGE_EOF
Usage:  decomp [-h] --processors <processor count> [-n <nemslice options>]
               [--root <root directory>] [--subdir <sub directory>] [--spread_suffix <suffix>]
               [--multikl] [--spectral] [--inertial] [--linear]
               [--rcb] [--rcb_ignore_z] [--rib] [--hsfc] [--random] [--scattered]
               [--weight <specification>] [...others...] meshfile.ext

Decomposition methods: [0..3] = speed ranking 0=fastest
              (See chaco and/or zoltan documentation for more details)
 --multikl:   [3] Multilevel KL. Gives high quality partitions.
                  Slow, high memory use.
 --spectral:  [2] use eigenvectors of a matrix constructed from the
                  graph to partition the graph.
 --hsfc       [?] Hilbert Space-Filling Curve (uses Zoltan)
 --rib:       [?] (default) Recursive Inertial Bisection (uses Zoltan)
 --rcb:       [?] Recursive Coordinate Bisection (uses Zoltan)
 --rcb_ignore_z:       [?] Recursive Coordinate Bisection (uses Zoltan), ignoring the z dimension coordinates
 --inertial:  [1] vertices divided into sets of equal mass by planes
                  orthogonal to the principle axis. Good, fast, medium memory
 --linear:    [0] vertices in order first n/p to proc 0, next to proc 1.
                  (fast and sometimes good).
 --random:    [0] vertices are assigned randomly to sets in a way that
                  preserves balance.
 --scattered: [0] vertices are handed out in order with the next vertex
                  going to whichever set is smallest.
 --brick      [-] special decomposition method; not useful for general use.
 --zpinch     [-] special decomposition method; not useful for general use.

 NOTE: The decomposition options can also be specified as "-l type"
       For example "--multikl" is the same as "-l multikl"

Other options:
--help              Print this message (-h)
--launch <cmd>      Prefix nem_slice call with this command to run it.
                    The cmd argument must be quoted if it contains options
                    for cmd.
--nolaunch          Prevent decomp from searching for launch script.
--noop              Only echo the nem_slice and nem_spread commands.
--nodal             Nodal Decomposition (-N)
--elemental         Elemental Decomposition (default, -E)
--input <file>      Optional nem_slice input file (-i)
--weight  <spec>    Specification of weighting options (-w)
--viz_mesh          Output mesh showing processor distribution of elements. (-V)
--spheres_linear    Decompose sphere elements using linear method (matches old behavior) (-s)
--processors <#p>   Specify number of processors (-p, -j)
--root <root_dir>   Root directory to begin the path to on the separate disks (-R)
--subdir <sub_dir>  Continuation of the path on the separate disks to the files (-S)
--spread_suffix <x> Suffix for parallel files (default: meshfile suffix) (-X)
--64                Force creation of 64-bit integer decomposed files (-6)
-n <options>        Specify additional nem_slice options

   Documentation: https://sandialabs.github.io/seacas-docs/sphinx/html/index.html#decomp"
   help-email: sierra-help@sandia.gov

DECOMP_USAGE_EOF
exit 1
}

function usage_old {
cat <<DECOMP_USAGE_EOF
Usage:  decomp [-h] -p <processor count> [-n <nemslice options>]
               [-R <root directory>] [-S <sub directory>] [-X <suffix>]
               [-l multikl] [-l spectral] [-l inertial] [-l linear]
               [-l rcb] [-l rcb_ignore_z] [-l rib] [-l hsfc] [-l random] [-l scattered]
               [...others...] meshfile.ext

Decomposition methods: [0..3] = speed ranking 0=fastest
              (See chaco and/or zoltan documentation for more details)
 -l multikl:   [3] Multilevel KL. Gives high quality partitions.
                  Slow, high memory use.
 -l spectral:  [2] use eigenvectors of a matrix constructed from the
                  graph to partition the graph.
 -l hsfc       [?] Hilbert Space-Filling Curve (uses Zoltan)
 -l rib:       [?] (default) Recursive Inertial Bisection (uses Zoltan)
 -l rcb:       [?] Recursive Coordinate Bisection (uses Zoltan)
 -l rcb_ignore_z:       [?] Recursive Coordinate Bisection (uses Zoltan) ignoring the z dimension coordinates
 -l inertial:  [1] vertices divided into sets of equal mass by planes
                  orthogonal to the principle axis. Good, fast, medium memory
 -l linear:    [0] vertices in order first n/p to proc 0, next to proc 1.
                  (fast and sometimes good).
 -l random:    [0] vertices are assigned randomly to sets in a way that
                  preserves balance.
 -l scattered: [0] vertices are handed out in order with the next vertex
                  going to whichever set is smallest.
 -l brick      [-] special decomposition method; not useful for general use.
 -l zpinch     [-] special decomposition method; not useful for general use.

Other options:
-N          Nodal Decomposition
-E          Elemental Decomposition (default)
-n options  Specify nem_slice options
-h          Print this message
-i          Optional nem_slice input file
-s          Decompose sphere elements using linear method (matches old behavior)
-p #p       Specify number of processors (or -j)
-R root_dir Root directory to begin the path to on the separate disks
-S sub_dir  Continuation of the path on the separate disks to the files
-X suffix   Suffix for parallel files (default: meshfile suffix)
-V          Output mesh showing processor distribution of elements
-6          Force creation of 64-bit integer decomposed files

   Documentation: https://sandialabs.github.io/seacas-docs/sphinx/html/index.html#decomp"
   help-email: sierra-help@sandia.gov

DECOMP_USAGE_EOF
exit 1
}

function old_getopt {
echo ${txtred}
cat <<DECOMP_EOF
########################################################################
The "getopt" executable that is available on this system is an older
version that is not compatible with the needs of the "decomp" tool.
If possible, you should update your getopt to a newer version and make
sure that the new getopt is in your path.

Below are some options for getting the current getopt version:
* If on a Mac: "sudo port install getopt" or "brew install gnu-getopt"
* Search the internet for "getopt-1.1.5" or "getopt-1.1.4"; download and build

Enter "-h" for the modified options that this version supports.
Enter "-H" for the options that the standard version supports.
########################################################################

DECOMP_EOF
echo ${txtrst}
}

########################################################################
function execute_loadbalance {
    basename=$1
    processors=$2
    decomp_method=$3
    nem_slice_flag=$4
    NEM_SLICE=$5
    input=$6
    LAUNCH=$7
    decomp_type=$8
    weighting=$9

    nemesis=$basename.$suffix_mesh.nem
    genesis=$basename.$suffix_mesh

    # Check for valid executable...
    if [ ! -x $NEM_SLICE ]
    then
        echo
        echo ${txtred}
        echo "ERROR: Could not find or execute $NEM_SLICE"
        echo "ERROR: Configuration error"
        echo ${txtrst}
        echo

        return -1
    fi

    date=$(date '+%m/%d/%y')
    time=$(date '+%H:%M:%S')
    prob_dir=$(pwd | sed "s/.*\///g")

    if [ -n "$input" ] && [ -e "$input" ]
    then
        echo ${txtblu}
        echo "Executing:"
        echo "   $LAUNCH $NEM_SLICE $decomp_type $spheres $do_viz $decomp_method $weighting $nem_slice_flag -o $nemesis -m mesh=$processors -a $input $genesis"
        echo ${txtrst}
        (${NOOP:+echo }$LAUNCH $NEM_SLICE $decomp_type $spheres $do_viz $decomp_method $weighting $nem_slice_flag -o $nemesis -m mesh=$processors -a $input $genesis)
        load_rc=$?
    else
        echo ${txtblu}
        echo "Executing:"
        echo "   $LAUNCH $NEM_SLICE $decomp_type $spheres $do_viz $decomp_method $weighting $nem_slice_flag -o $nemesis -m mesh=$processors $genesis"
        echo ${txtrst}
        (${NOOP:+echo }$LAUNCH $NEM_SLICE $decomp_type $spheres $do_viz $decomp_method $weighting $nem_slice_flag -o $nemesis -m mesh=$processors $genesis)
        load_rc=$?
    fi
    return $load_rc
}

########################################################################
# Create the nem_spread input file (basename.pex) and execute nem_spread...
function execute_spread {
    basename=$1
    rootdir=$2
    subdir=$3
    numproc=$4
    NEM_SPREAD=$5

    nemesis=$basename.$suffix_mesh.nem
    genesis=$basename.$suffix_mesh

    pexsh=$basename.$suffix_mesh.pex

    # Check for valid executable...
    if [ ! -x $NEM_SPREAD ]
    then
        echo
        echo ${txtred}
        echo "ERROR: Could not find or execute $NEM_SPREAD"
        echo "ERROR: Configuration error"
        echo ${txtrst}
        echo
        return -1
    fi

    if [ "$subdir" != "." ]
    then
        dir=$rootdir/$subdir
        if [ ! -e $dir ]
        then
            echo "making directory $dir"
            mkdir -p $dir
        fi
    fi

    if [ -e $pexsh ]
    then
        /bin/rm -rf $pexsh
    fi

    echo "Input FEM file                   = $genesis"                  >  $pexsh
    echo "LB file                          = $nemesis"                  >> $pexsh
    echo "Parallel Results File Base Name  = $basename"                 >> $pexsh
    echo "File Extension for Spread Files  = .$suffix_spread"           >> $pexsh
    echo "Number of Processors             = $numproc "                 >> $pexsh
    if [ "$debug_level" != "0" ]
    then
    echo "Debug                            = $debug_level"              >> $pexsh
    fi
    echo "------------------------------------------------------------" >> $pexsh
    echo "                Parallel I/O section"                         >> $pexsh
    echo "------------------------------------------------------------" >> $pexsh
    echo "Parallel Disk Info= number=1, offset=1, zeros, nosubdirectory">> $pexsh
    echo "Parallel file location = root=$rootdir, subdir=$subdir"       >> $pexsh

    echo ${txtblu}
    echo "Executing:"
    echo "   $NEM_SPREAD $use64 $pexsh"
    echo ${txtrst}
    ${NOOP:+echo }$NEM_SPREAD $use64 $pexsh
}

########################################################################
# initialize variables
# Text color variables
if [[ $TERM == *"xterm"* ]] || [[ $TERM == "screen" ]]; then
    txtund=$(tput sgr 0 1)    # Underline
    txtbld=$(tput bold)       # Bold
    txtred=$(tput setaf 1)    # Red
    txtgrn=$(tput setaf 2)    # Green
    txtylw=$(tput setaf 3)    # Yellow
    txtblu=$(tput setaf 4)    # Blue
    txtpur=$(tput setaf 5)    # Purple
    txtcyn=$(tput setaf 6)    # Cyan
    txtwht=$(tput setaf 7)    # White
    txtrst=$(tput sgr0)       # Text reset
else
    export TERM=dumb
    txtund=""
    txtbld=""
    txtred=""
    txtgrn=""
    txtylw=""
    txtblu=""
    txtpur=""
    txtcyn=""
    txtwht=""
    txtrst=""
fi

basename=''
processors=''
debug_level='0'
decomp_type="-e"
decomp_method="-l rib"
nem_slice_flag="-c"
weighting=''
rootdir="$(pwd)/"
subdir="."
suffix_spread=''
spheres="-S"
do_viz=""
curdir=$(pwd)
use64=""

ACCESSBIN=$(dirname "$0")
ACCESSBIN=$(cd "${ACCESSBIN}" && pwd)
NEM_SLICE=${ACCESSBIN}/nem_slice
NEM_SPREAD=${ACCESSBIN}/nem_spread
NOOP=

########################################################################
# Test that the getopt being used will handle the long options...
TEST=$(getopt -T)
if [ $? != 4 ] ; then
  old_getopt
  OLDGETOPT=1
else
  OLDGETOPT=0
fi

if [ $# -eq 0 ] ; then
    usage
fi

########################################################################
# decomp options:
if [ $OLDGETOPT -eq 1 ] ; then
   TEMP=$(getopt Hh6ENn:i:o:p:j:r:R:S:X:sl:vVd: "$@")
else
   TEMP=$(getopt -o Hh6ENn:i:o:p:j:r:R:S:X:sl:vVd:w: -a \
       --long input: \
       --long launch: \
       --long nolaunch \
       --long noop \
       --long help \
       --long nem_slice_flag \
       --long 64 \
       --long processors: \
       --long rootdir: \
       --long subdir: \
       --long spread_suffix: \
       --long spheres_linear \
       --long viz_mesh \
       --long verbose \
       --long debug: \
       --long nodal \
       --long elemental \
       --long multikl \
       --long rcb \
       --long rcb_ignore_z \
       --long rib \
       --long hsfc \
       --long spectral \
       --long inertial \
       --long linear \
       --long random \
       --long scattered \
       --long brick \
       --long zpinch \
       --long weight: \
       -n 'decomp' -- "$@")
fi

if [ $? != 0 ] ; then usage ; exit 1 ; fi

# Note the quotes around `$TEMP': they are essential!
eval set -- "$TEMP"
while true ; do
    case "$1" in
    (-h|--help)
        usage ; shift ;;
    (-H|--help)
        usage_new ; shift ;;
    (-v|--verbose)
        set -x ; shift ;;
    (-V|--viz_mesh)
        do_viz="-y" ; shift ;;
    (-d|--debug)
        debug_level="$2" ; shift 2 ;;
    (-n|--nem_slice_flag)
        nem_slice_flag="$2" ; shift 2 ;;
    (-N|--nodal)
        decomp_type="-n" ; shift ;;
    (-E|--elemental)
        decomp_type="-e" ; shift ;;
    (-i|--input)
        input=$2 ; shift 2 ;;
    (-w|--weight)
        weighting="-w $2" ; shift 2 ;;
    (-p|-j|--processors)
        processors=$2 ; shift 2 ;;
    (-R|--rootdir)
        rootdir=$2 ; shift 2 ;;
    (-S|--subdir)
        subdir=$2 ; shift 2 ;;
    (-X|--spread_suffix)
        suffix_spread=$2 ; shift 2 ;;
    (-s|--spheres_linear)
        spheres="" ; shift ;;
    (-6|--64)
        use64="-64" ; shift ;;
    (--noop)
        NOOP="--noop" ; shift ;;
    (--launch)
        LAUNCH="$2" ; shift 2 ;;
    (--nolaunch)
        LAUNCH="nolaunch" ; shift ;;
    (-l)
        decomp_method="-l $2" ; shift 2 ;;

# Decomposition Method options...
    (--multikl)   decomp_method="-l multikl" ; shift ;;
    (--rcb)       decomp_method="-l rcb" ; shift ;;
    (--rcb_ignore_z)       decomp_method="-l rcb,ignore_z" ; shift ;;
    (--rib)       decomp_method="-l rib" ; shift ;;
    (--hsfc)      decomp_method="-l hsfc" ; shift ;;
    (--spectral)  decomp_method="-l spectral" ; shift ;;
    (--inertial)  decomp_method="-l inertial" ; shift ;;
    (--linear)    decomp_method="-l linear" ; shift ;;
    (--random)    decomp_method="-l random" ; shift ;;
    (--scattered) decomp_method="-l scattered" ; shift ;;
    (--brick)     decomp_method="-l brick" ; shift ;;
    (--zpinch)    decomp_method="-l zpinch" ; shift ;;
    (--) shift ; break ;;
    (*) echo "${txtred}ERROR: unrecognized option $1${txtrst}" ; shift ;;
    esac
done

# Figure out how to run nem_slice where we are.
if [ "${LAUNCH:-}" == "nolaunch" ]
then
    # Ensure launch is NOT used --nolaunch was passed.
    LAUNCH=
else
    # If a specific launch script was passed and runs on the node we are on, use it.
    # If not, if Sierra launch command is in path and runs on the node we are on
    # we will try to use that.
    # If not we will run nem_slice directly (ie, w/o launch).
    if [ -z "${LAUNCH:-}" ]
    then
        LAUNCH=$(command -v launch) 
    fi
    if [ -n "${LAUNCH}" ]
    then
        if [ $(basename $LAUNCH) == launch ]
        then
            # If 'launch' was found or passed,
            # Always run nem_slice with -n 1
            LAUNCH="$LAUNCH -n 1"
        fi
        # If the launch script doesn't work to run something simple, don't use it.
        $LAUNCH >/dev/null 2>&1 true || LAUNCH=

        if [ -z "$PBS_JOBID" -a -z "$LSB_JOBID" -a -z "$SLURM_JOB_ID" ]
        then
            # Unset these if we we are NOT in an allocation.
            unset OMPI_MCA_btl_openib_ib_retry_count
            unset OMPI_MCA_btl_openib_ib_timeout
            export OMPI_MCA_mtl=^psm
            export OMPI_MCA_pml=^cm
            export OMPI_MCA_btl=tcp,vader,self
            if [ -n "${SIERRA_LIBRARY_PATH}" ]
            then
                export DYLD_LIBRARY_PATH=$SIERRA_LIBRARY_PATH
            fi
        fi
    fi
fi

########################################################################
if [ $# -eq 0 ] ; then
    echo ${txtred}
    echo "ERROR:******************************************************************"
    echo "ERROR:"
    echo "ERROR:  No 'meshfile' specified."
    echo "ERROR:"
    echo "ERROR:******************************************************************"
    echo ${txtrst}
    usage
else
    if [ -e "$1" ]
    then
        file=$1
        suffix_mesh=${file##*.}
        if test -z "$suffix_spread"
        then
            suffix_spread=$suffix_mesh
        fi
        basename=${file%.*}
    else
    echo ${txtred}
        echo "ERROR:******************************************************************"
        echo "ERROR:"
        echo "ERROR: $1 does not exist."
        echo "ERROR:"
        echo "ERROR:******************************************************************"
        echo ${txtrst}
        usage
    fi
fi

if [ -n "$input" ]
then
   if [ ! -e "$input" ]
   then
    echo ${txtred}
        echo "ERROR:******************************************************************"
        echo "ERROR:"
        echo "ERROR: Input file $input does not exist."
        echo "ERROR:"
        echo "ERROR:******************************************************************"
        echo ${txtrst}
        usage
   fi
fi

if test -z "$processors"
then
    echo
    echo ${txtred}
    echo "ERROR:***************************************************"
    echo "ERROR:"
    echo "ERROR: Must specify number of processors --processors (-p -j) option"
    echo "ERROR:"
    echo "ERROR:***************************************************"
    echo ${txtrst}
    echo
    usage
fi

if [ $processors -eq 1 ]
then
    echo ${txtgrn}
    echo "INFO: No decomposition needed for --processors = 1"
    echo ${txtrst}
    echo
    exit 0
fi

########################################################################
execute_loadbalance "$basename" "$processors" "$decomp_method" "$nem_slice_flag" "$NEM_SLICE" "$input" "$LAUNCH" "$decomp_type" "$weighting"
if [ $? -ne 0 ]
then
    echo ${txtred}
    echo "ERROR:******************************************************************"
    echo "ERROR:"
    echo "ERROR     During nem_slice execution. Check error output above and rerun"
    echo "ERROR:"
    echo "ERROR:******************************************************************"
    echo ${txtrst}
    exit 1
else
    echo "${txtgrn}...nem_slice successful execution${txtrst}"
fi

########################################################################
execute_spread "$basename" "$rootdir" "$subdir" "$processors" "$NEM_SPREAD"
if [ $? -ne 0 ]
then
    echo ${txtred}
    echo "ERROR:******************************************************************"
    echo "ERROR:"
    echo "ERROR    During nem_spread execution. Check error output above and rerun"
    echo "ERROR:"
    echo "ERROR:******************************************************************"
    echo ${txtrst}
    exit 1
else
    echo "${txtgrn}...nem_spread successful execution${txtrst}"
fi

echo "${txtgrn}DECOMP Successful Execution${txtrst}"