#!/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"
    echo "eobs eupd ecen efcs epos"
    echo "postsnd awips gempak"
    exit 1

fi

step=$1

# Cray information
export launcher="aprun"
export npe_node_max=24

# Configure MPI environment
export MPI_BUFS_PER_PROC=2048
export MPI_BUFS_PER_HOST=2048
export MPI_GROUP_MAX=256
export MPI_MEMMAP_OFF=1
export MP_STDOUTMODE="ORDERED"
export NTHSTACK=1024000000
export OMP_STACKSIZE="2048M"
export KMP_AFFINITY="disabled"
export job=${LSB_JOBNAME:-$step}
export jobid=${job}.${LSB_JOBID:-$$}

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

    nth_max=$(($npe_node_max / $npe_node_prep))

    export POE=${POE:-"YES"}
    export sys_tp="Cray-XC40"

elif [ $step = "anal" ]; then

    nth_max=$(($npe_node_max / $npe_node_anal))

    export NTHREADS_GSI=${nth_gsi:-$nth_max}
    [[ $NTHREADS_GSI -gt $nth_max ]] && export NTHREADS_GSI=$nth_max
    export APRUN_GSI="$launcher -j 1 -n ${npe_gsi:-$npe_anal} -N $npe_node_anal -d $NTHREADS_GSI -cc depth"

    export NTHREADS_CALCINC=${nth_calcinc:-1}
    [[ $NTHREADS_CALCINC -gt $nth_max ]] && export NTHREADS_CALCINC=$nth_max
    export APRUN_CALCINC="$launcher -j 1 -n \$ncmd -N 1 -d $NTHREADS_CALCINC -cc depth"

    if [ ${USE_CFP:-"NO"} = "YES" ]; then
        export APRUNCFP="$launcher -q -j 1 -n \$ncmd -N \$ncmd_max -d 1 cfp"
    fi

    export NTHREADS_CYCLE=${nth_cycle:-12}
    [[ $NTHREADS_CYCLE -gt $npe_node_max ]] && export NTHREADS_CYCLE=$npe_node_max
    npe_node_cycle=$(($npe_node_max / $NTHREADS_CYCLE))
    npe_cycle=${ntiles:-6}
    [[ $npe_node_cycle -gt $npe_cycle ]] && npe_node_cycle=$npe_cycle
    export APRUN_CYCLE="$launcher -j 1 -n $npe_cycle -N $npe_node_cycle -d $NTHREADS_CYCLE -cc depth"

    export NTHREADS_GAUSFCANL=1
    npe_node_gausfcanl=$(($npe_node_max / $NTHREADS_GAUSFCANL))
    npe_gausfcanl=${npe_gausfcanl:-1}
    [[ $npe_node_gausfcanl -gt $npe_gausfcanl ]] && npe_node_gausfcanl=$npe_gausfcanl
    export APRUN_GAUSFCANL="$launcher -j 1 -n $npe_gausfcanl -N $npe_node_gausfcanl -d $NTHREADS_GAUSFCANL -cc depth"

elif [ $step = "eobs" ]; then

    nth_max=$(($npe_node_max / $npe_node_eobs))

    export NTHREADS_GSI=${nth_gsi:-$nth_max}
    [[ $NTHREADS_GSI -gt $nth_max ]] && export NTHREADS_GSI=$nth_max
    export APRUN_GSI="$launcher -j 1 -n ${npe_gsi:-$npe_eobs} -N $npe_node_eobs -d $NTHREADS_GSI -cc depth"

    if [ ${USE_CFP:-"NO"} = "YES" ]; then
        export APRUNCFP="$launcher -q -j 1 -n \$ncmd -N \$ncmd_max -d 1 cfp"
    fi

elif [ $step = "eupd" ]; then

    nth_max=$(($npe_node_max / $npe_node_eupd))

    export NTHREADS_ENKF=${nth_enkf:-$nth_max}
    [[ $NTHREADS_ENKF -gt $nth_max ]] && export NTHREADS_ENKF=$nth_max
    export APRUN_ENKF="$launcher -j 1 -n ${npe_enkf:-$npe_eupd} -N $npe_node_eupd -d $NTHREADS_ENKF -cc depth"

    if [ ${USE_CFP:-"NO"} = "YES" ]; then
        export APRUNCFP="$launcher -q -j 1 -n \$ncmd -N \$ncmd_max -d 1 cfp"
    fi

elif [ $step = "fcst" ]; then

    nth_max=$(($npe_node_max / $npe_node_fcst))

    export NTHREADS_FV3=${nth_fv3:-$nth_max}
    [[ $NTHREADS_FV3 -gt $nth_max ]] && export NTHREADS_FV3=$nth_max
    export cores_per_node=$npe_node_max
    export APRUN_FV3="$launcher -j 1 -n ${npe_fv3:-$npe_fcst} -N $npe_node_fcst -d $NTHREADS_FV3 -cc depth"

    export NTHREADS_REGRID_NEMSIO=${nth_regrid_nemsio:-1}
    [[ $NTHREADS_REGRID_NEMSIO -gt $nth_max ]] && export NTHREADS_REGRID_NEMSIO=$nth_max
    export APRUN_REGRID_NEMSIO="$launcher -j 1 -n $LEVS -N $npe_node_fcst -d $NTHREADS_REGRID_NEMSIO -cc depth"

    export NTHREADS_REMAP=${nth_remap:-2}
    [[ $NTHREADS_REMAP -gt $nth_max ]] && export NTHREADS_REMAP=$nth_max
    export APRUN_REMAP="$launcher -j 1 -n ${npe_remap:-$npe_fcst} -N $npe_node_fcst -d $NTHREADS_REMAP -cc depth"

