#!/usr/bin/env bash
# Prepare an isolated VBench-v1 environment for RTX 6000D (sm_120 / CUDA 13).
set -Eeuo pipefail

CONDA=${CONDA:-/root/.miniconda3/bin/conda}
SOURCE_ENV=${SOURCE_ENV:-sglang}
TARGET_ENV=${TARGET_ENV:-vbench}
TARGET_PREFIX=${TARGET_PREFIX:-/root/.miniconda3/envs/$TARGET_ENV}
VBENCH_REPO=${VBENCH_REPO:-/data/wxy/VBench}
LOG_ROOT=${LOG_ROOT:-/data/wxy/vbench_setup_logs}
WHEEL_DIR=${WHEEL_DIR:-/data/wxy/vbench_wheels}
PIP_INDEX_URL=${PIP_INDEX_URL:-https://mirrors.aliyun.com/pypi/simple}
export PIP_INDEX_URL

mkdir -p "$LOG_ROOT"
log() { printf '[%s] %s\n' "$(date '+%F %T')" "$*"; }

[[ -x "$CONDA" ]] || { log "missing conda: $CONDA"; exit 1; }
[[ -d "$VBENCH_REPO/vbench" ]] || { log "missing VBench repo: $VBENCH_REPO"; exit 1; }

if [[ ! -x "$TARGET_PREFIX/bin/python" ]]; then
  log "cloning $SOURCE_ENV to isolated environment $TARGET_ENV"
  "$CONDA" create --name "$TARGET_ENV" --clone "$SOURCE_ENV" --yes
fi

PYTHON="$TARGET_PREFIX/bin/python"
PIP=("$PYTHON" -m pip)
export CUDA_HOME=${CUDA_HOME_OVERRIDE:-/usr/local/cuda-13.2}
export TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST:-12.0}

log "installing only modules missing from the cloned CUDA-13 environment"
# --no-deps prevents pip from redownloading large packages already importable
# under another distribution name (notably the working cv2 build).
"${PIP[@]}" install --no-index --no-deps --no-build-isolation \
  "$WHEEL_DIR/matplotlib-3.11.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl" \
  "$WHEEL_DIR/openai-clip-1.0.1.tar.gz" \
  "$WHEEL_DIR/decord-0.6.0-py3-none-manylinux2010_x86_64.whl" \
  "$WHEEL_DIR/pyiqa-0.1.16-py3-none-any.whl" \
  "$WHEEL_DIR/lvis-0.5.3-py3-none-any.whl" \
  "$WHEEL_DIR/fairscale-0.4.13.tar.gz" \
  "$WHEEL_DIR/fvcore-0.1.5.post20221221.tar.gz" \
  "$WHEEL_DIR/boto3-1.43.80-py3-none-any.whl" \
  "$WHEEL_DIR/pycocoevalcap-1.2-py3-none-any.whl"

log "installing Detectron2 from source (isolated environment only)"
"${PIP[@]}" install --no-deps --no-build-isolation \
  'detectron2@git+https://github.com/facebookresearch/detectron2.git'

log "validating imports and RTX 6000D CUDA execution"
cd "$VBENCH_REPO"
"$PYTHON" - <<'PY'
import importlib
import torch

assert torch.cuda.is_available()
x = torch.randn(1024, 1024, device="cuda")
y = x @ x
torch.cuda.synchronize()
print("torch", torch.__version__, "cuda", torch.version.cuda,
      "capability", torch.cuda.get_device_capability(), "matmul", tuple(y.shape))
for name in ("vbench", "decord", "clip", "pyiqa", "detectron2"):
    importlib.import_module(name)
    print("import", name, "ok")
PY

touch "$LOG_ROOT/DONE"
log "VBench environment ready: $TARGET_PREFIX"