#!/bin/bash

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

if [[ ! -d $LOG_DIR ]]; then 
    mkdir -p $LOG_DIR 
fi

# Use this to conditionally check TF version
pip install packaging
if [[ $(python -c "import tensorflow as tf; from packaging.version import Version; is_tf27_or_greater = Version(tf.__version__) >= Version('2.7'); print(is_tf27_or_greater)") == 'True' ]]; then
    PROFILER_LOG_DIR="/logs/plugins/profile"
else
    PROFILER_LOG_DIR="/logs/train/plugins/profile"
fi

set -e

export KERAS_BACKEND="tensorflow"

if [ -z "$KERAS_BACKEND" ]; then
    echo "usage: testBackendWithKeras theano OR testBackendWithKeras tensorflow OR testBackendWithKeras mxnet OR cntk"
    exit 1
fi

# tf1 test
if [[ $(python -c 'import tensorflow as tf; print(tf.__version__)') == 1.*.* ]]; then
    echo "Running tf1 keras test!"
    ${BIN_DIR}/testBackendWithKeras || exit 1
    exit 0
fi

TRAINING_LOG=${LOG_DIR}/${KERAS_BACKEND}_train_mnist.log

echo "Training mnist using $KERAS_BACKEND... This may take a few minutes. You can follow progress on the log file : $TRAINING_LOG"
set +e

[ -d $PROFILER_LOG_DIR ] && echo "Cleanning existing logs" && rm -rf $PROFILER_LOG_DIR

python ${BIN_DIR}/testTensorflow.py >$TRAINING_LOG 2>&1
RETURN_VAL=`echo $?`
set -e

echo "Verify tensorflow profiler can generate logs."

# profiler: tf > 2.* and < 2.2 generates local.trace where as tf > 2.2 generates *.pb
if [[ $(python -c "import tensorflow as tf; from packaging.version import Version; is_less_than_tf22 = Version(tf.__version__) < Version('2.2'); print(is_less_than_tf22)") == 'True' ]]; then
    LOG_RESULT=$(find $PROFILER_LOG_DIR -name "local.trace")
else
    LOG_RESULT=$(find $PROFILER_LOG_DIR -name "*.pb")
fi

echo "Checking the logs result: $LOG_RESULT"

[ -z "$LOG_RESULT" ] && echo "Cannot find Profiler logs!" && exit 1

if [ ${RETURN_VAL} -eq 0 ]; then
    echo "Training mnist Complete using $KERAS_BACKEND."
    cat $TRAINING_LOG

    # TensorFlow time to time might output different warnings after printing the accuracy. This makes the test flaky.
    if [ "$KERAS_BACKEND" !=  "tensorflow" ]; then
        ACCURACY=0.50
        for i in $(sed "s/Test accuracy: //g" $TRAINING_LOG); do
            ACCURACY=$i
        done

        ACCURACY=`python -c "print(str($ACCURACY*100))"`
        ACCURACY=${ACCURACY%.*}

        if [[ $ACCURACY -lt 95 ]]; then
          echo "Failed Validation Accuracy using $KERAS_BACKEND: $ACCURACY"
          echo "Training Log: "
          cat $TRAINING_LOG
          exit 1
        fi
        echo "Successful Validation Accuracy using $KERAS_BACKEND: $ACCURACY"
    fi
else
    echo "Training mnist Failed using $KERAS_BACKEND. LOG:"
    cat $TRAINING_LOG
    exit 1
fi

#Delete the downloaded mnist database
set +e
rm -rf ${HOME_DIR}/.keras/datasets
set -e
exit 0