#!/usr/bin/env bash
# Copyright (c) Huawei Technologies Co., Ltd. 2020. All rights reserved.
###### Environment variable settings, need to set according to your own device ######

# ASCEND_OPP_PATH: The installation path of the OPP package, where "/usr/local/Ascend/opp" is the
#                  default installation path. If user defines the installation path, please modify it.
# Uncomment and modify it when you specified installation path of OPP.
#export ASCEND_OPP_PATH=/usr/local/Ascend/opp

# ASCEND_TENSOR_COMPLIER_INCLUDE: The path of the header file of the ATC package, where "/usr/local/Ascend/atc/include" is the
#                                 default installation path. If user defines the installation path, please modify it.
# Uncomment and modify it when you specified installation path of ATC.
#export ASCEND_TENSOR_COMPLIER_INCLUDE=/usr/local/Ascend/atc/include

# TOOLCHAIN_DIR: The path of the cross compilation tool, where "/usr/local/Ascend/toolkit/toolchain/hcc" is the
#                default installation path. If user defines the installation path, please modify it.
# Uncomment and modify it when you specified installation path of toolkit.
#export TOOLCHAIN_DIR=/usr/local/Ascend/toolkit/toolchain/hcc

# AICPU_KERNEL_TARGET: The name of the custom operator kernels library, which "cust_aicpu_kernels" is the
#                      default installation name. If user defines the installation name, please modify it.
# Uncomment and modify it when you specified the installation name of custom operator kernels library.
#export AICPU_KERNEL_TARGET=cust_aicpu_kernels

# The SYSTEM_INFO flag is the name of the form of the operator package generated by the compiler, which can be customized by the user, E.g:
#   a. If the operating system version is CentOS and the architecture is aarch64, it can be set to centos_aarch64, and the operator package name generated by compilation is custom_opp_centos_aarch64.run.
#   b. If the operating system version is CentOS and the architecture is x86_64, it can be set to centos_x86_64, and the name of the operator package generated by compilation is custom_opp_centos_x86_64.run.
#   c. If the SYSTEM_INFO environment variable is not set, the default value is used: ubuntu_x86_64, and the operator package name generated by compilation is custom_opp_ubuntu_x86_64.run.
# Uncomment and modify it when you need to specify os and architecture.
#export SYSTEM_INFO=centos_x86_64

###### The following logic can be used without modification ######

# parse input parameters
clean=n
while true
do
  case "$1" in
  -h | --help)
  echo "  --help                show the help info of using this shell"
  echo "  --kernel_so_name=*    so name of aicpu custom operators "
  echo "  --toolchain_dir=*     specify cross compile toolchain directory"
  echo "  clean                 clean compile result directory"
  exit 0
  ;;
  --so_name=*)
  so_name=`echo ${1} | cut -d"=" -f2`
  length=${#so_name}
  aicpu_target_name=${so_name:3:($length - 6)}
  export AICPU_KERNEL_TARGET=$aicpu_target_name
  echo "aicpu_target_name=${aicpu_target_name}"
  shift
  ;;
  clean)
  clean=y
  shift
  ;;
  --toolchain_dir=*)
  tool_dir=`echo ${1} | cut -d"=" -f2`
  export TOOLCHAIN_DIR=$tool_dir
  echo "specified TOOLCHAIN_DIR=$TOOLCHAIN_DIR"
  shift
  ;;
  -*)
  echo "  --help                show the help info of using this shell"
  echo "  --toolchain_dir=*     specify cross compile toolchain directory"
  echo "  --so_name=*           so name of aicpu custom operators generated by compilation"
  echo "  clean                 clean compile result directory"
  exit 0
  ;;
  *)
  break
  ;;
  esac
done

# The following logic can be used without modification ######
log() {
  cur_date=`date +"%Y-%m-%d %H:%M:%S"`
  echo "[$cur_date] "$1
}

