#!/bin/bash

doguide=0

if [[ ! -e /opt/intel/oneapi/setvars.sh ]]; then
    echo "You don't have the intel oneAPI base toolkit installed."
    doguide=1
else
    echo "Would you like to review the links for resources for setting up oneAPI/your drivers?"
    read -p "Y/n: " yn
    if [[ $yn == [yY] ]]; then
        doguide=1
    fi
fi

if [[ $doguide == 1 ]]; then
    echo "First, follow this guide for installing your GPU drivers, if you haven't already:"
    echo "https://dgpu-docs.intel.com/installation-guides/ubuntu/ubuntu-jammy-arc.html"
    echo "Afterwards, follow this guide for installing the oneAPI base toolkit:"
    echo "https://www.intel.com/content/www/us/en/developer/tools/oneapi/base-toolkit-download.html?operatingsystem=linux&distributions=aptpackagemanager"
fi

# Python should be 3.10 specifically for now (3.9 might work but let's use 3.10)
python3.10 --version
if [[ $? != 0 ]]; then
    echo "Python 3.10 is needed, and will be installed..."
    sudo apt install python3.10
fi

# Dependency spam, most of this should be needed. libjpeg9 exists because otherwise torchvision complains.
sudo apt-get install build-essential intel-oneapi-mkl intel-level-zero-gpu level-zero intel-opencl-icd intel-media-va-driver-non-free libmfx1 libgl-dev intel-oneapi-compiler-dpcpp-cpp libjpeg9 numactl


# get local dependencies
git submodule init
git submodule update --remote
# setup venv
python3.10 -m venv venv
source ./venv/bin/activate
python3 -m pip install --upgrade pip # just to be safe

# IPEX
echo ""
echo "Currently, Intel does not provide appropriate wheels: no torchaudio wheels whatsoever, and the latest stable IPEX wheels use too much VRAM."
echo "Building the current master from source, or acquiring wheels from unofficial sources, is necessary."
echo ""
echo "Would you like the setup to install known working, custom built, home grown, corn fed wheels for you, also installing/using megatools?"
read -p "(No exits setup) Y/n:" yn
if [[ $yn == [yY] ]]; then
    sudo apt install megatools
    mkdir temp_wheels
    cd ./temp_wheels
    sudo apt-get install megatools
    megadl https://mega.nz/folder/LBgQSTyS#BPjGq8WEpjoZ-uQF7deqTg
    python -m pip install requests pillow expecttest hypothesis
    python -m pip install --force-reinstall torch-1.13.0a0+git49444c3-cp310-cp310-linux_x86_64.whl
    python -m pip install --force-reinstall --no-deps torchvision-0.14.1a0+5e8e2f1-cp310-cp310-linux_x86_64.whl
    python -m pip install --force-reinstall --no-deps torchaudio-0.13.1+b90d798-cp310-cp310-linux_x86_64.whl
    python -m pip install --force-reinstall intel_extension_for_pytorch-1.13.120+git947ab23-cp310-cp310-linux_x86_64.whl
    cd ..
    rm -rf ./temp_wheels
else
    echo "Exiting..."
    exit 1
fi

# Luckily, it seems that Intel's CCL is sufficiently good
python -m pip install oneccl_bind_pt -f https://developer.intel.com/ipex-whl-stable-xpu

# install requirements
python3 -m pip install -r ./modules/tortoise-tts/requirements.txt # install TorToiSe requirements
python3 -m pip install -e ./modules/tortoise-tts/ # install TorToiSe
python3 -m pip install -r ./modules/dlas/requirements.txt # instal DLAS requirements
python3 -m pip install -e ./modules/dlas/ # install DLAS
python3 -m pip install -r ./requirements.txt # install local requirements

rm *.bat

deactivate