#!/bin/bash

# fix keys
sudo killall apt apt-get
sudo rm /var/lib/apt/lists/lock
sudo rm /var/cache/apt/archives/lock
sudo rm /var/lib/dpkg/lock
sudo rm /var/lib/dpkg/lock-frontend
sudo dpkg --configure -a
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -

# install libraries
sudo apt-get -y install libreadline-dev 
sudo apt-get -y install software-properties-common
sudo apt-get -y install libprotobuf-dev 
sudo apt-get -y install protobuf-compiler
sudo apt-get -y install libpng-dev
sudo apt-get -y install libjpeg-dev

# install cmake
sudo apt-get -y purge cmake
cd ~
git clone https://github.com/Kitware/CMake.git
cd CMake
cmake -DCMAKE_USE_OPENSSL=OFF
./bootstrap; make; sudo make install

# clone torch
git clone https://github.com/torch/distro.git ~/torch --recursive
cd ~/torch

# fix torch for cuda 10
rm -fr cmake/3.6/Modules/FindCUDA*
cd extra/cutorch

# patch
echo 'diff --git a/lib/THC/THCAtomics.cuh b/lib/THC/THCAtomics.cuh
index 400875c..ccb7a1c 100644
--- a/lib/THC/THCAtomics.cuh
+++ b/lib/THC/THCAtomics.cuh
@@ -94,6 +94,7 @@ static inline __device__ void atomicAdd(long *address, long val) {
 }
 
 #ifdef CUDA_HALF_TENSOR
+#if !(__CUDA_ARCH__ >= 700 || !defined(__CUDA_ARCH__) )
 static inline  __device__ void atomicAdd(half *address, half val) {
   unsigned int * address_as_ui =
       (unsigned int *) ((char *)address - ((size_t)address & 2));
@@ -117,6 +118,7 @@ static inline  __device__ void atomicAdd(half *address, half val) {
    } while (assumed != old);
 }
 #endif
+#endif' >> atomic.patch
patch -p1 < atomic.patch

# install torch dependencies
cd ~/torch
bash install-deps;

# install torch
./clean.sh
export TORCH_NVCC_FLAGS="-D__CUDA_NO_HALF_OPERATORS__"
yes | ./install.sh

# add torch to path
source ~/.bashrc 
. /home/paperspace/torch/install/bin/torch-activate

# packages
sudo rm -rf ~/.cache/luarocks/
luarocks install loadcaffe
cd ~/torch/extra/cutorch
export TORCH_NVCC_FLAGS="-D__CUDA_NO_HALF_OPERATORS__"
luarocks make rocks/cutorch-scm-1.rockspec
cd ~
luarocks install cudnn
luarocks install cunn
luarocks install image

# get cudnn (https://developer.nvidia.com/cudnn) 
#wget https://s3.amazonaws.com/open-source-william-falcon/cudnn-9.0-linux-x64-v7.1.tgz  
#sudo tar -xzvf cudnn-9.0-linux-x64-v7.1.tgz  
#sudo cp cuda/include/cudnn.h /usr/local/cuda/include
#sudo cp cuda/lib64/libcudnn* /usr/local/cuda/lib64
#sudo chmod a+r /usr/local/cuda/include/cudnn.h /usr/local/cuda/lib64/libcudnn*

# get right cudnn luarock (assuming cudnn 7.*)
git clone https://github.com/soumith/cudnn.torch.git -b R7 && cd cudnn.torch && luarocks make cudnn-scm-1.rockspec && cd .. && rm -rf cudnn.torch

# get neural-style
cd ~
git clone https://github.com/jcjohnson/neural-style
cd neural-style
sh models/download_models.sh