#!/bin/bash

if [ $# -lt 1 ]; then
    echo $0
    echo "Not enough args"
    echo "Use as: $0 <framework>, where framework is one of (mxnet, pytorch, tensorflow)"
    exit 1
fi

FRAMEWORK=$1
# Assuming that if a version is released to pypi, it will have a corresponding release tag on GitHub.
# Not using master because there could be breaking changes in tests between versions.
SMDEBUG_RELEASE_TAG=$(python -c "import smdebug; print(smdebug.__version__)")

# We need to filter only the word from the output of the previous line because
# import statements might contain log statements which corrupt the logic below
SMDEBUG_RELEASE_TAG=$(echo $SMDEBUG_RELEASE_TAG |  perl -lane 'print $F[-1]')

# declare -A only works with bash 4 and above, and creates an "associate array" in bash
if [ "$(printf '%s\n' "0.8.0" "$SMDEBUG_RELEASE_TAG" | sort -V | head -n1)" = "0.8.0" ]; then
    declare -A ZCC_TEST_NAME=( ["mxnet"]="test_mxnet_gluon_integration.py" \
                           ["pytorch"]="test_pytorch_integration.py" \
                           ["tensorflow"]="test_tensorflow_integration.py" \
                           ["tensorflow2"]="test_tensorflow2_integration.py")
else
    declare -A ZCC_TEST_NAME=( ["mxnet"]="mxnet_gluon_integration_test.py" \
                           ["pytorch"]="pytorch_integration_tests.py" \
                           ["tensorflow"]="tensorflow_integration_tests.py" \
                           ["tensorflow2"]="tensorflow2_integration_tests.py")
fi

HOME_DIR=/test
BIN_DIR=${HOME_DIR}/bin
LOG_DIR=${HOME_DIR}/logs

echo "Simply verify if smdebug is installed correctly."
echo "Cloning git repo for release tag $SMDEBUG_RELEASE_TAG"
git clone -b ${SMDEBUG_RELEASE_TAG} https://github.com/awslabs/sagemaker-debugger.git sagemaker-debugger || exit 1

echo "Changing to sagemaker-debugger."
cd sagemaker-debugger

git pull

echo "Creating mock test folder."
mkdir -p mock
cp -r tests mock/
cp -r examples mock/
cd mock/

echo "Installing pre-requisites."
pip install pytest

echo "Running framework tests."


if [[ $FRAMEWORK == "tensorflow" ]]; then
    pip install tensorflow-datasets
    pytest tests/${FRAMEWORK} \
                            -k "not test_keras and not test_mnist_s3 and not test_mnist_s3_multi_save_configs" \
                            --ignore=tests/mxnet/test_training_end.py
    exit_code=$?
    if [ "$exit_code" -eq "0" ]
    then
        echo "SMDebug framework tests succeeded"
    else
        echo "SMDebug framework tests failed"
        exit $exit_code
    fi
elif [[ $FRAMEWORK == "mxnet" ]]; then
    pytest tests/${FRAMEWORK}  --ignore=tests/mxnet/test_training_end.py
    exit_code=$?
    if [ "$exit_code" -eq "0" ]
    then
        echo "SMDebug tests succeeded for $FRAMEWORK"
    else
        echo "SMDebug tests failed for $FRAMEWORK"
        exit $exit_code
    fi
elif [[ $FRAMEWORK == "tensorflow2" ]]; then
    tf_version=$(python -c "import tensorflow as tf; print(tf.__version__)")
    for test_file in tests/${FRAMEWORK}/*
    do
        if [[ "$tf_version" == *"2.7."* && "$test_file" == *"test_keras_mirrored.py" ]]; then
            echo "skipping test_keras_mirrored.py for TF2.7 as it is flaky"
        else
            pytest $test_file
            exit_code=$?
            if [ "$exit_code" -eq "0" ]; then
                echo "SMDebug test succeeded: $test_file"
            elif [ "$exit_code" -eq "5" ]; then
                echo "SMDebug no tests collected: $test_file"
            else
                echo "SMDebug test failed: $test_file"
                exit $exit_code
            fi
        fi
    done
else
    for test_file in tests/${FRAMEWORK}/*
    do
        pytest $test_file
        exit_code=$?
        if [ "$exit_code" -eq "0" ]
        then
            echo "SMDebug test succeeded: $test_file"
        elif [ "$exit_code" -eq "5" ]
        then
            # pytest enum ExitCode.NO_TESTS_COLLECTED = 5
            # See: https://github.com/pytest-dev/pytest/blob/050c2df737f7c6c3b1de7a6f375de107edcb6873/src/_pytest/config/__init__.py#L98
            echo "SMDebug no tests collected: $test_file"
        else
            echo "SMDebug test failed: $test_file"
            exit $exit_code
        fi
    done
fi

echo "Setting up SMDebug environment."
echo "{\"LocalPath\": \"newlogsRunTest1/test_hook_from_json_config\"}" | tee /tmp/test_hook_from_json_config.json
export SMDEBUG_CONFIG_FILE_PATH=/tmp/test_hook_from_json_config.json

if [ ! ${ZCC_TEST_NAME[$FRAMEWORK]} ]; then
    echo "$FRAMEWORK is not a valid entry in list of frameworks"
    exit 1
fi
if [[ ${FRAMEWORK} == "tensorflow2" ]]; then
    echo "Running Gradtape test"
    if [ "$(printf '%s\n' "0.7.2" "$SMDEBUG_RELEASE_TAG" | sort -V | head -n1)" = "0.7.2" ]; then
        if [ "$(printf '%s\n' "0.8.0" "$SMDEBUG_RELEASE_TAG" | sort -V | head -n1)" = "0.8.0" ]; then
            pytest tests/zero_code_change/test_${FRAMEWORK}_gradtape_integration.py
            exit_code=$?
        else
            python tests/zero_code_change/${FRAMEWORK}_gradtape_integration_tests.py
            exit_code=$?
        fi
        if [ "$exit_code" -eq "0" ]
        then
            echo "SMDebug gradtape integration test succeeded"
        else
            echo "SMDebug gradtape integration test failed"
            exit $exit_code
        fi
    fi
fi

echo "Running ZCC tests."

pytest tests/zero_code_change/${ZCC_TEST_NAME[$FRAMEWORK]}

if [ "$?" -eq "0" ]
then
    echo "SMDebug ZCC tests succeeded."
else
    echo "SMDebug ZCC tests Failed."
    exit 1
fi

echo "All SMDebug tests succeeded!"
exit 0