#!/bin/bash

echo "Running the sanity check for the profiler enabled smdebug binary."
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
HOME_DIR=/test
BIN_DIR=${HOME_DIR}/bin
LOG_DIR=${HOME_DIR}/logs

echo "Verify whether smdebug with profiler is installed correctly."
SMDEBUG_RELEASE_TAG=$(python -c "import smdebug; print(smdebug.__version__)")

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

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

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 == "tensorflow2" ]] || [[ $FRAMEWORK == "pytorch" ]]; then
    for test_file in tests/profiler/${FRAMEWORK}/*
    do
        python -m pytest --assert=plain $test_file
        exit_code=$?
        if [ "$exit_code" -eq "0" ]
        then
            echo "SMProfiler 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 "SMProfiler no tests collected: $test_file"
        else
            echo "SMProfiler test failed: $test_file"
            exit $exit_code
        fi
    done
    for test_file in  tests/profiler/core/*
    do
        python -m pytest --assert=plain $test_file
        exit_code=$?
        if [ "$exit_code" -eq "0" ]
        then
            echo "SMProfiler core 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 "SMProfiler core no tests collected: $test_file"
        else
            echo "SMProfiler core test failed: $test_file"
            exit $exit_code
        fi
    done
else
    echo "SMProfiler is current not supported for this $FRAMEWORK"
fi

echo "Running ZCC tests."

echo "Setting up SMDebug and SMProfiler 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 [[ ${FRAMEWORK} == "tensorflow2" ]]; then
    export SMPROFILER_CONFIG_PATH=tests/core/json_configs/test_tf2_profiler_config_parser_by_time.json
    echo "Setting the SMPROFILER_CONFIG_PATH to $SMPROFILER_CONFIG_PATH"
    echo "Running Gradtape test"
    python -m pytest --assert=plain tests/zero_code_change/test_${FRAMEWORK}_gradtape_integration.py
    exit_code=$?
    if [ "$exit_code" -eq "0" ]
    then
        echo "SMProfiler gradtape integration test succeeded"
    else
        echo "SMProfiler gradtape integration test failed"
        exit $exit_code
    fi
fi

if [[ ${FRAMEWORK} == "pytorch" ]]; then
    export SMPROFILER_CONFIG_PATH=tests/core/json_configs/test_pytorch_profiler_config_parser.json
    echo "Setting the SMPROFILER_CONFIG_PATH to $SMPROFILER_CONFIG_PATH"
    echo "Running pytorch zcc integration test"
    python -m pytest --assert=plain "tests/zero_code_change/test_pytorch_integration.py::test_pytorch[True-False]"
    exit_code=$?
    if [ "$exit_code" -eq "0" ]
    then
        echo "SMProfiler pytorch [True-False] zcc integration test succeeded"
    else
        echo "SMProfiler pytorch [True-False] zcc integration test failed"
        exit $exit_code
    fi
    # Disabling test_pytorch[False-False]. See https://github.com/awslabs/sagemaker-debugger/issues/580
    : '
    python -m pytest --assert=plain "tests/zero_code_change/test_pytorch_integration.py::test_pytorch[False-False]"
    exit_code=$?
    if [ "$exit_code" -eq "0" ]
    then
        echo "SMProfiler pytorch [False-False] zcc integration test succeeded"
    else
        echo "SMProfiler pytorch [False-False] zcc integration test failed"
        exit $exit_code
    fi
    '
fi


echo "All SMprofiler tests succeeded!"
exit 0