#!/bin/bash -e

#
# Global config
#

set -o pipefail

GCC_BUILD_LIST="OpenMP,Pthread,Serial,OpenMP_Serial,Pthread_Serial"
INTEL_BUILD_LIST="OpenMP,Pthread,Serial,OpenMP_Serial,Pthread_Serial"
CLANG_BUILD_LIST="Pthread,Serial,Pthread_Serial"
CUDA_BUILD_LIST="Cuda_OpenMP,Cuda_Pthread,Cuda_Serial"

GCC_WARNING_FLAGS="-Wall,-Wshadow,-pedantic,-Werror,-Wsign-compare,-Wtype-limits,-Wignored-qualifiers,-Wempty-body,-Wclobbered,-Wuninitialized"
CLANG_WARNING_FLAGS="-Wall,-Wshadow,-pedantic,-Werror,-Wsign-compare,-Wtype-limits,-Wuninitialized"
INTEL_WARNING_FLAGS="-Wall,-Wshadow,-pedantic,-Werror,-Wsign-compare,-Wtype-limits,-Wuninitialized"
CUDA_WARNING_FLAGS=""

BASE_MODULE_LIST="<COMPILER_NAME>/<COMPILER_VERSION>/base,hwloc/1.10.1/<COMPILER_NAME>/<COMPILER_VERSION>/base"
CUDA_MODULE_LIST="<COMPILER_NAME>/<COMPILER_VERSION>,gcc/4.7.2/base"

export OMP_NUM_THREADS=4

declare -i NUM_RESULTS_TO_KEEP=7

RESULT_ROOT_PREFIX=TestAll

source /projects/modulefiles/utils/sems-modules-init.sh
source /projects/modulefiles/utils/kokkos-modules-init.sh

SCRIPT_KOKKOS_ROOT=$( cd "$( dirname "$0" )" && cd .. && pwd )

#
# Handle arguments
#

DEBUG=False
ARGS=""
CUSTOM_BUILD_LIST=""
DRYRUN=False
BUILD_ONLY=False
declare -i NUM_JOBS_TO_RUN_IN_PARALLEL=3
TEST_SCRIPT=False

while [[ $# > 0 ]]
do
key="$1"
case $key in
--kokkos-path*)
KOKKOS_PATH="${key#*=}"
;;
--build-list*)
CUSTOM_BUILD_LIST="${key#*=}"
;;
--debug*)
DEBUG=True
;;
--build-only*)
BUILD_ONLY=True
;;
--test-script*)
TEST_SCRIPT=True
;;
--num*)
NUM_JOBS_TO_RUN_IN_PARALLEL="${key#*=}"
;;
--dry-run*)
DRYRUN=True
;;
--help)
echo "test_all_sandia <ARGS> <OPTIONS>:"
echo "--kokkos-path=/Path/To/Kokkos: Path to the Kokkos root directory"
echo "    Defaults to root repo containing this script"
echo "--debug: Run tests in debug. Defaults to False"
echo "--test-script: Test this script, not Kokkos"
echo "--num=N: Number of jobs to run in parallel "
echo "--dry-run: Just print what would be executed"
echo "--build-only: Just do builds, don't run anything"
echo "--build-list=BUILD,BUILD,BUILD..."
echo "    Provide a comma-separated list of builds instead of running all builds"
echo "    Valid items:"
echo "      OpenMP, Pthread, Serial, OpenMP_Serial, Pthread_Serial"
echo "      Cuda_OpenMP, Cuda_Pthread, Cuda_Serial"
echo ""
echo "ARGS: list of expressions matching compilers to test"
echo "  supported compilers"
echo "    gcc/4.7.2"
echo "    gcc/4.8.4"
echo "    gcc/4.9.2"
echo "    gcc/5.1.0"
echo "    intel/14.0.4"
echo "    intel/15.0.2"
echo "    clang/3.5.2"
echo "    clang/3.6.1"
echo "    cuda/6.5.14"
echo "    cuda/7.0.28"
echo "    cuda/7.5.18"
echo ""
echo "Examples:"
echo "  Run all tests"
echo "  % test_all_sandia"
echo ""
echo "  Run all gcc tests"
echo "  % test_all_sandia gcc"
echo ""
echo "  Run all gcc/4.7.2 and all intel tests"
echo "  % test_all_sandia gcc/4.7.2 intel"
echo ""
echo "  Run all tests in debug"
echo "  % test_all_sandia --debug"
echo ""
echo "  Run gcc/4.7.2 and only do OpenMP and OpenMP_Serial builds"
echo "  % test_all_sandia gcc/4.7.2 --build-list=OpenMP,OpenMP_Serial"
echo ""
echo "If you want to kill the tests, do:"
echo "  hit ctrl-z"
echo "  % kill -9 %1"
echo
exit 0
;;
*)
# args, just append
ARGS="$ARGS $1"
;;
esac
shift
done

