#!/usr/bin/env bash

set -eo pipefail

if [ "${BUILD_CONFIG}" = "2.0" ]; then
    TF_VERSION=v2.0.0-beta1
    TF_PATCH=0
    TF_CLONE=1
    BAZEL_VERSION=0.24.1
    BAZEL_BIN=none # force download
    BAZEL_COPT_FLAGS+=" --config=v2"
elif [ "${BUILD_CONFIG}" = "1.14" ]; then
    TF_VERSION=v1.14.0
    TF_PATCH=0
    TF_CLONE=1
    BAZEL_VERSION=0.24.1
    BAZEL_BIN=none
fi

# Release tag or branch name
TF_VERSION=${TF_VERSION:-"v1.13.1"}
TF_PATCH=${TF_PATCH:-"1"}
# Shorthand
# : ${TF_VERSION:="v1.13.1"}
# : ${TF_PATCH:="1"}

# Clone a fresh copy
if [ -n "${TF_CLONE}" ]; then
    TF_SRCDIR="/tmp/build-tflite/tensorflow/${TF_VERSION}/src"
    TF_OUTDIR="/tmp/build-tflite/tensorflow/${TF_VERSION}/out"
    if [ ! -d "${TF_SRCDIR}" ]; then
        git clone --branch ${TF_VERSION} --depth 1 https://github.com/tensorflow/tensorflow ${TF_SRCDIR}
    fi

    # Download patches
    if [ "${TF_PATCH}" = "1" ]; then
        TF_PATCH_DIR="/tmp/build-tflite/tensorflow/${TF_VERSION}/patches"
        TF_PATCHES=("0001-tflite-allow-fp16-for-fp32-models.patch" "0002-tflite-label-image-log-timing-dist.patch")
        for patch in "${TF_PATCHES[@]}"; do
            wget -q -nc --show-progress https://gist.github.com/muendelezaji/47b32812a15b622f05019b0e433de9b1/raw/${patch} -P ${TF_PATCH_DIR}
        done
    fi
fi

# Alternatively, provide path to cloned TensorFlow repo
TF_SRCDIR="${TF_SRCDIR:-$(readlink -f "src")}"
TF_OUTDIR="${TF_OUTDIR:-$(readlink -f "out")}"           # Build outputs folder
TF_PATCH_DIR="${TF_PATCH_DIR:-$(readlink -f "patches")}" # Folder with patches
# Shorthand
# : ${TF_SRCDIR:=$(readlink -f "src")}
# : ${TF_OUTDIR:=$(readlink -f "out")}
# : ${TF_PATCH_DIR:=$(readlink -f "patches")}

# Sanity checks
if [ ! -d "${TF_SRCDIR}" ]; then
    echo "ERROR: Cannot find TensorFlow source directory. Ensure it exists or set TF_CLONE to do a new clone."
    exit 1
elif [ ! -f "${TF_SRCDIR}/configure" -o ! -f "${TF_SRCDIR}/BUILD" ]; then
    echo "ERROR: Missing configure or BUILD file in ${TF_SRCDIR}. Ensure it is a valid Bazel workspace."
    exit 1
elif [ "${TF_PATCH}" = "1" -a ! -d "${TF_PATCH_DIR}" ]; then
    echo "ERROR: Patches directory was not found. It must exist when TF_PATCH is set to 'yes'."
    exit 1
elif [ -z "${ANDROID_NDK_HOME}" -o -z "${ANDROID_SDK_HOME}" ]; then
    echo "ERROR: Cannot locate Android NDK or SDK. Ensure both ANDROID_NDK_HOME and ANDROID_SDK_HOME are set."
    exit 1
fi

pushd ${TF_SRCDIR}

# Apply required patches
if [ "${TF_PATCH}" = "1" ]; then
    for patch in $(find ${TF_PATCH_DIR} -type f -name "*.patch" | sort); do
        git apply ${patch}
    done
    git config --local user.name temp
    git config --local user.email temp@example.com
    git commit -am "[TFLite] All patched up and ready to go"
fi

# Bazel
BAZEL_BIN=${BAZEL_BIN:-"$(command -v bazel || true)"}
BAZEL_JOBS=${BAZEL_JOBS:-"4"}
# BAZEL_PATCH="no"
# BAZEL_AVAIL_RAM=2048 # MB
# BAZEL_AVAIL_CPU=4.0  # number of cpu cores (1.0 representing a single full core)
# BAZEL_AVAIL_IO=1.0   # workstation I/O capability (with 1.0 representing average workstation)
# if [ -n "${BAZEL_AVAIL_RAM}" -a -n "${BAZEL_AVAIL_CPU}" -a -n "${BAZEL_AVAIL_IO}" ]; then
#     BAZEL_RESOURCES="--local_resources ${BAZEL_AVAIL_RAM},${BAZEL_AVAIL_CPU},${BAZEL_AVAIL_IO}"
# fi
if [ -n "${BAZEL_DRYRUN}" ]; then
    BAZEL_DEBUG_FLAGS+=" --nobuild"
fi
if [ -n "${BAZEL_VERBOSE}" ]; then
    BAZEL_DEBUG_FLAGS+=" -s --explain=${TF_OUTDIR}/explain.log"
