#!/bin/bash

export http_proxy="http://192.168.10.100:10809"
export https_proxy="http://192.168.10.100:10809"
export ALL_PROXY="socks5://192.168.10.100:10808"


#
# setup environment
#
export DEBIAN_FRONTEND=noninteractive


#
# install prerequisites (many of these are for numpy)
#
apt-get update && \
    apt-get install -y --no-install-recommends \
            python3-pip \
		  python3-dev \
		  libopenblas-dev \
		  libopenmpi2 \
            openmpi-bin \
            openmpi-common \
		  gfortran \
    && rm -rf /var/lib/apt/lists/*

pip3 install setuptools Cython wheel
pip3 install numpy --verbose


#
# PyTorch (for JetPack 4.4 production release)
#
#  PyTorch v1.2.0 https://nvidia.box.com/shared/static/lufbgr3xu2uha40cs9ryq1zn4kxsnogl.whl (torch-1.2.0-cp36-cp36m-linux_aarch64.whl)
#  PyTorch v1.3.0 https://nvidia.box.com/shared/static/017sci9z4a0xhtwrb4ps52frdfti9iw0.whl (torch-1.3.0-cp36-cp36m-linux_aarch64.whl)
#  PyTorch v1.4.0 https://nvidia.box.com/shared/static/c3d7vm4gcs9m728j6o5vjay2jdedqb55.whl (torch-1.4.0-cp36-cp36m-linux_aarch64.whl)
#  PyTorch v1.5.0 https://nvidia.box.com/shared/static/3ibazbiwtkl181n95n9em3wtrca7tdzp.whl (torch-1.5.0-cp36-cp36m-linux_aarch64.whl)
#
export PYTORCH_URL="https://nvidia.box.com/shared/static/9eptse6jyly1ggt9axbja2yrmj6pbarc.whl"
export PYTORCH_WHL="torch-1.6.0-cp36-cp36m-linux_aarch64.whl"

wget --quiet --show-progress --progress=bar:force:noscroll --no-check-certificate ${PYTORCH_URL} -O ${PYTORCH_WHL} && \
    pip3 install ${PYTORCH_WHL} --verbose && \
    echo rm ${PYTORCH_WHL}


#
# torchvision
# - PyTorch v1.0 : torchvision v0.2.2
# - PyTorch v1.1 : torchvision v0.3.0
# - PyTorch v1.2 : torchvision v0.4.0
# - PyTorch v1.3 : torchvision v0.4.2
# - PyTorch v1.4 : torchvision v0.5.0
# - PyTorch v1.5 : torchvision v0.6.0
# - PyTorch v1.6 : torchvision v0.7.0
#
export TORCHVISION_VERSION="v0.7.0"
export PILLOW_VERSION="pillow<7"
export TORCH_CUDA_ARCH_LIST="5.3;6.2;7.2"

echo "torchvision version = $TORCHVISION_VERSION" && echo "pillow version = $PILLOW_VERSION" && echo "TORCH_CUDA_ARCH_LIST = $TORCH_CUDA_ARCH_LIST"

apt-get update && \
    apt-get install -y --no-install-recommends \
		  git \
		  build-essential \
            libjpeg-dev \
		  zlib1g-dev \
    && rm -rf /var/lib/apt/lists/*

git clone -b ${TORCHVISION_VERSION} https://github.com/pytorch/vision torchvision && \
    cd torchvision && \
    export BUILD_VERSION="${TORCHVISION_VERSION#v*}" && \
    python3 setup.py install && \
    cd ../ && \
    rm -rf torchvision && \
    pip3 install "${PILLOW_VERSION}"


#
# torchaudio
#

export TORCHAUDIO_VERSION="v0.6.0"

apt-get update && \
    apt-get install -y --no-install-recommends \
		  sox \
		  libsox-dev \
		  libsox-fmt-all \
    && rm -rf /var/lib/apt/lists/*

git clone -b ${TORCHAUDIO_VERSION} https://github.com/pytorch/audio torchaudio && \
    cd torchaudio && \
    python3 setup.py install && \
    cd ../ && \
    rm -rf torchaudio


#
# PyCUDA
#
export PATH="/usr/local/cuda/bin:${PATH}"
export LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH}"
echo "$PATH" && echo "$LD_LIBRARY_PATH"

pip3 install pycuda --verbose