#!/bin/bash
#
# pyTorch install script for NVIDIA Jetson TX1/TX2,
# from a fresh flashing of JetPack 2.3.1 / JetPack 3.0 / JetPack 3.1
#
# for the full source, see jetson-reinforcement repo:
#   https://github.com/dusty-nv/jetson-reinforcement/blob/master/CMakePreBuild.sh
#
# note:  pyTorch documentation calls for use of Anaconda,
#        however Anaconda isn't available for aarch64.
#        Instead, we install directly from source using setup.py
sudo apt-get install python-pip

# upgrade pip
pip install -U pip
pip --version
# pip 9.0.1 from /home/ubuntu/.local/lib/python2.7/site-packages (python 2.7)


#cd $HOME
# clone pyTorch repo
git clone http://github.com/pytorch/pytorch
cd pytorch
git submodule update --init

# install prereqs
sudo pip install -U setuptools
sudo pip install -r requirements.txt

# Develop Mode:
#python setup.py build_deps
#sudo python setup.py develop

cd torch
ln -s _C.so lib_C.so
cd lib
#ln -s libTH.so.1 libTH.so
#ln -s libTHC.so.1 libTHC.so
cd ../../

git clone https://github.com/pytorch/examples
sudo pip install -r examples/reinforcement_learning/requirements.txt 

git clone https://github.com/pytorch/vision
cd vision
sudo python setup.py install

sudo apt-get install swig
sudo pip install box2D

cd ../../

echo "[Pre-build]  pyTorch setup complete"


# Install Mode:  (substitute for Develop Mode commands)
sudo python setup.py install

# Verify CUDA (from python interactive terminal)
# import torch
# print(torch.cuda.is_available())
# a = torch.cuda.FloatTensor(2)
# print(a)
# b = torch.randn(2).cuda()
# print(b)
# c = a + b
# print(c)