#!/bin/bash
# 此脚本用来编译安装TensorRT-8.2.0软件
# TensorRT-8.2.0软件介绍
#  
#  Usage:
#     bash /path/to/build.sh 
# 使用build.sh 脚本样例 bash /path/to/build.sh 
# 使用build.sh 需要root权限或者sudo权限,否则无法安装依赖
#
##############################################################
set -e
# 导入公共方法脚本
source "$(
  cd "$(dirname "$0")" || exit 1
  pwd
)/common.sh"

function install_development() {
  sudo yum install -y time patch  environment-modules systemd-devel zlib-devel zlib bc gcc-c++
  source /etc/profile
  module purge
}

# 创建环境变量脚本
function write_TensorRT_run() {
  cat >"$1/TensorRT-8.2.0/run.sh" <<EOF
#!/bin/bash
sudo yum install -y m4 environment-modules time wget tar git zlib zlib-devel autoconf automake libtool gcc-c++

source /etc/profile &> /dev/null
current_dir="\$(cd "\$(dirname "\${BASH_SOURCE[0]}")" && pwd -P)"

for f in "\${current_dir}"/*; do
  if [[ -d "\${f}" && -f "\$(ls "\${f}"/*modulefiles 2>&1)" ]]; then
    module use "\${f}" && module load "\${f}"/*modulefiles
  fi
done
cd \$current_dir
if [[ ! -d "$installpath/hmpi-1.1.1" ]]; then
    mkdir -p $installpath/hmpi-1.1.1
    cp -r \$current_dir/hmpi-1.1.1/* $installpath/hmpi-1.1.1/
fi
echo -e "\033[1;32;1mTensorRT environment initialization completed.\033[0m"
EOF

  chmod +x "$1"/TensorRT-8.2.0/run.sh
}

# 安装依赖软件
<<install_dependency_software
  description: 安装依赖软件
  Parameters:
    param1 - 脚本名称
    param2 - 依赖名称
    param3 - modulefiles名称
install_dependency_software
function install_dependency_software() {
  bash "$dep_dir"/"$1".sh "$buildpath" "$installpath"
  module use "$installpath"/"$2"
  module load "$installpath"/"$2"/"$3"_modulefiles
}

# 安装毕晟编译器
function install_bisheng() {
  install_dependency_software "bisheng-compiler-2.1.0" "bisheng-compiler-2.1.0" "bisheng"
  module use "$installpath"/bisheng-compiler-2.1.0
  module load "$installpath"/bisheng-compiler-2.1.0/bisheng_modulefiles
}

# 安装hmpi-1.1.1
function install_hmpi() {
  install_dependency_software "hmpi-1.1.1" "hmpi-1.1.1" "hmpi"
}

# 安装cuda-11.4 
function install_cuda() {
  bash "$dep_dir"/cuda-11.4.sh "$buildpath" "$installpath" "$dep_dir" "$(get_current_dir)"
  module use "$installpath"/cuda-11.4
  module load "$installpath"/cuda-11.4/cuda_modulefiles
}

# 安装cudnn 
function install_cudnn() {
  bash "$dep_dir"/cudnn-8.2.sh "$buildpath" "$installpath" "$dep_dir" "$(get_current_dir)"
}

#安装TensorRT-GA
function install_TensorRT-GA() {
  bash "$dep_dir"/TensorRT-GA-8.2.0.6.sh "$buildpath" "$installpath" "$dep_dir" "$(get_current_dir)"
}


#安装 TensorRT-OSS
function install_TensorRT-OSS() {
  bash "$4"/TensorRT-8.2.0.sh "$1" "$2" "$4"
  mkdir -p "$3"/TensorRT-8.2.0 && cp -r "$2"/* "$3"/TensorRT-8.2.0
}


# 将生成的hpc文件打包到pkg文件夹下
function tar_pkg() {
  echo "# 正在打包 TensorRT 文件 路径为$1"
  cd "$1" && tar -zcf TensorRT-8.2.0-hpc.tar.gz TensorRT-8.2.0
  echo -e "\033[1;32;1m# TensorRT-8.2.0 已经打包成功 文件路径为$1/TensorRT-8.2.0-hpc.tar.gz\033[0m"
}

function main() {

 # 校验执行脚本时的入参
  check_build_param "$1"
  print_current_running_script_info "build.sh"

  # 构建路径
  buildpath=$(create_dir "build")

  # 安装路径
  installpath=$(create_dir "install")

  # 依赖路径
  dep_dir=$(get_current_dir)/dependency

  # 打包路径
  pkg_dir=$(get_current_dir)/pkg

  install_development

  # 安装毕昇编译器
  install_bisheng
   if [ $? -ne 0 ]; then
    echo -e "\033[31m# 正在退出\033[0m"
    exit 1
  fi

  # 安装hmpi
  install_hmpi
   if [ $? -ne 0 ]; then
    echo -e "\033[31m# 正在退出\033[0m"
    exit 1
  fi

  install_cuda

  install_cudnn

  install_TensorRT-GA

  install_TensorRT-OSS "$buildpath" "$installpath" "$pkg_dir" "$(get_current_dir)"

  write_TensorRT_run "$pkg_dir"

  tar_pkg "$pkg_dir"
  echo " # ------------------------------------------------------------------------"
  echo "# 运行软件前请执行以下命令"
  echo -e "\033[1;32;1m# source ""${pkg_dir}""/TensorRT-8.2.0/run.sh\033[0m"
  echo "# 来加载 TensorRT-8.2.0 运行时环境变量"
  echo "# ------------------------------------------------------------------------"
  rm -rf "$buildpath" "$installpath"
}

main "$@"