#!/usr/bin/env bash
set -euo pipefail

# Self-contained bundle version (2026-09-19): code = ../myGR00T (this bundle), outputs = ../runs.
# Only the conda env, the LeRobot dataset and the base checkpoint are referenced by absolute path.
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
BUNDLE="$(dirname "$HERE")"

# Set this to the dataset directory name under TRAIN_DATASET_BASE.
DATASETNAME="${DATASETNAME:-PickPlaceCounterToCabinet_150eps_augmented_groot15_v3_100_clean}"

MYGROOT_ROOT="${MYGROOT_ROOT:-$BUNDLE/myGR00T}"
PYTHON_BIN="${PYTHON_BIN:-/data/nvidia/gripper_augmentator/conda-envs/mygr00t/bin/python}"
CONDA_BIN="${CONDA_BIN:-/data/nvidia/gripper_augmentator/conda-envs/mygr00t/bin}"
TRAIN_DATASET_BASE="${TRAIN_DATASET_BASE:-/lp-dev/jonghoon/robocasa_full/gr00t_views}"
OUTPUT_BASE="${OUTPUT_BASE:-$BUNDLE/runs}"
BASE_CKPT="${BASE_CKPT:-/lp-dev/jonghoon/myGR00T_outputs/pnpcountertocab_all502_gbs64_wandb_60k_save5k_20260427_193901/checkpoint-60000}"

TASK_SLUG="${TASK_SLUG:-pnpcountertocab}"
RUN_ID="${RUN_ID:-$(date -u +%Y%m%d_%H%M%S)}"
MAX_STEPS="${MAX_STEPS:-30000}"
SAVE_STEPS="${SAVE_STEPS:-5000}"
RESUME="${RESUME:-0}"
GPUS="${GPUS:-0,1,2,3}"
PER_GPU_BATCH="${PER_GPU_BATCH:-16}"
DATALOADER_NUM_WORKERS="${DATALOADER_NUM_WORKERS:-8}"
DATA_CONFIG="${DATA_CONFIG:-single_panda_gripper}"
EMBODIMENT_TAG="${EMBODIMENT_TAG:-new_embodiment}"
BACKBONE_MODEL_TYPE="${BACKBONE_MODEL_TYPE:-eagle}"
BACKBONE_SELECT_LAYER="${BACKBONE_SELECT_LAYER:-12}"
REPORT_TO="${REPORT_TO:-tensorboard}"   # wandb needs WANDB_API_KEY (or WANDB_API_KEY_FILE)

DATASET_PATH="${DATASET_PATH:-${TRAIN_DATASET_BASE}/${DATASETNAME}}"
RUN_NAME="${RUN_NAME:-${TASK_SLUG}_${DATASETNAME}_newemb_from60k_gbs64_${MAX_STEPS}steps_${RUN_ID}}"
OUTPUT_DIR="${OUTPUT_DIR:-${OUTPUT_BASE}/${RUN_NAME}}"

export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-$GPUS}"
export PATH="${CONDA_BIN}:${PATH}"
export WANDB_PROJECT="${WANDB_PROJECT:-gr00t-robocasa}"
export WANDB_ENTITY="${WANDB_ENTITY:-alphabet1}"
export WANDB_DIR="${WANDB_DIR:-$BUNDLE/logs/wandb}"
export WANDB_CACHE_DIR="${WANDB_CACHE_DIR:-$BUNDLE/logs/wandb/cache}"
export WANDB_CONFIG_DIR="${WANDB_CONFIG_DIR:-$BUNDLE/logs/wandb/config}"
export NO_ALBUMENTATIONS_UPDATE=1
# The mygr00t env's editable install of `gr00t` points at a deleted checkout; import the bundle copy instead.
export PYTHONPATH="$MYGROOT_ROOT${PYTHONPATH:+:$PYTHONPATH}"
export HF_HOME="${HF_HOME:-/lp-dev/jonghoon/.cache/huggingface}"
export HF_HUB_OFFLINE="${HF_HUB_OFFLINE:-1}"
export TOKENIZERS_PARALLELISM=false

# Optional, safer than hard-coding a secret in this script.
WANDB_API_KEY_FILE="${WANDB_API_KEY_FILE:-}"
if [[ -z "${WANDB_API_KEY:-}" && -n "$WANDB_API_KEY_FILE" && -f "$WANDB_API_KEY_FILE" ]]; then
    WANDB_API_KEY="$(tr -d '[:space:]' < "$WANDB_API_KEY_FILE")"
    export WANDB_API_KEY
fi

IFS=',' read -r -a GPU_LIST <<< "$CUDA_VISIBLE_DEVICES"
NUM_GPUS="${NUM_GPUS:-${#GPU_LIST[@]}}"

if [[ ! -d "$DATASET_PATH" ]]; then
    echo "FATAL: dataset path does not exist: $DATASET_PATH" >&2
    exit 1
fi
if [[ ! -f "$BASE_CKPT/model.safetensors.index.json" ]]; then
    echo "FATAL: checkpoint is missing or incomplete: $BASE_CKPT" >&2
    exit 1
fi

mkdir -p "$OUTPUT_DIR" "$WANDB_DIR" "$WANDB_CACHE_DIR" "$WANDB_CONFIG_DIR"
cd "$MYGROOT_ROOT"

exec > >(tee -a "$OUTPUT_DIR/train.log") 2>&1

echo "run_name=$RUN_NAME"
echo "output_dir=$OUTPUT_DIR"
echo "dataset_name=$DATASETNAME"
echo "dataset_path=$DATASET_PATH"
echo "base_ckpt=$BASE_CKPT"
echo "cuda_visible_devices=$CUDA_VISIBLE_DEVICES"
echo "num_gpus=$NUM_GPUS"
echo "per_gpu_batch=$PER_GPU_BATCH"
echo "global_batch=$((NUM_GPUS * PER_GPU_BATCH))"
echo "max_steps=$MAX_STEPS"
echo "save_steps=$SAVE_STEPS"
echo "resume=$RESUME"
date -u

RESUME_ARGS=()
if [[ "$RESUME" == "1" ]]; then
    RESUME_ARGS+=(--resume)
fi

SEED_ARGS=()
if [[ -n "${DATA_SEED:-}" ]]; then
    SEED_ARGS+=(--data-seed "$DATA_SEED")
fi

exec "$PYTHON_BIN" "$MYGROOT_ROOT/scripts/gr00t_finetune.py" \
    --run-name "$RUN_NAME" \
    --dataset-path "$DATASET_PATH" \
    --output-dir "$OUTPUT_DIR" \
    --data-config "$DATA_CONFIG" \
    --dataloader-num-workers "$DATALOADER_NUM_WORKERS" \
    --embodiment-tag "$EMBODIMENT_TAG" \
    --backbone-model-type "$BACKBONE_MODEL_TYPE" \
    --backbone-select-layer "$BACKBONE_SELECT_LAYER" \
    --base-model-path "$BASE_CKPT" \
    --num-gpus "$NUM_GPUS" \
    --batch-size "$PER_GPU_BATCH" \
    --max-steps "$MAX_STEPS" \
    --save-steps "$SAVE_STEPS" \
    --report-to "$REPORT_TO" \
    "${SEED_ARGS[@]}" \
    "${RESUME_ARGS[@]}" \
    "$@"