#!/usr/bin/env python3
"""openpi config.py に trim2 + 補正データ用の学習設定 pi05_mecharm_trim_corr_lora を足す (2026-10-31 依頼 45).

依頼 46: データを corr_v2 (角度の抜けを速度 20 の動きで埋め直したもの) と合わせた
mecharm_redblock_box_teach_v1_trim2_corr_v2 に替えた. corr_v1 の設定が入っている config.py には足さずに止まる
(同じ名前の設定が 2 つになるため. openpi の config.py を git checkout してから入れ直す).
依頼 48: データを corr_v3 (着いたあとの止まりを詰めたもの) と合わせた ..._trim2_corr_v3 に替え,
学習率の予定を続きからの学習用に弱めた (warmup 100 / peak 2.5e-5 / decay_steps 1_500. 下げきる先 decay_lr 5e-6 は
trim2 と同じ). corr_v1 / corr_v2 の設定が入っている config.py には足さずに止まる.

install_addon.py と add_trim_config.py (lab/Log11/pod/) を実行した後に 1 回実行:
  python3 add_trim_corr_config.py /workspace/openpi --steps 4000 --init base
  python3 add_trim_corr_config.py /workspace/openpi --steps 1500 --init /workspace/ckpt/3999/params

pi05_mecharm_trim_lora の設定をそのまま写し, データだけを替える:
  name             pi05_mecharm_trim_lora                     -> pi05_mecharm_trim_corr_lora
  repo_id          local/mecharm_redblock_box_teach_v1_trim2  -> local/mecharm_redblock_box_teach_v1_trim2_corr_v2
  num_train_steps  4_000                                      -> --steps (決めるのはエイイチ. 既定値は置かない)
  lr_schedule      warmup 300 / peak 5e-5 / decay_steps 5_000 -> warmup 100 / peak 2.5e-5 / decay_steps 1_500
                   (decay_lr 5e-6 はそのまま. 依頼 48)
  weight_loader    pi05_base                                  -> --init base ならそのまま (一から)
                                                                 --init <params> なら trim2 の学習済みから続ける
batch 32 などほかは trim2 と同じ. 正規化の値は計算し直さず trim2 の値を置く (pod_prep_trim_corr.sh).
何度実行しても同じ結果 (二重に入れない). 入っている設定と --steps / --init が違うときは止まる.
"""
import argparse
import re
import sys
from pathlib import Path

NAME = "pi05_mecharm_trim_corr_lora"
DATA_TAG = "teach_v1_trim2_corr_v3"
REPO_ID = f"local/mecharm_redblock_box_{DATA_TAG}"
OLD_TAGS = ("teach_v1_trim2_corr_v1", "teach_v1_trim2_corr_v2")   # 前に入れたことのある版 (依頼 45 / 46)
MARK_A = f"    # >>> r2r2r mecharm {DATA_TAG} (train config)\n"
MARK_B = f"    # <<< r2r2r mecharm {DATA_TAG} (train config)\n"
# 学習率の予定 (依頼 48. 3999 からの続きの学習 1,500 ステップ用)
LR_REPS = [
    ("warmup_steps=300,", "warmup_steps=100,"),
    ("peak_lr=5e-5,", "peak_lr=2.5e-5,"),
    ("decay_steps=5_000,", "decay_steps=1_500,"),
]
DECAY_LR = "decay_lr=5e-6,"     # 下げきる先. trim2 と同じ (変えない)
BASE_LOADER = 'weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),'


def make_block(s, steps, init):
    A = "    # >>> r2r2r mecharm teach_v1_trim2 (train config)\n"
    B = "    # <<< r2r2r mecharm teach_v1_trim2 (train config)\n"
    assert s.count(A) == 1 and s.count(B) == 1, "先に install_addon.py と add_trim_config.py を実行する"
    block = s[s.index(A):s.index(B) + len(B)]
    reps = [
        (A, MARK_A),
        (B, MARK_B),
        ('name="pi05_mecharm_trim_lora"', f'name="{NAME}"'),
        ('repo_id="local/mecharm_redblock_box_teach_v1_trim2"', f'repo_id="{REPO_ID}"'),
        ("num_train_steps=4_000,", f"num_train_steps={steps},"),
        *LR_REPS,
    ]
    if init != "base":
        reps.append((BASE_LOADER, f'weight_loader=weight_loaders.CheckpointWeightLoader("{init}"),'))
    for a, b in reps:
        assert block.count(a) == 1, f"見つからない/重複: {a!r}"
        block = block.replace(a, b)
    assert block.count(DECAY_LR) == 1, "trim2 の下げきる先 (decay_lr) が見つからない"
    return block


def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("openpi_root")
    ap.add_argument("--steps", type=int, required=True, help="学習のステップ数 (num_train_steps)")
    ap.add_argument("--init", required=True,
                    help="base = π0.5 base から一から / それ以外 = 続きから学習する params のパス "
                         "(例 /workspace/ckpt/3999/params)")
    args = ap.parse_args(argv)
    if args.steps <= 0:
        sys.exit("[ERROR] --steps は 1 以上")
    if args.init != "base" and not args.init.rstrip("/").endswith("params"):
        sys.exit(f"[ERROR] --init には checkpoint の params フォルダを渡す: {args.init}")
    cfg = Path(args.openpi_root).resolve() / "src" / "openpi" / "training" / "config.py"
    s = cfg.read_text()
    old = [t for t in OLD_TAGS if f"    # >>> r2r2r mecharm {t} (train config)\n" in s]
    if old:
        sys.exit(f"[ERROR] 前の版 ({', '.join(old)}) の {NAME} が入っている. "
                 "openpi の config.py を git checkout してから入れ直す")
    if MARK_A in s:
        have = s[s.index(MARK_A):]
        have = have[:have.index(MARK_B)]
        m = re.search(r"num_train_steps=(\d+),", have)
        same_steps = m and int(m.group(1)) == args.steps
        same_init = (BASE_LOADER in have) if args.init == "base" else (f'("{args.init}")' in have)
        # codex 推奨 (依頼 48): データ名・学習率の予定も入っているものと同じかを見る
        same_rest = (f'repo_id="{REPO_ID}"' in have and all(b in have for _, b in LR_REPS)
                     and DECAY_LR in have)
        if not same_rest:
            sys.exit("[ERROR] 入っている設定のデータ名・学習率の予定が違う. openpi の config.py を git checkout してから入れ直す")
        if not (same_steps and same_init):
            sys.exit("[ERROR] 別の --steps / --init で入っている. openpi の config.py を git checkout してから入れ直す")
        print("already installed (config.py は変更なし)")
        return 0
    block = make_block(s, args.steps, args.init)
    B = "    # <<< r2r2r mecharm teach_v1_trim2 (train config)\n"
    j = s.index(B) + len(B)
    cfg.write_text(s[:j] + block + s[j:])
    print(f"installed: {NAME} (repo_id={REPO_ID}, {args.steps} steps, init={args.init})")
    return 0


if __name__ == "__main__":
    sys.exit(main())