# set kokkos path
if [ -z "$KOKKOS_PATH" ]; then
    KOKKOS_PATH=$SCRIPT_KOKKOS_ROOT
else
    # Ensure KOKKOS_PATH is abs path
    KOKKOS_PATH=$( cd $KOKKOS_PATH && pwd )
fi

# set build type
if [ "$DEBUG" = "True" ]; then
    BUILD_TYPE=debug
else
    BUILD_TYPE=release
fi

# If no args provided, do all compilers
if [ -z "$ARGS" ]; then
    ARGS='?'
fi

# Format: (compiler module-list build-list exe-name warning-flag)
COMPILERS=("gcc/4.7.2 $BASE_MODULE_LIST $GCC_BUILD_LIST g++ $GCC_WARNING_FLAGS"
           "gcc/4.8.4 $BASE_MODULE_LIST $GCC_BUILD_LIST g++ $GCC_WARNING_FLAGS"
           "gcc/4.9.2 $BASE_MODULE_LIST $GCC_BUILD_LIST g++ $GCC_WARNING_FLAGS"
           "gcc/5.1.0 $BASE_MODULE_LIST $GCC_BUILD_LIST g++ $GCC_WARNING_FLAGS"
           "intel/14.0.4 $BASE_MODULE_LIST $INTEL_BUILD_LIST icpc $INTEL_WARNING_FLAGS"
           "intel/15.0.2 $BASE_MODULE_LIST $INTEL_BUILD_LIST icpc $INTEL_WARNING_FLAGS"
           "clang/3.5.2 $BASE_MODULE_LIST $CLANG_BUILD_LIST clang++ $CLANG_WARNING_FLAGS"
           "clang/3.6.1 $BASE_MODULE_LIST $CLANG_BUILD_LIST clang++ $CLANG_WARNING_FLAGS"
           "cuda/6.5.14 $CUDA_MODULE_LIST $CUDA_BUILD_LIST $KOKKOS_PATH/config/nvcc_wrapper $CUDA_WARNING_FLAGS"
           "cuda/7.0.28 $CUDA_MODULE_LIST $CUDA_BUILD_LIST $KOKKOS_PATH/config/nvcc_wrapper $CUDA_WARNING_FLAGS"
           "cuda/7.5.18 $CUDA_MODULE_LIST $CUDA_BUILD_LIST $KOKKOS_PATH/config/nvcc_wrapper $CUDA_WARNING_FLAGS"
           )

# Process args to figure out which compilers to test
COMPILERS_TO_TEST=""
for ARG in $ARGS; do
    for COMPILER_DATA in "${COMPILERS[@]}"; do
        ARR=($COMPILER_DATA)
        COMPILER=${ARR[0]}
        if [[ "$COMPILER" = $ARG* ]]; then
            if [[ "$COMPILERS_TO_TEST" != *${COMPILER}* ]]; then
                COMPILERS_TO_TEST="$COMPILERS_TO_TEST $COMPILER"
            else
                echo "Tried to add $COMPILER twice"
            fi
        fi
    done
done

#
# Functions
#

# get_compiler_name <COMPILER>
get_compiler_name() {
    echo $1 | cut -d/ -f1
}