project_path=$(cd "$(dirname "$0")"; pwd)
# clean compile result
if [ "x$clean" == "xy" ] 2>/dev/null; then
  rm -rf $project_path/caffe.proto 2>/dev/null
  rm -rf $project_path/caffe.proto.origin 2>/dev/null
  rm -rf $project_path/build_out 2>/dev/null
  rm -rf $project_path/framework/caffe_plugin/proto/caffe 2>/dev/null
  log "[INFO] Clean successfully."
  exit 0
fi

# if specified cross compile toolchain directory is not exist, log error and exit
if [[ ! -z "$TOOLCHAIN_DIR" ]] && [[ ! -d "$TOOLCHAIN_DIR" ]];then
    log "[ERROR] Specified cross compile toolchain directory is not exist"
    exit 1
fi

# set aicpu kernel implement compiler target, default to be custom_cpu_kernels
if [[ -z "${aicpu_target_name}" ]] && [[ -z "${AICPU_KERNEL_TARGET}" ]]; then
    log "[INFO] no so_name is specified, use default compile target: libcust_cpu_kernels.so!"
    aicpu_target_name="cust_cpu_kernels"
    export AICPU_KERNEL_TARGET=$aicpu_target_name
    #exit 1
fi

chmod -R 755 $project_path/cmake/util/
mkdir -p $project_path/build_out

###### Environment variable settings, need to set according to your own device ######

# ASCEND_TENSOR_COMPLIER_INCLUDE: The path of the header file of the ATC package, where "/usr/local/Ascend/atc/include" is the
#                                 default installation path. If user defines the installation path, please modify it.
# Uncomment and modify it when you specified installation path of ATC.
if [[ -z "${ASCEND_TENSOR_COMPLIER_INCLUDE}" ]]; then
    if [[ -d "/usr/local/Ascend/atc/include" ]];then
        export ASCEND_TENSOR_COMPLIER_INCLUDE=/usr/local/Ascend/atc/include
    else
        log "[ERROR] ENV ASCEND_TENSOR_COMPLIER_INCLUDE is not set"
        exit 1
    fi
else
    if [[ ! -d "${ASCEND_TENSOR_COMPLIER_INCLUDE}" ]];then
        log "[ERROR] ENV ASCEND_TENSOR_COMPLIER_INCLUDE=${ASCEND_TENSOR_COMPLIER_INCLUDE} dir is not exist"
        exit 1
    fi
fi

# The SYSTEM_INFO flag is the name of the form of the operator package generated by the compiler, which can be customized by the user, E.g:
#   a. If the operating system version is CentOS and the architecture is aarch64, it can be set to centos_aarch64, and the operator package name generated by compilation is custom_opp_centos_aarch64.run.
#   b. If the operating system version is CentOS and the architecture is x86_64, it can be set to centos_x86_64, and the name of the operator package generated by compilation is custom_opp_centos_x86_64.run.
#   c. If the SYSTEM_INFO environment variable is not set, the default value is used: ubuntu_x86_64, and the operator package name generated by compilation is custom_opp_ubuntu_x86_64.run.
# Uncomment and modify it when you need to specify os and architecture.
if [[ -z "${SYSTEM_INFO}" ]]; then
    export SYSTEM_INFO=centos_x86_64
fi

# STEP 1, Check plugin files and proto files
# Check for caffe_plugin files, If there is no plugin file, do not compile caffe proto.
caffe_plugin_cpp_files_num=$(ls $project_path/framework/caffe_plugin/*.cpp 2> /dev/null | wc -l)
if [[ "$caffe_plugin_cpp_files_num" == 0 ]];then
  log "[INFO] No caffe plugin files."
fi  # endif compile caffe proto

# STEP 2, CMake and make
cd $project_path/build_out
rm -rf *.run
log "[INFO] Cmake begin."
cmake ..
if [ $? -ne 0 ]; then
  log "[ERROR] Please check cmake result."
  exit 1
fi

log "[INFO] Make begin."
make -j
if [ $? -ne 0 ]; then
  log "[ERROR] Please check make result."
  exit 1
fi
log "[INFO] Build successfully, the target is in $project_path/build_out."