#!/usr/bin/env bash
# pod_prep_trim_corr.sh — Pod で trim2 + corr_v3 の学習の準備をする (2026-10-31 依頼 46 / 48). 学習は回さない.
#
# 前提: lab_start.sh で Pod を準備してある (openpi・install_addon.py・add_trim_config.py・/workspace/ckpt/3999)
#       Mac から次の 2 つを送ってある:
#         /workspace/data/mecharm_redblock_box_teach_v1_trim2_corr_v3_raw_episodes.tar
#         /workspace/r2r2r/pi05/  (この依頼の版. add_trim_corr_config.py と convert_to_lerobot.py を使う)
#
#   bash /workspace/r2r2r/pi05/openpi_addon/pod_prep_trim_corr.sh <steps> <init> [v3|v4|v5]
#   例: bash ... 1500 /workspace/ckpt/3999/params
#   依頼 57: 3 つ目に v4 を付けると trim2 + corr_v4 (..._trim2_corr_v4・設定 pi05_mecharm_trim_corr_v4_lora) の準備.
#     例: bash ... 1500 /workspace/ckpt/corr_v1_1499/params v4
#     続きの元は corr_v1_1499 (lab_train.sh --model corr2 の 3c が取る). 3999 は取らないので, 正規化の値は
#     1499 の assets にある値 (3999 の trim2 の値と sha256 が同じ. 依頼 52 で確かめた) を, 同じ sha256 で確かめて置く
#   依頼 58: v4 のデータは ..._trim2_corr_v4b (8 本の corr_v4 から 182545 を外した 7 本. 合わせて 26 本). 設定の名前は同じ
#   依頼 97: v5 は trim2 + corr_v5 (..._trim2_corr_v5・設定 pi05_mecharm_trim_corr_v5_lora. 合わせて 30 本).
#     続きの元は corr_v2_1499 (lab_train.sh --model corr3 の 3c が取る). 正規化の値は 1499 の assets にある値
#     (trim2 の値と sha256 が同じ. 依頼 97 で KIOXIA の写しを確かめた) を, 同じ sha256 で確かめて置く
#     例: bash ... 1500 /workspace/ckpt/corr_v2_1499/params v5
#   v3 (既定) は今までと同じ
#
# すること (何度流しても同じ結果):
#   1. tar を /workspace/data/<データ名>/ に広げ, エピソードの本数を出す
#   2. openpi の config.py に pi05_mecharm_trim_corr_lora を足す (--steps / --init)
#   3. LeRobot 形式に変換 (repo_id local/<データ名>)
#   4. 正規化の値は計算し直さず, 3999 の学習で使った trim2 の値をそのまま置く (依頼 48)
#      元: 3999 のモデルの assets (lab_start.sh 3c の hf download が置く)
#        /workspace/ckpt/3999/assets/local/mecharm_redblock_box_teach_v1_trim2/norm_stats.json
#      先: openpi が学習で読む場所 (assets/<設定の名前>/<repo_id>/)
#        assets/pi05_mecharm_trim_corr_lora/local/<データ名>/norm_stats.json
#      Mac の写し (/Volumes/KIOXIA/r2r2r-data/checkpoints/pi05_mecharm_trim_lora_v1/3999/assets/...) と
#      sha256 が同じかを確かめる
#   5. 学習を始めるコマンドを出す (流さない)
set -euo pipefail

STEPS="${1:?学習のステップ数 (例 1500)}"
INIT="${2:?続きから学習する params (例 /workspace/ckpt/3999/params) / 一からなら base}"
VARIANT="${3:-v3}"
case "$VARIANT" in
    v3) NAME=mecharm_redblock_box_teach_v1_trim2_corr_v3
        CONFIG=pi05_mecharm_trim_corr_lora
        NORM_SRC=/workspace/ckpt/3999/assets/local/mecharm_redblock_box_teach_v1_trim2/norm_stats.json
        EXP=trim_corr_v3; N_MAC=24 ;;
    v4) NAME=mecharm_redblock_box_teach_v1_trim2_corr_v4b
        CONFIG=pi05_mecharm_trim_corr_v4_lora
        NORM_SRC=/workspace/ckpt/corr_v1_1499/assets/local/mecharm_redblock_box_teach_v1_trim2_corr_v3/norm_stats.json
        EXP=trim_corr_v4; N_MAC=26 ;;
    v5) NAME=mecharm_redblock_box_teach_v1_trim2_corr_v5
        CONFIG=pi05_mecharm_trim_corr_v5_lora
        NORM_SRC=/workspace/ckpt/corr_v2_1499/assets/local/mecharm_redblock_box_teach_v1_trim2_corr_v4b/norm_stats.json
        EXP=trim_corr_v5; N_MAC=30 ;;
    *) echo "[ERROR] 3 つ目は v3 か v4 か v5: $VARIANT" >&2; exit 2 ;;
esac
TAR=/workspace/data/${NAME}_raw_episodes.tar
DATA=/workspace/data/$NAME
OPENPI=/workspace/openpi
ADDON=/workspace/r2r2r/pi05/openpi_addon
NORM_SHA256=80346837eb12ffa6e683416f0fa41a196519fa8459d1f0b666b06cd621c7f61b   # Mac の写しの sha256 (v3・v4・v5 とも同じ trim2 の値)
NORM_DST=assets/$CONFIG/local/$NAME/norm_stats.json                           # $OPENPI からの相対

[ -f "$TAR" ] || { echo "[ERROR] tar がない: $TAR" >&2; exit 1; }
[ -f "$NORM_SRC" ] || { echo "[ERROR] trim2 の正規化の値がない: $NORM_SRC (3c で続きの元のモデルを取る)" >&2; exit 1; }
sha_of() { sha256sum "$1" | cut -d' ' -f1; }
[ "$(sha_of "$NORM_SRC")" = "$NORM_SHA256" ] \
    || { echo "[ERROR] $NORM_SRC が 3999 の学習の trim2 の値 (sha256 $NORM_SHA256) と違う" >&2; exit 1; }
if [ "$INIT" != base ] && [ ! -d "$INIT" ]; then
    echo "[ERROR] 続きから学習する params がない: $INIT" >&2; exit 1
fi
# shellcheck disable=SC1091
[ -f "$HOME/.local/bin/env" ] && source "$HOME/.local/bin/env"

echo "== 1. tar を広げる → $DATA"
# codex 必須: 前に広げたものへ上書きすると, tar を作り直して減ったエピソードやコマが残る.
# 空の一時フォルダに広げて確かめてから入れ替える (前のものは消す. 元は tar にある)
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" -gt 0 ] || { echo "[ERROR] tar にエピソードがない: $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 init=$INIT)"
python3 "$ADDON/add_trim_corr_config.py" "$OPENPI" --variant "$VARIANT" --steps "$STEPS" --init "$INIT"

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. 正規化の値は trim2 の値をそのまま置く (計算し直さない)"
mkdir -p "$(dirname "$NORM_DST")"
cp "$NORM_SRC" "$NORM_DST"
[ "$(sha_of "$NORM_DST")" = "$NORM_SHA256" ] || { echo "[ERROR] 置いた値が違う: $NORM_DST" >&2; exit 1; }
echo "   $NORM_SRC → $OPENPI/$NORM_DST (sha256 ${NORM_SHA256:0:12}…)"

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"
