#!/usr/bin/env bash

##########################################################
# Disallow running with sudo or su
##########################################################
if [[ "$EUID" -eq 0 ]]
  then printf "\033[1;101mPlease do not run this script as root (no su or sudo)! \033[0m \n";
  exit;
fi
MACHINE_TYPE=`uname -m`
if [ ${MACHINE_TYPE} == 'aarch64' ];
  then printf "\033[1;101mThis application currently only supports x86_64 Chromebooks \033[0m \n";
  exit;
fi

###############################################################
## HELPERS
###############################################################
title() {
    printf "\033[1;42m";
    printf '%*s\n'  "${COLUMNS:-$(tput cols)}" '' | tr ' ' ' ';
    printf '%-*s\n' "${COLUMNS:-$(tput cols)}" "  # $1" | tr ' ' ' ';
    printf '%*s'  "${COLUMNS:-$(tput cols)}" '' | tr ' ' ' ';
    printf "\033[0m";
    printf "\n\n";
}

breakLine() {
    printf "\n";
    printf '%*s\n' "${COLUMNS:-$(tput cols)}" '' | tr ' ' -;
    printf "\n\n";
    sleep .5;
}

notify() {
    printf "\n";
    printf "\033[1;46m $1 \033[0m \n";
}

curlToFile() {
    notify "Downloading: $1 ----> $2";
    sudo curl -fSL "$1" -o "$2";
}

###############################################################
## REGISTERED VARIABLES
###############################################################
installedTensorflow=0;


###############################################################
## INSTALLATION
###############################################################

#INSTALL MINICONDA WITH TENSORFLOW
installConda() {
title "Installing Conda with Tensorflow and Keras and PyTorch";
    cd /tmp;
    curl -O https://repo.anaconda.com/archive/Anaconda3-2019.10-Linux-x86_64.sh;
    bash Anaconda3-2019.10-Linux-x86_64.sh -b -u -p $HOME/miniconda;
    rm Anaconda3-2019.10-Linux-x86_64.sh;
    cd ~;
    sudo apt  install jupyter-core
    ~/miniconda/bin/conda install -y -c conda-forge tensorflow keras;
    ~/miniconda/bin/conda install -y pytorch-cpu torchvision-cpu -c pytorch
    ~/miniconda/bin/conda install -c conda-forge jupyterlab
    ~/miniconda/bin/conda config --set auto_activate_base false
    ~/miniconda/bin/conda create -y -n tensorflow_env tensorflow jupyter keras pytorch-cpu torchvision-cpu -c pytorch;
    #conda config --set auto_activate_base true
    ~/miniconda/bin/conda init bash
    echo "conda activate tensorflow_env" > StartTFenv.sh
    chmod +x StartTFenv.sh
    source ~/.bashrc;

}

installCondaCUDA() {
title "Installing Conda with Tensorflow-GPU and Keras and PyTorch";
    cd /tmp;
    curl -O https://repo.anaconda.com/archive/Anaconda3-2019.10-Linux-x86_64.sh;
    bash Anaconda3-2019.10-Linux-x86_64.sh -b -u -p $HOME/miniconda;
    rm Anaconda3-2019.10-Linux-x86_64.sh;
    cd ~;
    sudo apt  install jupyter-core
    ~/miniconda/bin/conda install -y -c conda-forge tensorflow keras;
    ~/miniconda/bin/conda install -y pytorch-cpu torchvision-cpu -c pytorch
    ~/miniconda/bin/conda install -c conda-forge jupyterlab
    ~/miniconda/bin/conda config --set auto_activate_base false
    ~/miniconda/bin/conda create -y -n tensorflow_env tensorflow-gpu jupyter keras pytorch-cpu torchvision-cpu -c pytorch;
    #conda config --set auto_activate_base true
    ~/miniconda/bin/conda init bash
    echo "conda activate tensorflow_env" > StartTFenv.sh
    chmod +x StartTFenv.sh
    source ~/.bashrc;

}