# get_compiler_version <COMPILER>
get_compiler_version() {
    echo $1 | cut -d/ -f2
}

# Do not call directly
get_compiler_data() {
    local compiler=$1
    local item=$2
    local compiler_name=$(get_compiler_name $compiler)
    local compiler_vers=$(get_compiler_version $compiler)

    local compiler_data
    for compiler_data in "${COMPILERS[@]}" ; do
        local arr=($compiler_data)
        if [ "$compiler" = "${arr[0]}" ]; then
            echo "${arr[$item]}" | tr , ' ' | sed -e "s/<COMPILER_NAME>/$compiler_name/g" -e "s/<COMPILER_VERSION>/$compiler_vers/g"
            return 0
        fi
    done

    # Not found
    echo "Unreconized compiler $compiler" >&2
    exit 1
}

#
# For all getters, usage: <GETTER> <COMPILER>
#

get_compiler_modules() {
    get_compiler_data $1 1
}

get_compiler_build_list() {
    get_compiler_data $1 2
}

get_compiler_exe_name() {
    get_compiler_data $1 3
}

get_compiler_warning_flags() {
    get_compiler_data $1 4
}

run_cmd() {
    echo "RUNNING: $*"
    if [ "$DRYRUN" != "True" ]; then
	eval "$* 2>&1"
    fi
}

# report_and_log_test_results <SUCCESS> <DESC> <PHASE>
report_and_log_test_result() {
    # Use sane var names
    local success=$1; local desc=$2; local phase=$3;

    if [ "$success" = "0" ]; then
	echo "  PASSED $desc"
        touch $PASSED_DIR/$desc
    else
	echo "  FAILED $desc" >&2
        echo $phase > $FAILED_DIR/$desc
        cat ${desc}.${phase}.log
    fi
}

setup_env() {
    local compiler=$1
    local compiler_modules=$(get_compiler_modules $compiler)

    module purge

    local mod
    for mod in $compiler_modules; do
	module load $mod 2>&1
        # It is ridiculously hard to check for the success of a loaded
        # module. Module does not return error codes and piping to grep
        # causes module to run in a subshell.
        module list 2>&1 | grep "$mod" >& /dev/null || return 1
    done

    return 0
}

# single_build_and_test <COMPILER> <BUILD> <BUILD_TYPE>
single_build_and_test() {
    # Use sane var names
    local compiler=$1; local build=$2; local build_type=$3;

    # set up env
    mkdir -p $ROOT_DIR/$compiler/"${build}-$build_type"
    cd $ROOT_DIR/$compiler/"${build}-$build_type"
    local desc=$(echo "${compiler}-${build}-${build_type}" | sed 's:/:-:g')
    setup_env $compiler >& ${desc}.configure.log || { report_and_log_test_result 1 ${desc} configure && return 0; }

    # Set up flags
    local compiler_warning_flags=$(get_compiler_warning_flags $compiler)
    local compiler_exe=$(get_compiler_exe_name $compiler)

    if [[ "$build_type" = hwloc* ]]; then
        local extra_args=--with-hwloc=$(dirname $(dirname $(which hwloc-info)))
    fi

    if [[ "$build_type" = *debug* ]]; then
        local extra_args="$extra_args --debug"
        local cxxflags="-g $compiler_warning_flags"
    else
        local cxxflags="-O3 $compiler_warning_flags"
    fi

    if [[ "$compiler" == cuda* ]]; then
        cxxflags="--keep --keep-dir=$(pwd) $cxxflags"
        export TMPDIR=$(pwd)
    fi

    # cxxflags="-DKOKKOS_USING_EXPERIMENTAL_VIEW $cxxflags"

    echo "  Starting job $desc"

    if [ "$TEST_SCRIPT" = "True" ]; then
        local rand=$[ 1 + $[ RANDOM % 10 ]]
        sleep $rand
        if [ $rand -gt 5 ]; then
            run_cmd ls fake_problem >& ${desc}.configure.log || { report_and_log_test_result 1 $desc configure && return 0; }
        fi
    else
        run_cmd ${KOKKOS_PATH}/generate_makefile.bash --with-devices=$build --compiler=$(which $compiler_exe) --cxxflags=\"$cxxflags\" $extra_args >& ${desc}.configure.log || { report_and_log_test_result 1 ${desc} configure && return 0; }
        run_cmd make build-test >& ${desc}.build.log || { report_and_log_test_result 1 ${desc} build && return 0; }
        if [[ "$BUILD_ONLY" == False ]]; then
            run_cmd make test >& ${desc}.test.log || { report_and_log_test_result 1 ${desc} test && return 0; }
        fi
    fi

    report_and_log_test_result 0 $desc

    return 0
}

