#!/bin/bash

# Create a conda environment for EveryVoice development, automating all the
# manual instructions in README.md

# Default versions:
CUDA_VERSION=11.8
PYTHON_VERSION=3.10

usage() {
    for msg in "$@"; do
        echo $msg >&2
    done
    cat <<==EOF== >&2

Usage: ./make-everyvoice-env [options]

  Create a conda environment for EveryVoice, automating all the instructions in
  README.md

Options:
  -h, --help              Print this help message

Torch pre-compilation options:
  --cuda CUDA_VERSION     Install torch compiled for CUDA_VERSION
                          Default: --cuda $CUDA_VERSION
  --cpu                   Install torch for use on CPU only
  --python PYTHON_VERSION Specify the Python version to use
                          Default: --python $PYTHON_VERSION

Target Environment Sepcification:
  -n, --name ENV_NAME     Name of the conda environment to create
                          Default: --name EveryVoice
  -p, --prefix ENV_PATH   Path of the conda environment to create

==EOF==
    exit
}

error_exit() {
   echo -n "ERROR: " >&2
   for msg in "$@"; do
      echo $msg >&2
   done
   echo "Use -h for help." >&2
   exit 1
}

arg_check() {
   if [ $2 -le $1 ]; then
      error_exit "Missing argument to $3 option."
   fi
}

ENV_PREFIX=
ENV_NAME=
while [[ $# -gt 0 ]]; do
    case "$1" in
    -h|--help)   usage;;
    --cuda)      arg_check 1 $# $!; CUDA_VERSION=$2; shift;;
    --cpu)       CUDA_VERSION=cpu;;
    --python)    arg_check 1 $# $!; PYTHON_VERSION=$2; shift;;
    -n|--name)   arg_check 1 $# $!; ENV_NAME=$2; shift;;
    -p|--prefix) arg_check 1 $# $!; ENV_PREFIX=$2; shift;;
    -*)          error_exit "Unknown option $1.";;
    *)           break;;
    esac
    shift
done

# This can only be run from the root of an EveryVoice sandbox
if [[ "$0" != make-everyvoice-env && "$0" != ./make-everyvoice-env ]]; then
    error_exit "make-everyvoice-env only works from the root of an EveryVoice sandbox."
fi

# submodules have to have already been initialized
if git submodule status | grep -q "^-"; then
    error_exit "Please init the submodules with \"git submodule update --init\"."
fi

# Interpret the (mutually exclusive) env name or prefix
if [[ ! $ENV_PREFIX ]]; then
    [[ $ENV_NAME ]] || ENV_NAME=EveryVoice
    ENV_OPTION="--name \"$ENV_NAME\""
    ENV_DESC="called \"$ENV_NAME\""
    ENV2ACTIVATE=$ENV_NAME
else
    [[ $ENV_NAME ]] && error_exit "Please specify only one of --name and --prefix."
    ENV_OPTION="--prefix \"$ENV_PREFIX\""
    ENV_DESC="at path \"$ENV_PREFIX\""
    ENV2ACTIVATE=$ENV_PREFIX
fi

# Setup conda aliases
if __conda_setup="$(conda shell.bash hook)"; then
    eval "$__conda_setup"
else
    error_exit "Cannot initialize conda."
fi
unset __conda_setup

# Don't overwrite an existing env
if conda activate "$ENV2ACTIVATE" >& /dev/null; then
    error_exit "Environment \"$ENV2ACTIVATE\" already exists, please use a different name or path."
fi

if [[ $CUDA_VERSION == CPU || $CUDA_VERSION == cpu ]]; then
    CUDA_TAG=cpu
    CUDA_DESC="use on CPU only"
else
    if which nvidia-smi >& /dev/null && nvidia-smi | grep -q CUDA; then
        if nvidia-smi | grep -q "CUDA Version: $CUDA_VERSION "; then
            : # CUDA version OK
        else
            echo "Warning: Mismatched CUDA version found. Specified: CUDA_VERSION=$CUDA_VERSION. Found:"
            nvidia-smi | grep CUDA
            echo "Please make sure the CUDA version available at runtime will match $CUDA_VERSION."
        fi
    else
        echo "Please make sure the CUDA version installed on your system matches CUDA_VERSION=$CUDA_VERSION."
    fi
    CUDA_TAG=cu$(echo $CUDA_VERSION | sed 's/\.//g')
    CUDA_DESC="CUDA $CUDA_VERSION"
fi
echo "Creating EveryVoice conda environment $ENV_DESC for $CUDA_DESC using Python $PYTHON_VERSION."

echo -n "Proceed (y/[n])? "
read proceed
if [[ "$proceed" =~ ^[y|Y] ]]; then
    echo Proceeding
else
    echo Quitting
    exit 1
fi

r() {
    cmd=$*
    echo "\$ $cmd"
    eval $cmd
    return $?
}

set -o errexit

r conda create -y $ENV_OPTION python=$PYTHON_VERSION ffmpeg
eval "$(conda shell.bash hook)"
r conda activate "$ENV2ACTIVATE"

# Recent versions of lit do not always compile cleanly, but we can fall back to
# the one with a wheel on pytorch.org if necessary
if ! r pip install lit --find-links https://download.pytorch.org/whl/torch_stable.html; then
    echo Falling back to installing an older lit with a known pre-compiled wheel
    r pip install lit==15.0.7 --find-links https://download.pytorch.org/whl/torch_stable.html
fi

r CUDA_TAG=$CUDA_TAG pip install -r requirements.torch.txt --find-links https://download.pytorch.org/whl/torch_stable.html
r pip install cython

# pycountry and pyworld don't always compile cleanly, but we can fall back to
# using conda-forge if necessary
PY_COUNTRY_WORLD=$(grep -o "\(pycountry\|pyworld\)==[0-9.]\+" pyproject.toml)
if ! r pip install $PY_COUNTRY_WORLD; then
    echo Falling back to installing pycountry and pyworld from conda-forge
    r conda install -y $PY_COUNTRY_WORLD -c conda-forge
fi

r conda install -y sox -c conda-forge
r pip install -e .[dev]
echo ""
echo "Environment creation completed with success"

echo ""
echo "Configuring your sandbox in case you want to contribute to the project."
if ! pre-commit install; then
    echo "Error running \"pre-commit install\". Your \"$ENV2ACTIVATE\" environment is good, but if you want to submit contributions to the project, please rerun \"pre-commit install\" in your sandbox."
fi
if ! gitlint install-hook; then
    echo ""
    echo "Error running \"gitlint install-hook\". Your \"$ENV2ACTIVATE\" environment is good, but if you want to submit contributions to the project, please rerun \"gitlint install-hook\" in your sandbox."
fi
# Try to install pre-commit and gitlint hooks in the submodules but don't complain on failure
git submodule foreach 'pre-commit install || true' || true
git submodule foreach 'gitlint install-hook || true' || true

echo ""
echo "SUCCESS!"
echo "EveryVoice environment \"$ENV2ACTIVATE\" successfully created."
echo "Run \"conda activate $ENV2ACTIVATE\" to activate it."
echo "Run \"everyvoice/run_tests.py all\" to validate it."
