#!/usr/bin/env bash
# Build the `thyme` conda env. Idempotent — safe to re-run.
#   bash demistify_project/setup_env.sh
set -uo pipefail

ENV_PREFIX=${ENV_PREFIX:-/mnt/bn/tns-algo-video-public-my2/yijiangli/conda/thyme}
DATA_ROOT=${DATA_ROOT:-/mnt/bn/tns-algo-video-public-my2/yijiangli/data}
THYME_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"

export PIP_CACHE_DIR="$DATA_ROOT/.cache/pip"
mkdir -p "$PIP_CACHE_DIR"

PIP="$ENV_PREFIX/bin/pip"
PY="$ENV_PREFIX/bin/python"

step() { echo; echo "### $* ###"; }

step "torch 2.6.0 + cu124"
$PIP install --no-cache-dir torch==2.6.0 torchvision==0.21.0 \
  --index-url https://download.pytorch.org/whl/cu124 || exit 1

step "numpy pin (ms-swift framework.txt wants <2.0)"
$PIP install "numpy<2.0" || exit 1

step "thyme / ms-swift (editable) + framework requirements"
$PIP install -e "$THYME_ROOT" || exit 1

step "training + eval extras"
# opencv-python >=4.11 declares numpy>=2, which collides with framework.txt's numpy<2.0.
# 4.10.0.84 is the newest build that still supports the numpy 1.x ABI.
# qwen-vl-utils >=0.0.13 dropped vision_process.IMAGE_FACTOR, which ms-swift's
# patch_qwen_vl_utils() sets -> AttributeError before training starts.
# 0.0.11 is the newest release that still exposes every attribute swift patches.
$PIP install \
  "qwen-vl-utils==0.0.11" \
  "opencv-python==4.10.0.84" \
  deepspeed \
  timm \
  huggingface_hub \
  accelerate \
  nest_asyncio \
  hf_transfer \
  || exit 1
# hf_transfer: this cluster exports HF_HUB_ENABLE_HF_TRANSFER=1 globally, and
# huggingface_hub hard-errors on any download if the package is absent.
# nest_asyncio is imported by vlmeval/dataset/utils/sarena.py but missing from
# VLMEvalKit's requirements.txt; without it `import vlmeval` fails outright.

step "flash-attn (prebuilt wheel matched to torch/abi)"
ABI=$($PY -c "import torch;print('TRUE' if torch._C._GLIBCXX_USE_CXX11_ABI else 'FALSE')")
FA_VER=2.7.4.post1
FA_WHL="flash_attn-${FA_VER}+cu12torch2.6cxx11abi${ABI}-cp310-cp310-linux_x86_64.whl"
FA_URL="https://github.com/Dao-AILab/flash-attention/releases/download/v${FA_VER}/${FA_WHL}"
echo "abi=$ABI url=$FA_URL"
$PIP install --no-cache-dir "$FA_URL" || echo "!! flash-attn wheel failed (training needs it; eval uses sdpa)"

step "VLMEvalKit requirements (best-effort, one package at a time)"
REQ="$THYME_ROOT/eval/VLMEvalKit/requirements.txt"
if ! $PIP install -r "$REQ"; then
  echo "bulk install failed -> falling back to per-package"
  grep -vE '^\s*(#|$)' "$REQ" | while read -r pkg; do
    $PIP install "$pkg" || echo "!! SKIPPED: $pkg"
  done
fi

step "re-pin constraints clobbered by transitive deps"
# VLMEvalKit's requirements pull unpinned numpy/opencv, which drag numpy>=2 back in.
$PIP install "numpy<2.0" "opencv-python==4.10.0.84" "transformers>=4.49,<4.53" \
  "datasets>=3.0,<3.4" "trl>=0.15,<0.18" "peft>=0.11,<0.16" || exit 1

step "pip check (decord-on-this-platform warning is expected and harmless)"
$PIP check

step "done — versions"
$PIP list 2>/dev/null | grep -iE '^(torch|torchvision|transformers|flash|deepspeed|numpy|qwen|opencv|datasets|trl|peft|accelerate|ms-swift|swift)'