# wait_for_jobs <NUM-JOBS>
wait_for_jobs() {
    local -i max_jobs=$1
    local -i num_active_jobs=$(jobs | wc -l)
    while [ $num_active_jobs -ge $max_jobs ]
    do
        sleep 1
        num_active_jobs=$(jobs | wc -l)
        jobs >& /dev/null
    done
}

# run_in_background <COMPILER> <BUILD> <BUILD_TYPE>
run_in_background() {
    local compiler=$1

    local -i num_jobs=$NUM_JOBS_TO_RUN_IN_PARALLEL
    if [[ "$BUILD_ONLY" == True ]]; then
        num_jobs=8
    else
        if [[ "$compiler" == cuda* ]]; then
            num_jobs=1
        fi
    fi
    wait_for_jobs $num_jobs

    single_build_and_test $* &
}

# build_and_test_all <COMPILER>
build_and_test_all() {
    # Get compiler data
    local compiler=$1
    if [ -z "$CUSTOM_BUILD_LIST" ]; then
	local compiler_build_list=$(get_compiler_build_list $compiler)
    else
	local compiler_build_list=$(echo "$CUSTOM_BUILD_LIST" | tr , ' ')
    fi

    # do builds
    local build
    for build in $compiler_build_list
    do
	run_in_background $compiler $build $BUILD_TYPE

        # If not cuda, do a hwloc test too
        if [[ "$compiler" != cuda* ]]; then
            run_in_background $compiler $build "hwloc-$BUILD_TYPE"
        fi
    done

    return 0
}

get_test_root_dir() {
    local existing_results=$(find . -maxdepth 1 -name "$RESULT_ROOT_PREFIX*" | sort)
    local -i num_existing_results=$(echo $existing_results | tr ' ' '\n' | wc -l)
    local -i num_to_delete=${num_existing_results}-${NUM_RESULTS_TO_KEEP}

    if [ $num_to_delete -gt 0 ]; then
        /bin/rm -rf $(echo $existing_results | tr ' ' '\n' | head -n $num_to_delete)
    fi

    echo $(pwd)/${RESULT_ROOT_PREFIX}_$(date +"%Y-%m-%d_%H.%M.%S")
}

wait_summarize_and_exit() {
    wait_for_jobs 1

    echo "#######################################################"
    echo "PASSED TESTS"
    echo "#######################################################"

    \ls -1 $PASSED_DIR | sort

    echo "#######################################################"
    echo "FAILED TESTS"
    echo "#######################################################"

    local failed_test
    local -i rv=0
    for failed_test in $(\ls -1 $FAILED_DIR)
    do
        echo $failed_test "("$(cat $FAILED_DIR/$failed_test)" failed)"
        rv=$rv+1
    done

    exit $rv
}

#
# Main
#

ROOT_DIR=$(get_test_root_dir)
mkdir -p $ROOT_DIR
cd $ROOT_DIR

PASSED_DIR=$ROOT_DIR/results/passed
FAILED_DIR=$ROOT_DIR/results/failed
mkdir -p $PASSED_DIR
mkdir -p $FAILED_DIR

echo "Going to test compilers: " $COMPILERS_TO_TEST
for COMPILER in $COMPILERS_TO_TEST; do
    echo "Testing compiler $COMPILER"
    build_and_test_all $COMPILER
done

wait_summarize_and_exit
