#!/usr/bin/env bash
# Check everything run_aug32.sh needs before committing a multi-hour batch to a GPU.
# Exits non-zero if anything required is missing.
set -uo pipefail

HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
VACE="$(cd "$HERE/.." && pwd)"
fail=0

ok()   { printf '  \033[32mOK  \033[0m %s\n' "$1"; }
bad()  { printf '  \033[31mMISS\033[0m %s\n' "$1"; fail=1; }
have() { [ -e "$1" ] && ok "$2" || bad "$2  ($1)"; }
runs() { [ -x "$1" ] && ok "$2  ($("$1" -V 2>&1))" || bad "$2  ($1)"; }

echo "환경 (가상환경 — 새로 만들 필요 없음)"
runs "${VACE_PY_ROBOCASA:-/lp-dev/jonghoon/robocasa_calib/envs/robocasa/bin/python}" "batch driver (robocasa, mujoco)"
runs "${VACE_PY_WAN:-/lp-dev/jonghoon/VACE/venv/bin/python}" "WAN inference (torch+flash_attn)"
have "${VACE_RENDER_SCRIPT_DIR:-$HERE/vendor}/render_mesh_views_side.py" "render_mesh_views_side.py (vendored)"
"${VACE_RENDER_PYTHON:-${VACE_PY_ROBOCASA:-/lp-dev/jonghoon/robocasa_calib/envs/robocasa/bin/python}}" \
    -c "import trimesh,pyrender" 2>/dev/null && ok "trimesh + pyrender (mesh view rendering)" \
    || bad "trimesh + pyrender missing from the render python"

echo
echo "VACE 코드와 가중치"
have "$VACE/VACE/vace/vace_wan_inference.py" "VACE fork (ali-vilab/VACE + batch flags)"
have "$VACE/VACE/models/Wan2.1-VACE-1.3B"    "Wan2.1-VACE-1.3B weights (via models symlink)"
have "$VACE/VACE/models/VACE-Annotators"     "VACE-Annotators"
have "$VACE/vace_object_insert_batch.py"     "batch driver"

echo
echo "데이터셋"
have "${VACE_ROBOCASA_ROOT:-/home/nvidia/jonghoon/robocasa_full/pickplace_target_human}" "소스 에피소드 (영상 + extras)"
have "${VACE_AUG_MATERIALS:-/home/nvidia/jonghoon/robocasa_full/aug_materials/robocasa}"  "GT target-object 마스크"
for r in $(echo "${VACE_ASSET_ROOTS:-/lp-dev/jonghoon/robocasa_calib/repos/robocasa/robocasa/models/assets/objects/lightwheel:/lp-dev/jonghoon/robocasa_calib/repos/robocasa/robocasa/models/assets/objects/objaverse}" | tr ":" " "); do
  have "$r" "메쉬 $(basename "$r")"
done

echo
echo "GT 마스크 커버리지 (prep-worker 가 요구 — 없으면 그 에피소드는 실패)"
A="${VACE_AUG_MATERIALS:-/home/nvidia/jonghoon/robocasa_full/aug_materials/robocasa}"
for t in PickPlaceCounterToCabinet PickPlaceCounterToStove PickPlaceDrawerToCounter \
         PickPlaceSinkToCounter PickPlaceToasterToCounter; do
  n=0
  for d in "$A/$t"/ep*; do
    [ -f "$d/mask_gt_target_object_robot0_agentview_left.mp4" ] &&
    [ -f "$d/mask_gt_target_object_robot0_agentview_right.mp4" ] &&
    [ -f "$d/mask_gt_target_object_robot0_eye_in_hand.mp4" ] && n=$((n+1))
  done 2>/dev/null
  printf '  %-30s %4d episodes\n' "$t" "$n"
done

echo
echo "GPU"
nvidia-smi --query-gpu=index,memory.used,memory.total,utilization.gpu --format=csv,noheader 2>/dev/null |
  awk -F', ' '{printf "  gpu %s  %s / %s  util %s\n",$1,$2,$3,$4}'
echo "  필요: 약 27 GB (Wan2.1-VACE-1.3B, --offload_model False 기준)"

echo
[ $fail -eq 0 ] && echo "모두 준비됨." || echo "누락 항목 있음 — 위 MISS 확인."
exit $fail