#!/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>
#   例: bash ... 1500 /workspace/ckpt/3999/params
#
# すること (何度流しても同じ結果):
#   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}"
NAME=mecharm_redblock_box_teach_v1_trim2_corr_v3
CONFIG=pi05_mecharm_trim_corr_lora
TAR=/workspace/data/${NAME}_raw_episodes.tar
DATA=/workspace/data/$NAME
OPENPI=/workspace/openpi
ADDON=/workspace/r2r2r/pi05/openpi_addon
NORM_SRC=/workspace/ckpt/3999/assets/local/mecharm_redblock_box_teach_v1_trim2/norm_stats.json
NORM_SHA256=80346837eb12ffa6e683416f0fa41a196519fa8459d1f0b666b06cd621c7f61b   # Mac の写しの sha256
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 (lab_start.sh の 3c で 3999 を取る)" >&2; exit 1; }
sha_of() { sha256sum "$1" | cut -d' ' -f1; }
[ "$(sha_of "$NORM_SRC")" = "$NORM_SHA256" ] \
    || { echo "[ERROR] $NORM_SRC が 3999 の学習の値 (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 で作ったときは 24 本)"

cd "$OPENPI"
echo "== 2. 学習設定 $CONFIG を足す (steps=$STEPS init=$INIT)"
python3 "$ADDON/add_trim_corr_config.py" "$OPENPI" --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=trim_corr_v3 --overwrite 2>&1 | tee /workspace/logs/train_trim_corr_v3.log"