#!/usr/bin/env bash
# Параллельная unitree-фаза: build_shards.py (v3) однопоточный по своей природе,
# так что параллелим ПО СЕТАМ (xargs -P). Каждый сет: build -> pack -> validate,
# tmp-шарды удаляются сразу. Дисковый предохранитель: ждать, пока свободно <30GB.
set -uo pipefail
cd /mnt/storage/tinyvla_wds
export HF_HOME=/mnt/storage/tinyvla_wds/.hf_home
source venv/bin/activate
S=tinyvla_h100/scripts
mkdir -p data/wds/unitree data/shards_tmp logs

UNITREE_G1="G1_Dex1_StackRygBlock_Dataset_Sim G1_Dex1_DiverseManip_DualArm_256x256 \
G1_Dex1_MountCamera_Dataset G1_Dex3_ToastedBread_Dataset G1_Dex1_DiverseManip_SingleArm_256x256 \
G1_Brainco_PickDoll_Dataset G1_Dex3_PickDoll_Dataset G1_Dex3_BlockStacking_Dataset \
G1_Dex3_GraspSquare_Dataset G1_Dex3_CameraPackaging_Dataset G1_Brainco_GraspOreo_Dataset"
UNITREE_Z1="Z1_StackBox_Dataset Z1_Dual_Dex1_PourCoffee_Dataset Z1_Dual_Dex1_FoldClothes_Dataset"

one() {
  name=$1
  [ -f data/wds/unitree/$name/manifest.json ] && { echo "skip $name (done)"; return 0; }
  case " Z1_StackBox_Dataset Z1_Dual_Dex1_PourCoffee_Dataset Z1_Dual_Dex1_FoldClothes_Dataset " in
    *" $name "*) EMB=19;; *) EMB=15;;
  esac
  while [ "$(df --output=avail -BG /mnt | tail -1 | tr -dc 0-9)" -lt 30 ]; do
    echo "$name: disk low, waiting"; sleep 120
  done
  echo "=== $name (embodiment $EMB) $(date +%H:%M:%S)"
  rm -rf data/shards_tmp/$name
  python $S/build_shards.py --repo-id unitreerobotics/$name --skip-bad-episodes \
      --out data/shards_tmp/$name --hz 10 --cache-dir .hf_home/tmpcache_$name \
    && python $S/pack_wds.py --src data/shards_tmp/$name \
      --out data/wds/unitree/$name --chunk 50 --maxcount 3000 --embodiment-id $EMB \
    && python $S/validate_wds.py data/wds/unitree/$name --batches 0 \
    || { echo "FAILED $name"; rm -rf data/wds/unitree/$name; }
  rm -rf data/shards_tmp/$name .hf_home/tmpcache_$name
}
export -f one
export S HF_HOME

printf "%s\n" $UNITREE_G1 $UNITREE_Z1 | xargs -P 4 -I{} bash -c 'one {}' 2>&1

python $S/hf_push_wds.py --src data/wds/unitree --repo AlexWortega/wds-unitree
echo "=== UNITREE DONE $(date +%H:%M:%S)"