#!/usr/bin/env bash
# Verifies that GR00T 1.5 fine-tuning runs from THIS bundle: imports, dataset load,
# then a REAL 20-step single-GPU fine-tune on the 150-episode augmented dataset
# (per-GPU batch 2, no LoRA, same code path as the full run) writing a checkpoint
# into runs/smoke_<timestamp>/.  Needs one GPU with ~25 GB free.
#   ./smoke_test.sh <gpu_index>          e.g. ./smoke_test.sh 3
set -euo pipefail
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"; B="$(dirname "$HERE")"
GPU="${1:?usage: smoke_test.sh <gpu_index>}"
PY=/data/nvidia/gripper_augmentator/conda-envs/mygr00t/bin/python
DS="${DATASET_PATH:-/lp-dev/jonghoon/robocasa_full/gr00t_views/PickPlaceCounterToCabinet_150eps_augmented_groot15_v3_100_clean}"
CK="${BASE_CKPT:-/lp-dev/jonghoon/myGR00T_outputs/pnpcountertocab_all502_gbs64_wandb_60k_save5k_20260427_193901/checkpoint-60000}"
pass(){ echo "  PASS  $*"; }; fail(){ echo "  FAIL  $*"; exit 1; }
export PYTHONPATH="$B/myGR00T" HF_HOME=/lp-dev/jonghoon/.cache/huggingface HF_HUB_OFFLINE=1 NO_ALBUMENTATIONS_UPDATE=1

echo "== static =="
[ -x "$PY" ] && pass "env python $PY" || fail "mygr00t env missing"
(cd "$B/myGR00T" && "$PY" -c "import gr00t, os; assert gr00t.__file__.startswith('$B/'), gr00t.__file__; from gr00t.experiment.data_config import DATA_CONFIG_MAP; assert 'single_panda_gripper' in DATA_CONFIG_MAP") && pass "gr00t imports from the bundle, data config single_panda_gripper registered" || fail "gr00t import"
[ -f "$DS/meta/modality.json" ] && pass "dataset $DS" || fail "dataset missing"
[ -f "$CK/model.safetensors.index.json" ] && pass "base checkpoint $CK" || fail "checkpoint missing"
(cd "$B/myGR00T" && "$PY" - <<PYEOF
from gr00t.data.dataset import LeRobotSingleDataset
from gr00t.data.schema import EmbodimentTag
from gr00t.experiment.data_config import DATA_CONFIG_MAP
dc = DATA_CONFIG_MAP["single_panda_gripper"]
ds = LeRobotSingleDataset("$DS", dc.modality_config(), transforms=dc.transform(), embodiment_tag=EmbodimentTag("new_embodiment"), video_backend="decord")
s = ds[0]; print({k: tuple(v.shape) for k, v in s.items() if hasattr(v, "shape")}); assert len(ds) > 1000
PYEOF
) && pass "LeRobotSingleDataset loads a sample (videos decoded)" || fail "dataset load"

echo "== 20-step fine-tune on GPU $GPU =="
RUN=smoke_$(date -u +%Y%m%d_%H%M%S); OUT="$B/runs/$RUN"
CUDA_VISIBLE_DEVICES="$GPU" GPUS="$GPU" NUM_GPUS=1 PER_GPU_BATCH=2 MAX_STEPS=20 SAVE_STEPS=20 DATALOADER_NUM_WORKERS=4 \
  REPORT_TO=tensorboard RUN_NAME="$RUN" OUTPUT_DIR="$OUT" DATASET_PATH="$DS" BASE_CKPT="$CK" \
  bash "$HERE/train_groot15_single_dataset.sh" > "$B/logs/$RUN.log" 2>&1 || { tail -30 "$B/logs/$RUN.log"; fail "fine-tune exited non-zero (log: $B/logs/$RUN.log)"; }
[ -f "$OUT/checkpoint-20/model.safetensors.index.json" ] || [ -f "$OUT/checkpoint-20/model.safetensors" ] && pass "checkpoint-20 written to $OUT" || fail "no checkpoint-20 in $OUT"
grep -q "train_loss\|'loss'" "$B/logs/$RUN.log" && pass "loss logged: $(grep -o "'loss': [0-9.]*" "$B/logs/$RUN.log" | tail -1)" || pass "run finished (no loss line found in log, check $B/logs/$RUN.log)"
echo "ALL PASSED  (log: $B/logs/$RUN.log)"