#!/bin/bash

##Required variables: adjust variable values according to the environments
#For Ubuntu
export ROCM_VERSION=5.4
export ANACONDA_PYTHON_VERSION=3.8
export UBUNTU_VERSION=20.04
export BUILD_ENVIRONMENT=pytorch-linux-focal-rocm5.4-py3.8
export DEBIAN_FRONTEND=noninteractive

#For CentOS Stream 9
#export ROCM_VERSION=5.4
#export ANACONDA_PYTHON_VERSION=3.9
#export CENTOS_VERSION=9
#export BUILD_ENVIRONMENT=pytorch-centos9-rocm5.4-py3.9
#export BASH_ENV=/etc/profile

##Optional variables
# AMD GPU targets
export PYTORCH_ROCM_ARCH=gfx906
# (optional) Install protobuf for ONNX
PROTOBUF=yes
# (optional) Install database packages like LMDB and LevelDB
DB=yes
# (optional) Install vision packages like OpenCV and ffmpeg
VISION=yes
# (optional) Install a specific clang version
#export CLANG_VERSION=
# (optional) Install a specific gcc version
#export GCC_VERSION=

# Git Installation
os_value=$(grep -oP '(?<=^ID=).+' /etc/os-release)
if [[ "$os_value" == *ubuntu* ]]; then
  apt update && apt -y install git
elif [[ "$os_value" == *centos* ]]; then
  yum -y remove git
  yum -y remove git-*
  yum install -y git
fi

#Git clone Pytorch from a branch
git clone --recursive https://github.com/pytorch/pytorch -b main

#Navigate to docker directory
docker_dir="pytorch/.ci/docker"
if [[ -d "${docker_dir}" ]]; then
  cd ${docker_dir}
else
  cd pytorch/.circleci/docker
fi

#Modify install_conda.sh file
sed -i 's/chown jenkins/#chown jenkins/g' ./common/install_conda.sh
sed -i 's/as_jenkins / /g' ./common/install_conda.sh

# Install common dependencies (so that this step can be cached separately)
cp ./common/install_base.sh install_base.sh
bash ./install_base.sh
rm install_base.sh

# Langpack, devtoolset and ninja installation for CentOS
if [[ "$os_value" == *centos* ]]; then
	yum install -y glibc-langpack-en
	dnf install -y rpmdevtools
	dnf --enablerepo=crb install -y ninja-build
fi

# Install clang
cp ./common/install_clang.sh install_clang.sh
bash ./install_clang.sh
rm install_clang.sh

# Install conda and other packages (e.g., numpy, pytest)
export PATH=/opt/conda/bin:$PATH
mkdir -p /opt/conda
cp requirements-ci.txt /opt/conda/requirements-ci.txt
cp ./common/install_conda.sh install_conda.sh
bash ./install_conda.sh
rm install_conda.sh && rm /opt/conda/requirements-ci.txt

# Install gcc
cp ./common/install_gcc.sh install_gcc.sh
bash ./install_gcc.sh
rm install_gcc.sh

# (optional) Install protobuf for ONNX
cp ./common/install_protobuf.sh install_protobuf.sh
if [ -n "${PROTOBUF}" ]
then bash ./install_protobuf.sh
fi
rm install_protobuf.sh
export INSTALLED_PROTOBUF=${PROTOBUF}

# (optional) Install database packages like LMDB and LevelDB
cp ./common/install_db.sh install_db.sh
if [ -n "${DB}" ]
then bash ./install_db.sh
fi
rm install_db.sh
export INSTALLED_DB=${DB}

# (optional) Install vision packages like OpenCV and ffmpeg
cp ./common/install_vision.sh install_vision.sh
if [ -n "${VISION}" ]
then bash ./install_vision.sh
fi
rm install_vision.sh
export INSTALLED_VISION=${VISION}

# Install magma
cp ./common/install_rocm_magma.sh install_rocm_magma.sh
bash ./install_rocm_magma.sh
rm install_rocm_magma.sh
export PATH=/opt/rocm/bin:$PATH
export PATH=/opt/rocm/hcc/bin:$PATH
export PATH=/opt/rocm/hip/bin:$PATH
export PATH=/opt/rocm/opencl/bin:$PATH
export PATH=/opt/rocm/llvm/bin:$PATH
export MAGMA_HOME=/opt/rocm/magma
export LANG=C.UTF-8
export LC_ALL=C.UTF-8

# Install ccache/sccache (do this last, so we get priority in PATH)
cp ./common/install_cache.sh install_cache.sh
export PATH=/opt/cache/bin:$PATH
bash ./install_cache.sh
rm install_cache.sh

# Build PyTorch
cd ../..
BUILD_FILE=".ci/pytorch/build.sh"
if [[ -f "${BUILD_FILE}" ]]; then
    bash ${BUILD_FILE}
else
    bash .jenkins/pytorch/build.sh
fi