fi
if [ -n "${BAZEL_DEBUG_FLAGS}" ]; then
    set -x
fi

# Download portable binary if not present
if [ ! -x "${BAZEL_BIN}" ]; then
    BAZEL_VERSION=${BAZEL_VERSION:-"0.24.1"}
    BAZEL_FILE="bazel-${BAZEL_VERSION}-linux-x86_64"
    BAZEL_DIR="/tmp/build-tflite/bazel/${BAZEL_VERSION}"

    wget -q -nc --show-progress https://github.com/bazelbuild/bazel/releases/download/${BAZEL_VERSION}/${BAZEL_FILE} -P ${BAZEL_DIR}
    chmod +x ${BAZEL_DIR}/${BAZEL_FILE}
    ln -sf ${BAZEL_DIR}/${BAZEL_FILE} ${BAZEL_DIR}/bazel
    export PATH="${BAZEL_DIR}:${PATH}"
    BAZEL_BIN="$(command -v bazel)"
fi

# Python
PYTHON_VERSION=${PYTHON_VERSION:-"2"}
export PYTHON_BIN_PATH="$(command -v python${PYTHON_VERSION})"
export USE_DEFAULT_PYTHON_LIB_PATH="1"

# Android
export ANDROID_BUILD_TOOLS_VERSION="28.0.3"
export ANDROID_API_LEVEL="28"
export ANDROID_NDK_API_LEVEL=${ANDROID_API_LEVEL}
NDK_VERSION="$(sed -En 's/Pkg.Revision = ([0-9]+).*/\1/p' ${ANDROID_NDK_HOME}/source.properties)"

# Select TensorFlow features
TF_BUILD_VARS+=" TF_NEED_CUDA=0"
TF_BUILD_VARS+=" TF_NEED_OPENCL_SYCL=0"
TF_BUILD_VARS+=" TF_NEED_ROCM=0"
TF_BUILD_VARS+=" TF_NEED_JEMALLOC=1"
TF_BUILD_VARS+=" TF_NEED_MPI=0"
TF_BUILD_VARS+=" TF_ENABLE_XLA=0"
TF_BUILD_VARS+=" TF_DOWNLOAD_CLANG=0"
TF_BUILD_VARS+=" TF_SET_ANDROID_WORKSPACE=1"
export ${TF_BUILD_VARS}

# Build options
# Disable unused default-on features
BAZEL_COPT_FLAGS+=" --config=noaws --config=nogcp --config=nohdfs --config=nokafka --config=noignite --config=nonccl --define tensorflow_mkldnn_contraction_kernel=0"
# Android config is equivalent to:
# --cpu=arm64-v8a --fat_apk_cpu=arm64-v8a --crosstool_top=//external:android/crosstool --host_crosstool_top=@bazel_tools//tools/cpp:toolchain
BAZEL_COPT_FLAGS+=" --config=android_arm64 --config=monolithic --config=opt --cxxopt=-std=c++11"
# Enable RUY
if [ "${TF_ENABLE_RUY}" = "1" ]; then
    BAZEL_COPT_FLAGS+=" --define tflite_with_ruy=true"
else
    BAZEL_COPT_FLAGS+=" --define tflite_with_ruy=false"
fi

# TFLite targets for Android
BAZEL_TARGETS+=" //tensorflow/lite:libtensorflowlite.so"
BAZEL_TARGETS+=" //tensorflow/lite/examples/label_image:label_image"
BAZEL_TARGETS+=" //tensorflow/lite/tools/benchmark:benchmark_model"

build_variant() {
    local variant="${1}"
    local ccflags="${2}"
    # Output path
    local variant_name="${TF_VERSION}_ruy${TF_ENABLE_RUY}_${variant}"
    local variant_path="${TF_OUTDIR}/${variant_name}"

    # GCC build optimisations - autoloaded during Bazel config step
    # For more info, see https://gcc.gnu.org/onlinedocs/gcc/ARM-Options.html
    export CC_OPT_FLAGS=" -O3 -DTFLITE_PROFILING_ENABLED ${ccflags}"

    # Configure and build
    if [ "${BAZEL_REBUILD}" != "1" ]; then
        ${BAZEL_BIN} clean
    fi
    ./configure
    ${BAZEL_BIN} build ${BAZEL_DEBUG_FLAGS} ${BAZEL_RESOURCES} -c opt --verbose_failures --jobs=${BAZEL_JOBS} ${BAZEL_COPT_FLAGS} ${BAZEL_TARGETS}

    if [ "${BAZEL_DRYRUN}" != "1" ]; then
        # Move to output dir
        mkdir -p ${variant_path}
        mv -f -t ${variant_path} \
            bazel-bin/tensorflow/lite/libtensorflowlite.so \
            bazel-bin/tensorflow/lite/examples/label_image/label_image \
            bazel-bin/tensorflow/lite/tools/benchmark/benchmark_model
        echo "Builds saved to: ${variant_path}"
    fi
}

# Generic v8-a with SIMD instructions
build_variant "v8a" "-march=armv8-a+simd"
# Optimised FP16 instructions
build_variant "v8.2a" "-march=armv8.2-a+fp16+dotprod"