#!/bin/ksh -x

if [ $# -ne 1 ]; then

    echo "Must specify an input argument to set runtime environment variables!"
    echo "argument can be any one of the following:"
    echo "anal fcst post vrfy metp"
    echo "eobs eupd ecen esfc efcs epos"
    echo "postsnd awips gempak"
    exit 1

fi

step=$1

# WCOSS2 information
export launcher="mpiexec -l"
export mpmd="--cpu-bind verbose,core cfp"

export npe_node_max=128

export job=${PBS_JOBNAME:-$step}
export jobid=${job}.${PBS_JOBID:-$$}

if [ $step = "prep" -o $step = "prepbufr" ]; then

    export POE=${POE:-"YES"}
    export BACK=${BACK:-"off"}
    export sys_tp="wcoss2"
    export launcher_PREP="mpiexec"

elif [ $step = "waveinit" -o $step = "waveprep" -o $step = "wavepostsbs" -o $step = "wavepostbndpnt" -o $step = "wavepostbndpntbll"  -o $step = "wavepostpnt" ]; then

    if [ $step = "waveprep" -a $CDUMP = "gfs" ]; then NTASKS=$NTASKS_gfs ; fi
    export wavempexec="$launcher -np"
    export wave_mpmd=${mpmd}

elif [ $step = "anal" ]; then

    export OMP_PLACES=cores
    export OMP_STACKSIZE=1G
    export FI_OFI_RXM_SAR_LIMIT=3145728

    export NTHREADS_GSI=$nth_anal
    export APRUN_GSI="$launcher -n ${npe_gsi:-${npe_anal:-$PBS_NP}} -ppn $npe_node_anal --cpu-bind depth --depth $NTHREADS_GSI"

    export NTHREADS_CALCINC=${nth_calcinc:-1}
    export APRUN_CALCINC="$launcher \$ncmd"

    export NTHREADS_CYCLE=${nth_cycle:-14}
    [[ $NTHREADS_CYCLE -gt $npe_node_max ]] && export NTHREADS_CYCLE=$npe_node_max
    npe_cycle=${ntiles:-6}
    export APRUN_CYCLE="$launcher -n $npe_cycle -ppn $npe_node_cycle --cpu-bind depth --depth $NTHREADS_CYCLE"

    export NTHREADS_GAUSFCANL=1
    npe_gausfcanl=${npe_gausfcanl:-1}
    export APRUN_GAUSFCANL="$launcher -n $npe_gausfcanl"

    export NTHREADS_CHGRES=${nth_echgres:-14}
    [[ $NTHREADS_CHGRES -gt $npe_node_max ]] && export NTHREADS_CHGRES=$npe_node_max
    export APRUN_CHGRES=""

    export CFP_MP=${CFP_MP:-"NO"}
    export USE_CFP=${USE_CFP:-"YES"}
    export APRUNCFP="$launcher -np \$ncmd $mpmd"

elif [ $step = "gldas" ]; then

    export NTHREADS_GLDAS=$nth_gldas
    export APRUN_GLDAS="$launcher -n $npe_gldas -ppn $npe_node_gldas --cpu-bind depth --depth $NTHREADS_GLDAS"

    export NTHREADS_GAUSSIAN=${nth_gaussian:-1}
    export APRUN_GAUSSIAN="$launcher -n $npe_gaussian -ppn $npe_node_gaussian --cpu-bind depth --depth $NTHREADS_GAUSSIAN"

    export USE_CFP=${USE_CFP:-"YES"}
    export APRUN_GLDAS_DATA_PROC="$launcher -np $npe_gldas $mpmd"

elif [ $step = "eobs" ]; then

    export OMP_PLACES=cores
    export OMP_STACKSIZE=1G
    export FI_OFI_RXM_SAR_LIMIT=3145728

    export NTHREADS_GSI=$nth_eobs
    export APRUN_GSI="$launcher -n ${npe_gsi:-${npe_eobs:-$PBS_NP}} -ppn $npe_node_eobs --cpu-bind depth --depth $NTHREADS_GSI"

    export CFP_MP=${CFP_MP:-"NO"}
    export USE_CFP=${USE_CFP:-"YES"}
    export APRUNCFP="$launcher -np \$ncmd $mpmd"

elif [ $step = "eupd" ]; then

    export OMP_PLACES=cores
    export OMP_PROC_BIND=true
    export OMP_STACKSIZE=2G
    export MPICH_COLL_OPT_OFF=1
    export FI_OFI_RXM_SAR_LIMIT=3145728

    export NTHREADS_ENKF=$nth_eupd
    export APRUN_ENKF="$launcher -n ${npe_enkf:-${npe_eupd:-$PBS_NP}} -ppn $npe_node_eupd --cpu-bind core --depth $NTHREADS_ENKF"

    export CFP_MP=${CFP_MP:-"NO"}
    export USE_CFP=${USE_CFP:-"YES"}
    export APRUNCFP="$launcher -np \$ncmd $mpmd"

elif [ $step = "fcst" ]; then

    export OMP_PLACES=cores
    export OMP_STACKSIZE=2048M
    export FI_OFI_RXM_RX_SIZE=40000
    export FI_OFI_RXM_TX_SIZE=40000

    export NTHREADS_FV3=$nth_fv3
    export cores_per_node=$npe_node_max
    if [ $CDUMP = "gdas" ]; then
      export APRUN_FV3="$launcher -n ${npe_fcst:-$PBS_NP} -ppn $npe_node_fcst --cpu-bind depth --depth $NTHREADS_FV3"
    else
      export APRUN_FV3="$launcher -n ${npe_fcst_gfs:-$PBS_NP} -ppn $npe_node_fcst_gfs --cpu-bind depth --depth $NTHREADS_FV3"
    fi
    export NTHREADS_REGRID_NEMSIO=${nth_regrid_nemsio:-1}
    export APRUN_REGRID_NEMSIO="$launcher -n $LEVS"

    export NTHREADS_REMAP=${nth_remap:-2}
    export APRUN_REMAP="$launcher -n ${npe_remap:-${npe_fcst:-$PBS_NP}}"

elif [ $step = "efcs" ]; then

    export MPICH_MPIIO_HINTS="*:romio_cb_write=disable"
    export FI_OFI_RXM_SAR_LIMIT=3145728
    export FI_OFI_RXM_RX_SIZE=40000
    export FI_OFI_RXM_TX_SIZE=40000

    export NTHREADS_FV3=$nth_efcs
    export cores_per_node=$npe_node_max
    export APRUN_FV3="$launcher -n ${npe_fv3:-${npe_efcs:-$PBS_NP}} -ppn $npe_node_efcs --cpu-bind depth --depth $NTHREADS_FV3"

    export NTHREADS_REGRID_NEMSIO=${nth_regrid_nemsio:-1}
    export APRUN_REGRID_NEMSIO="$launcher -n $LEVS"

elif [ $step = "post" ]; then

    export NTHREADS_NP=${nth_np:-1}
    export APRUN_NP="$launcher -n ${npe_np:-${npe_post:-$PBS_NP}} -ppn $npe_node_post --cpu-bind depth --depth $NTHREADS_NP"

    export NTHREADS_DWN=${nth_dwn:-1}
    export APRUN_DWN="$launcher -np ${npe_dwn:-$PBS_NP} $mpmd"

elif [ $step = "ecen" ]; then

    export NTHREADS_ECEN=$nth_ecen
    export APRUN_ECEN="$launcher -n ${npe_ecen:-$PBS_NP} -ppn $npe_node_ecen --cpu-bind depth --depth $NTHREADS_ECEN"

    export NTHREADS_CHGRES=${nth_chgres:-14}
    [[ $NTHREADS_CHGRES -gt $npe_node_max ]] && export NTHREADS_CHGRES=$npe_node_max
    export APRUN_CHGRES="time"

    export NTHREADS_CALCINC=${nth_calcinc:-1}
    export APRUN_CALCINC="$launcher -n ${npe_ecen:-$PBS_NP}"

    export NTHREADS_CYCLE=${nth_cycle:-14}
    [[ $NTHREADS_CYCLE -gt $npe_node_max ]] && export NTHREADS_CYCLE=$npe_node_max
    export APRUN_CYCLE="$launcher -n $npe_ecen -ppn $npe_node_cycle --cpu-bind depth --depth $NTHREADS_CYCLE"

elif [ $step = "esfc" ]; then

    export NTHREADS_ESFC=$nth_esfc
    export APRUN_ESFC="$launcher -n ${npe_esfc:-$PBS_NP} -ppn $npe_node_esfc --cpu-bind depth --depth $NTHREADS_ESFC"

    export NTHREADS_CYCLE=${nth_cycle:-14}
    [[ $NTHREADS_CYCLE -gt $npe_node_max ]] && export NTHREADS_CYCLE=$npe_node_max
    export APRUN_CYCLE="$launcher -n $npe_esfc -ppn $npe_node_cycle --cpu-bind depth --depth $NTHREADS_CYCLE"

elif [ $step = "epos" ]; then

    export NTHREADS_EPOS=$nth_epos
    export APRUN_EPOS="$launcher -n ${npe_epos:-$PBS_NP} -ppn $npe_node_epos --cpu-bind depth --depth $NTHREADS_EPOS"

elif [ $step = "fv3ic" ]; then

    export NTHREADS_CHGRES=${nth_chgres:-$npe_node_max}
    [[ $NTHREADS_CHGRES -gt $npe_node_max ]] && export NTHREADS_CHGRES=$npe_node_max
    export APRUN_CHGRES="time"

elif [ $step = "postsnd" ]; then

    export MPICH_MPIIO_HINTS_DISPLAY=1
    export OMP_NUM_THREADS=1
    export NTHREADS_POSTSND=${nth_postsnd:-1}
    export APRUN_POSTSND="$launcher -n $npe_postsnd --depth=$NTHREADS_POSTSND --cpu-bind depth"
    export NTHREADS_POSTSNDCFP=${nth_postsndcfp:-1}
    export APRUN_POSTSNDCFP="$launcher -np $npe_postsndcfp $mpmd"

elif [ $step = "awips" ]; then

    export NTHREADS_AWIPS=${nth_awips:-2}
    export APRUN_AWIPSCFP="$launcher -np ${npe_awips:-$PBS_NP} $mpmd"

elif [ $step = "gempak" ]; then

    export NTHREADS_GEMPAK=${nth_gempak:-1}
    export APRUN_GEMPAKCFP="$launcher -np \$ntasks $mpmd"

elif [ $step = "waveawipsbulls" ]; then

    unset PERL5LIB

fi