#!/usr/bin/env bash
set -euo pipefail

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
ENV_DIR="${ENV_DIR:-/workspace/envs/rlbench}"
THIRD_PARTY_DIR="${THIRD_PARTY_DIR:-/workspace/third_party}"

mkdir -p "$THIRD_PARTY_DIR"

python3.10 -m venv "$ENV_DIR"
source "$ENV_DIR/bin/activate"

python -m pip install --upgrade pip setuptools wheel
python -m pip install -r "$ROOT_DIR/setup/requirements_core.txt"

if [ ! -d "$THIRD_PARTY_DIR/PyRep" ]; then
  git clone https://github.com/markusgrotz/PyRep.git "$THIRD_PARTY_DIR/PyRep"
fi
if [ ! -d "$THIRD_PARTY_DIR/RLBench" ]; then
  git clone https://github.com/markusgrotz/RLBench.git "$THIRD_PARTY_DIR/RLBench"
fi
if [ ! -d "$THIRD_PARTY_DIR/YARR" ]; then
  git clone https://github.com/markusgrotz/YARR.git "$THIRD_PARTY_DIR/YARR"
fi
if [ ! -d "$THIRD_PARTY_DIR/AnyBimanual" ]; then
  git clone https://github.com/liyaxuanliyaxuan/AnyBimanual.git "$THIRD_PARTY_DIR/AnyBimanual"
fi

git -C "$THIRD_PARTY_DIR/PyRep" checkout b8bd1d7a3182adcd570d001649c0849047ebf197
git -C "$THIRD_PARTY_DIR/RLBench" checkout 8af748c51287989294e00c9c670e3330a0e35ed5
git -C "$THIRD_PARTY_DIR/YARR" checkout 6822ff78602c77878b27d4cfe759ce029c67bffb
git -C "$THIRD_PARTY_DIR/AnyBimanual" checkout 76024e48b0e9489101459e85bc909c126ec581b4

python -m pip install -e "$THIRD_PARTY_DIR/PyRep"
python -m pip install -e "$THIRD_PARTY_DIR/RLBench"
python -m pip install -e "$THIRD_PARTY_DIR/YARR"

source "$ROOT_DIR/setup/env_vars.sh"

echo "Environment bootstrapped."
echo "You still need a compatible CoppeliaSim install at \$COPPELIASIM_ROOT."
echo "After that, activate the env and source setup/env_vars.sh before running RLBench or AnyBimanual jobs."