#!/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
    read -p "Would you like to review the links for resources for setting up oneAPI/your drivers?" 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


# 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
python -m pip install torch==1.13.0a0+git6c9b55e torchvision==0.14.1a0 intel_extension_for_pytorch==1.13.120+xpu -f https://developer.intel.com/ipex-whl-stable-xpu

# Non-mandatory torchaudio install
pip list | grep torchaudio
if [[ $? == 1 ]]; then
    echo "torchaudio is not installed. You may install one yourself (venv location is ./venv/bin/activate)"
    read -p "Would you like the setup to install a custom built wheel for you, also installing/using megatools? (No exits setup) Y/N:" yn
    if [[ $yn == [yY] ]]; then
        sudo apt install megatools
        mkdir temp_torchaudio_wheel
        cd ./temp_torchaudio_wheel
        megadl https://mega.nz/file/jc4SnCgS#_7vProAZGAw1l0vojLk7Reo7UZEUCeeHT0ye12Haow8
        pip install --force-reinstall --no-deps torchaudio-0.13.1+b90d798-cp310-cp310-linux_x86_64.whl
        cd ..
        rm -rf ./temp_torchaudio_wheel
    else
        echo "Exiting..."
        exit 1
    fi
fi

# 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

sed -i 's|bitch|BEATCH\nsource /opt/intel/oneapi/setvars.sh|' ./start.sh

deactivate