#!/usr/bin/env bash
# Yukon launcher: pred_size=128 for finer voxel quantization (~4 image px
# per bin vs puget's 8 px at pred_size=64). Requires H100-class memory.
# bs=2 is a conservative starting point given the 8x voxel-count bump.
set -euo pipefail
if [[ $# -lt 3 ]]; then
  echo "usage: $0 <name> <data_dir> <iters> [extra args...]" >&2
  exit 2
fi
NAME="$1"; DATA_DIR="$2"; ITERS="$3"; shift 3

REPO_DIR="/home/robot-lab/smithbot_v4"
PYTHON="/home/robot-lab/cameron/raiden_fork.venv/bin/python"
export DINO_REPO_DIR="/home/robot-lab/cameron/dinov3_local"
export DINO_WEIGHTS_PATH="/home/robot-lab/cameron/dinov3_local/weights/dinov3_vits16plus_pretrain_lvd1689m-4057cbaa.pth"
RUN_DIR="/home/robot-lab/smithbot/runs/${NAME}"
CKPT_DIR="/home/robot-lab/smithbot/checkpoints/${NAME}"
LOG="/home/robot-lab/smithbot/logs/exp_${NAME}.log"
mkdir -p "$RUN_DIR" "$CKPT_DIR" "$(dirname "$LOG")"
export WANDB_DIR="$RUN_DIR"

echo "── smithbot_v4 YUKON train ${NAME} on $(hostname -s) (CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0}) ──"
echo "  data:    $DATA_DIR"
echo "  iters:   $ITERS  ·  pred_size=128 (4 px voxel bin)"
echo "  log:     $LOG"

cd "$RUN_DIR"
exec "$PYTHON" "$REPO_DIR/train.py" \
  --name "$NAME" \
  --data_root "$DATA_DIR" \
  --max_iters "$ITERS" \
  --v4 \
  --multiview scene,wrist \
  --batch_size 2 \
  --lr 3e-4 --lr_schedule cosine --lr_warmup 500 \
  --backbone_lr_mult 0.1 \
  --n_window 32 --n_height_bins 128 --n_gripper_bins 32 --n_rot_clusters 256 \
  --pred_size 128 \
  --feat_dim 32 --d_model 256 --mlp_layers 1 --mlp_hidden_mult 2.0 \
  --height_enc_dim 32 --voxel_mlp_hidden 64 \
  --cls_fusion concat --per_voxel_grip_rot \
  --temporal_head flat_mlp \
  --mv_past_n 8 --mv_past_enc_dim 64 \
  --cam_margin_m 0.01 \
  --img_size 504 --frame_stride 1 \
  --grip_lo -2.0 --grip_hi 0.5 \
  --ckpt_every_iters 200 --vis_every_iters 200 \
  "$@"