elif [ $step = "efcs" ]; then

    nth_max=$(($npe_node_max / $npe_node_efcs))

    export NTHREADS_FV3=${nth_fv3:-$nth_max}
    [[ $NTHREADS_FV3 -gt $nth_max ]] && export NTHREADS_FV3=$nth_max
    export cores_per_node=$npe_node_max
    export APRUN_FV3="$launcher -j 1 -n ${npe_fv3:-$npe_efcs} -N $npe_node_efcs -d $NTHREADS_FV3 -cc depth"

    export NTHREADS_REGRID_NEMSIO=${nth_regrid_nemsio:-1}
    [[ $NTHREADS_REGRID_NEMSIO -gt $nth_max ]] && export NTHREADS_REGRID_NEMSIO=$nth_max
    export APRUN_REGRID_NEMSIO="$launcher -j 1 -n $LEVS -N $npe_node_efcs -d $NTHREADS_REGRID_NEMSIO -cc depth"

elif [ $step = "post" ]; then

    nth_max=$(($npe_node_max / $npe_node_post))

    export NTHREADS_NP=${nth_np:-1}
    [[ $NTHREADS_NP -gt $nth_max ]] && export NTHREADS_NP=$nth_max
    export APRUN_NP="$launcher -j 1 -n ${npe_np:-$npe_post} -N $npe_node_post -d $NTHREADS_NP -cc depth"

    export NTHREADS_DWN=${nth_dwn:-1}
    [[ $NTHREADS_DWN -gt $nth_max ]] && export NTHREADS_DWN=$nth_max
    export APRUN_DWN="$launcher -j 1 -n $npe_dwn -N $npe_node_dwn -d $NTHREADS_DWN cfp"

elif [ $step = "ecen" ]; then

    nth_max=$(($npe_node_max / $npe_node_ecen))

    export NTHREADS_ECEN=${nth_ecen:-$nth_max}
    [[ $NTHREADS_ECEN -gt $nth_max ]] && export NTHREADS_ECEN=$nth_max
    export APRUN_ECEN="$launcher -j 1 -n $npe_ecen -N $npe_node_ecen -d $NTHREADS_ECEN -cc depth"

    export NTHREADS_CHGRES=${nth_chgres:-12}
    [[ $NTHREADS_CHGRES -gt $npe_node_max ]] && export NTHREADS_CHGRES=$npe_node_max
    export APRUN_CHGRES="$launcher -j 1 -n 1 -N 1 -d $NTHREADS_CHGRES -cc depth"

    export NTHREADS_CALCINC=${nth_calcinc:-1}
    [[ $NTHREADS_CALCINC -gt $nth_max ]] && export NTHREADS_CALCINC=$nth_max
    export APRUN_CALCINC="$launcher -j 1 -n $npe_ecen -N $npe_node_ecen -d $NTHREADS_CALCINC -cc depth"

elif [ $step = "epos" ]; then

    nth_max=$(($npe_node_max / $npe_node_epos))

    export NTHREADS_EPOS=${nth_epos:-$nth_max}
    [[ $NTHREADS_EPOS -gt $nth_max ]] && export NTHREADS_EPOS=$nth_max
    export APRUN_EPOS="$launcher -j 1 -n $npe_epos -N $npe_node_epos -d $NTHREADS_EPOS -cc depth"

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="$launcher -j 1 -n 1 -N 1 -d $NTHREADS_CHGRES -cc depth"

elif [ $step = "vrfy" ]; then

    export IOBUF_PARAMS="*:size=32M:count=4:verbose"
    export APRUNTRACK="$launcher -j1 -n1 -N1 -d1 -cc depth"

elif [ $step = "postsnd" ]; then
    export IOBUF_PARAMS="sigf*:size=128M:count=20:prefetch=0:verbose,gfs_collectiv*:size=128M:count=2:prefetch=0:verbose,*.snd:size=128M:count=3:prefetch=0:verbose,*.sfc:size=32M:count=3:prefetch=0:verbose,bufr.*:size=8M:count=20:prefetch=0:verbose"

    nth_max=$(($npe_node_max / $npe_node_postsnd))

    export NTHREADS_POSTSND=${nth_postsnd:-1}
    [[ $NTHREADS_POSTSND -gt $nth_max ]] && export NTHREADS_POSTSND=$nth_max
    export APRUN_POSTSND="$launcher -j 1 -n $npe_postsnd -N $npe_node_postsnd -d $NTHREADS_POSTSND -cc depth"

    export NTHREADS_POSTSNDCFP=${nth_postsndcfp:-1}
    [[ $NTHREADS_POSTSNDCFP -gt $nth_max ]] && export NTHREADS_POSTSNDCFP=$nth_max
    export APRUN_POSTSNDCFP="$launcher -j 1 -n $npe_postsndcfp -N $npe_node_postsndcfp -d $NTHREADS_POSTSNDCFP cfp"

elif [ $step = "awips" ]; then
    nth_max=$(($npe_node_max / $npe_node_awips))

    export NTHREADS_AWIPS=${nth_awips:-2}
    [[ $NTHREADS_AWIPS -gt $nth_max ]] && export NTHREADS_AWIPS=$nth_max
    export APRUN_AWIPSCFP="$launcher -j 1 -n $npe_awips -N $npe_node_awips -d $NTHREADS_AWIPS -cc depth cfp"
    
elif [ $step = "gempak" ]; then
    nth_max=$(($npe_node_max / $npe_node_gempak))

    export NTHREADS_GEMPAK=${nth_gempak:-3}
    [[ $NTHREADS_GEMPAK -gt $nth_max ]] && export NTHREADS_GEMPAK=$nth_max
    export APRUN_GEMPAKCFP="$launcher -j 1 -n $npe_gempak -N $npe_node_gempak -d $NTHREADS_GEMPAK cfp"
    
fi