#!/bin/bash

# Copyright (c) 2020 Samsung Electronics Co., Ltd. All Rights Reserved
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

set -e

DRIVER_PATH="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

VENV_ACTIVATE=${DRIVER_PATH}/venv/bin/activate
# NOTE please use venv's python instead of python after `source activation`.
# This script is called by debian maintainer script, i.e. `postinst`.
# Since debian maintainer script is called with sudo, `source activation` is ignored.
VENV_PYTHON=${DRIVER_PATH}/venv/bin/python

if [ ! -f ${VENV_ACTIVATE} ]; then
  # Create python virtual enviornment
  python3 -m venv "${DRIVER_PATH}/venv"
fi

# NOTE version
# - https://github.com/onnx/onnx/blob/master/docs/Versioning.md
# - https://github.com/onnx/onnx-tensorflow/blob/master/Versioning.md

VER_TENSORFLOW=2.3.0
VER_ONNX=1.10.1
VER_ONNX_TF=1.9.0

# Install tensorflow

PIP_TRUSTED_HOST="--trusted-host pypi.org "
PIP_TRUSTED_HOST+="--trusted-host files.pythonhost.org "
PIP_TRUSTED_HOST+="--trusted-host download.pytorch.org "

PIP_TIMEOUT="--default-timeout=1000 "

PIP_OPTIONS="${PIP_TIMEOUT} ${PIP_TRUSTED_HOST}"

# NOTE $ONE_PREPVENV_PIP_OPTION is to provide additional PIP options
# such as ceritificate file behind firewall
# ex) ONE_PREPVENV_PIP_OPTION="--cert SomePrivateCetificate.crt" ./one-prepare-venv
if [[ ! -z "$ONE_PREPVENV_PIP_OPTION" ]]; then
  PIP_OPTIONS+=" ${ONE_PREPVENV_PIP_OPTION} "
fi

# TODO remove version number of 'pip==20.2.1 setuptools==49.3.0'
# NOTE adding version is for temporary hotfix of setuptools 50.x.y version
${VENV_PYTHON} -m pip ${PIP_OPTIONS} install -U pip==20.2.1 setuptools==49.3.0
${VENV_PYTHON} -m pip ${PIP_OPTIONS} install tensorflow-cpu==${VER_TENSORFLOW}
${VENV_PYTHON} -m pip ${PIP_OPTIONS} install Pillow==6.2.2

# Install PyTorch and ONNX related
${VENV_PYTHON} -m pip ${PIP_OPTIONS} install torch==1.8.1+cpu -f https://download.pytorch.org/whl/torch_stable.html

# Provide install of custom onnx-tf
if [ -n "${EXT_ONNX_TF_WHL}" ]; then
  ${VENV_PYTHON} -m pip ${PIP_OPTIONS} install onnx==${VER_ONNX} ${EXT_ONNX_TF_WHL}
else
  ${VENV_PYTHON} -m pip ${PIP_OPTIONS} install onnx==${VER_ONNX} onnx-tf==${VER_ONNX_TF}
fi