#INSTALL OPENCV
installOpenCV() {
title "Building and installing OpenCV";
     OPENCV_VERSION='4.2.0'
     # 2. INSTALL THE DEPENDENCIES

     # Build tools:
     sudo apt-get install -y build-essential cmake

 
     sudo apt-get install -y libgtkglext1-dev libvtk6-dev

     # Media I/O:
     sudo apt-get install -y zlib1g-dev libjpeg-dev libwebp-dev libpng-dev libtiff5-dev libjasper-dev libopenexr-dev libgdal-dev

     # Video I/O:
     sudo apt-get install -y libdc1394-22-dev libavcodec-dev libavformat-dev libswscale-dev libtheora-dev libvorbis-dev libxvidcore-dev libx264-dev yasm libopencore-amrnb-dev libopencore-amrwb-dev libv4l-dev libxine2-dev

     # Parallelism and linear algebra libraries:
     sudo apt-get install -y libtbb-dev libeigen3-dev

     # Python:
     sudo apt-get install -y python-dev python-tk python-numpy python3-dev python3-tk python3-numpy

     # Java:
     sudo apt-get install -y ant default-jdk

     # Documentation:
     sudo apt-get install -y doxygen


     # 3. INSTALL THE LIBRARY

     wget https://github.com/opencv/opencv/archive/${OPENCV_VERSION}.zip
     unzip ${OPENCV_VERSION}.zip
     rm ${OPENCV_VERSION}.zip
     mv opencv-${OPENCV_VERSION} OpenCV
     cd OpenCV
     mkdir build
     cd build
     cmake -DWITH_OPENGL=ON -DFORCE_VTK=ON -DWITH_TBB=ON -DWITH_GDAL=ON -DWITH_XINE=ON  DBUILD_EXAMPLES=ON -DENABLE_PRECOMPILED_HEADERS=OFF ..
     make -j4
     sudo make install
     sudo ldconfig
}
installNCSudev() {
title "Setting udev Rules";
     sudo apt-get -y install udev cpio;
     sudo mkdir /etc/udev/rules.d/
     cd ~/
     cat <<EOF > 97-usbboot.rules
SUBSYSTEM=="usb", ATTRS{idProduct}=="2150", ATTRS{idVendor}=="03e7", GROUP="users", MODE="0666", ENV{ID_MM_DEVICE_IGNORE}="1"
SUBSYSTEM=="usb", ATTRS{idProduct}=="2485", ATTRS{idVendor}=="03e7", GROUP="users", MODE="0666", ENV{ID_MM_DEVICE_IGNORE}="1"
SUBSYSTEM=="usb", ATTRS{idProduct}=="f63b", ATTRS{idVendor}=="03e7", GROUP="users", MODE="0666", ENV{ID_MM_DEVICE_IGNORE}="1"
EOF
     sudo cp 97-usbboot.rules /etc/udev/rules.d/
     #sudo echo "ChromeOS" > /etc/lsb-release
     sudo udevadm control --reload-rules
     sudo udevadm trigger
     sudo ldconfig
}

installCoralEdgeTPU() {
title "Setting up Coral Edge TPU ";
sudo rm -rf /etc/apt/sources.list.d/coral-edgetpu.*
sudo apt-key adv --keyserver keyserver.ubuntu.com --recv 8B57C5C2836F4BEB
echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -
sudo apt-get update
sudo apt-get -y install libedgetpu1-std
}

installCoralEdgeTPUMAX() {
title "Setting up Coral Edge TPU Max ";
sudo rm -rf /etc/apt/sources.list.d/coral-edgetpu.*
echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -
sudo apt-get update
sudo apt-get -y install libedgetpu1-max
}

installTFSampleCode(){
title "Writing Sample - Tensorflow Hello World ";
     cd ~/
     cat <<EOF > SampleTFHelloWorld.py
import tensorflow as tf
# Simple hello world using TensorFlow by Keith I Myers
# https://KMyers.me | Keith@KMyers.me
print(hwinfo)


hello = tf.constant('Hello World, from TensorFlow!')

sess = tf.Session()

print(sess.run(hello))
hwinfo = tf.config.experimental.list_physical_devices(device_type=None)
EOF
}



###############################################################
## MAIN PROGRAM
###############################################################
sudo apt install -y dialog;

cmd=(dialog --backtitle "Crostini Machine Learning Setup Utility v00.1" \
--ascii-lines \
--clear \
--nocancel
--separate-output \
--checklist "Select installable packages:" 42 50 50);

options=(
    01 "Install Conda with TensorFlow and Keras" on
    02 "Install Conda with TensorFlow-GPU, CUDA and Keras(EXPERIMENTAL/Nvidia Only)" off
    03 "Build and Install OpenCV" on
    04 "Setup udev rules for Intel Movidius Neural Compute Stick 1/2" off
    05 "Setup Google Coral Edge TPU (Experimental)" off
    06 "Setup Google Coral Edge TPU High Powered Mode (Experimental)" off
    07 "Create Tensorflow Hello Sample Script" on
);

choices=$("${cmd[@]}" "${options[@]}" 2>&1 >/dev/tty);

#clear;

# Preparation
##########################################################
title "Installing Pre-Requisite Packages";
    cd ~/;
    sudo chown -R $(whoami) ~/
    sudo apt update;
    sudo apt dist-upgrade -y;
    sudo apt autoremove -y --purge;

    sudo apt install -y ca-certificates \
    apt-transport-https \
    software-properties-common \
    wget \
    git \
    curl \
    htop \
    mlocate \
    gnupg2 \
    cmake \
    libssh2-1-dev \
    libssl-dev \
    nano \
    vim \
    unzip;

    sudo updatedb;
breakLine;


title "Updating apt";
    sudo apt update;
    notify "The apt package manager is fully updated...";
breakLine;

for choice in ${choices}
do
    case ${choice} in
        01) installConda ;;
        02) installCondaCUDA ;;
        03) installOpenCV ;;
        04) installNCSudev ;;
        05) installCoralEdgeTPU ;;
        06) installCoralEdgeTPUMAX ;;
        07) installTFSampleCode ;;
    esac
done

# Clean
##########################################################
title "Finalising & Cleaning Up...";
    sudo chown -R $(whoami) ~/;
    sudo apt --fix-broken install -y;
    sudo apt dist-upgrade -y;
    sudo apt autoremove -y --purge;
breakLine;

notify "Installation Complete";


###############################################################
## POST INSTALLATION ACTIONS
###############################################################


  xdg-open 'https://kmyers.me/help/chromeos-crostini-ml/'
echo "";