#!/usr/bin/env bash
# pod_prep_deep.sh — Pod で deep_even10_v1 (位置 1〜5 を各 10 本にそろえた 50 本) の学習の準備をする (2026-11-09 依頼 133).
# 学習は回さない. lab_train.sh --model deep1b (依頼 135) の 4b が呼ぶ (deep1 は依頼 133 で呼んだ形).
#
# 前提: lab_train.sh 3a〜3f で Pod を準備してある (openpi・install_addon.py・add_trim_config.py). 続きの元のモデルは取らない
#       (π0.5 base から. 重みは学習の最初に openpi が gs://openpi-assets から取る)
#       Mac から次の 2 つを送ってある:
#         /workspace/data/mecharm_redblock_box_deep_even10_v1_raw_episodes.tar
#         /workspace/r2r2r/pi05/  (add_deep_config.py と convert_to_lerobot.py を使う)
#
#   bash /workspace/r2r2r/pi05/openpi_addon/pod_prep_deep.sh <steps> <variant>
#   例: bash ... 3000 b
#   <variant> は add_deep_config.py の --variant (依頼 135):
#     v1 = 設定 pi05_mecharm_deep_even10_lora・decay_steps 5_000 (依頼 133 の 2999. lab_train.sh --model deep1)
#     b  = 設定 pi05_mecharm_deep_even10_b_lora・decay_steps を <steps> と同じに (lab_train.sh --model deep1b)
#
# すること (何度流しても同じ結果):
#   1. tar を /workspace/data/<データ名>/ に広げ, エピソードの本数を確かめる (Mac で作ったときは 50 本)
#   2. openpi の config.py に <variant> の設定を足す (--steps --variant)
#   3. LeRobot 形式に変換 (repo_id local/<データ名>)
#   4. 正規化の値をこの 50 本で測り直す (RUNPOD_TRAIN.md 4. の compute_norm_stats.py をこの設定で).
#      置く場所は openpi が学習で読む assets/<設定の名前>/local/<データ名>/norm_stats.json.
#      最後に「NORM sha256=<64 桁> path=<openpi からの相対>」の1行を出す (lab_train.sh が由来に残し, 持ち帰る)
#   5. 学習を始めるコマンドを出す (流さない)
set -euo pipefail

STEPS="${1:?学習のステップ数 (例 3000)}"
VARIANT="${2:?add_deep_config.py の --variant (v1 か b)}"
NAME=mecharm_redblock_box_deep_even10_v1
case "$VARIANT" in          # 何かする前に確かめる (依頼 135)
    v1) CONFIG=pi05_mecharm_deep_even10_lora; EXP=deep_even10_v1 ;;
    b)  CONFIG=pi05_mecharm_deep_even10_b_lora; EXP=deep_even10_v1_b ;;
    *)  echo "[ERROR] <variant> は v1 か b: $VARIANT" >&2; exit 2 ;;
esac
N_MAC=50
TAR=/workspace/data/${NAME}_raw_episodes.tar
DATA=/workspace/data/$NAME
OPENPI=/workspace/openpi
ADDON=/workspace/r2r2r/pi05/openpi_addon
NORM_DST=assets/$CONFIG/local/$NAME/norm_stats.json                           # $OPENPI からの相対

[ -f "$TAR" ] || { echo "[ERROR] tar がない: $TAR" >&2; exit 1; }
# shellcheck disable=SC1091
[ -f "$HOME/.local/bin/env" ] && source "$HOME/.local/bin/env"

echo "== 1. tar を広げる → $DATA"
# 前に広げたものへ上書きしない (pod_prep_trim_corr.sh と同じ. 空の一時フォルダに広げて確かめてから入れ替える)
TMP="$DATA.tmp.$$"
rm -rf "$TMP"
mkdir -p "$TMP"
tar xf "$TAR" -C "$TMP"
N=$(find "$TMP/raw_episodes" -mindepth 2 -maxdepth 2 -name episode.npz | wc -l | tr -d ' ')
[ "$N" = "$N_MAC" ] || { echo "[ERROR] tar のエピソードが $N 本 (Mac で作ったときは $N_MAC 本): $TAR" >&2; rm -rf "$TMP"; exit 1; }
rm -rf "$DATA"
mv "$TMP" "$DATA"
echo "   エピソード $N 本 (Mac で作ったときは $N_MAC 本)"

cd "$OPENPI"
echo "== 2. 学習設定 $CONFIG を足す (steps=${STEPS}・variant ${VARIANT}. π0.5 base から)"
python3 "$ADDON/add_deep_config.py" "$OPENPI" --steps "$STEPS" --variant "$VARIANT"

echo "== 3. LeRobot 形式に変換 (local/$NAME)"
uv run python /workspace/r2r2r/pi05/convert_to_lerobot.py --raw-dir "$DATA/raw_episodes" --repo-id "local/$NAME"

echo "== 4. 正規化の値をこの $N 本で測り直す (compute_norm_stats.py --config-name $CONFIG)"
rm -f "$NORM_DST"           # 前の回の値を残さない (測れなかったときに古い値で学習しない)
uv run scripts/compute_norm_stats.py --config-name "$CONFIG"
[ -s "$NORM_DST" ] || { echo "[ERROR] 正規化の値ができていない: $OPENPI/$NORM_DST" >&2; exit 1; }
SHA=$(sha256sum "$NORM_DST" | cut -d' ' -f1)
echo "   $OPENPI/$NORM_DST (sha256 ${SHA:0:12}…)"
echo "NORM sha256=$SHA path=$NORM_DST"

echo "== 5. 準備できました. 学習はこのコマンドで始める (このスクリプトでは流さない):"
echo "   cd $OPENPI && XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 uv run scripts/train.py $CONFIG --exp-name=$EXP --overwrite 2>&1 | tee /workspace/logs/train_$EXP.log"