#!/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 が違うときは止まる.

依頼 57: --variant v4 で, corr_v4 (今日の手首の外れのなかった 3 本を足したもの) と合わせた ..._trim2_corr_v4 の設定
pi05_mecharm_trim_corr_v4_lora を足す (名前が違うので v3 の設定と並べて入る). 学習率の予定・ほかの値は v3 と同じ.
続きの元 (--init) は corr_v1_1499 (/workspace/ckpt/corr_v1_1499/params). 既定は --variant v3 (今までと同じ)
  python3 add_trim_corr_config.py /workspace/openpi --variant v4 --steps 1500 --init /workspace/ckpt/corr_v1_1499/params
依頼 58: v4 のデータを ..._trim2_corr_v4b (corr_v4 から 10-31 の 182545 を外した 7 本) に替えた. 設定の名前は
pi05_mecharm_trim_corr_v4_lora のまま. 8 本の ..._trim2_corr_v4 の設定が入っている config.py には足さずに止まる
依頼 97: --variant v5 で, corr_v5 (corr_v4b の 7 本 + 11-06 のゲームパッドで直した 4 本) と合わせた ..._trim2_corr_v5 の設定
pi05_mecharm_trim_corr_v5_lora を足す (v3・v4 と並べて入る). 学習率の予定・ほかの値は v3 と同じ.
続きの元 (--init) は corr_v2_1499 (/workspace/ckpt/corr_v2_1499/params)
  python3 add_trim_corr_config.py /workspace/openpi --variant v5 --steps 1500 --init /workspace/ckpt/corr_v2_1499/params
"""
import argparse
import re
import sys
from pathlib import Path

# 版ごとの 設定の名前・データ. old_tags = 入っていたら止まる前の版 (同じ設定の名前を使っていたもの)
VARIANTS = {
    "v3": {"name": "pi05_mecharm_trim_corr_lora", "data_tag": "teach_v1_trim2_corr_v3",
           "old_tags": ("teach_v1_trim2_corr_v1", "teach_v1_trim2_corr_v2")},   # 依頼 45 / 46 で入れた版
    "v4": {"name": "pi05_mecharm_trim_corr_v4_lora", "data_tag": "teach_v1_trim2_corr_v4b",
           "old_tags": ("teach_v1_trim2_corr_v4",)},   # 依頼 57 / 58 (8 本の corr_v4 → 182545 を外した 7 本の corr_v4b)
    "v5": {"name": "pi05_mecharm_trim_corr_v5_lora", "data_tag": "teach_v1_trim2_corr_v5", "old_tags": ()},  # 依頼 97
}


def variant(v):
    """版 v の (NAME, REPO_ID, OLD_TAGS, MARK_A, MARK_B)."""
    d = VARIANTS[v]
    tag = d["data_tag"]
    return (d["name"], f"local/mecharm_redblock_box_{tag}", d["old_tags"],
            f"    # >>> r2r2r mecharm {tag} (train config)\n", f"    # <<< r2r2r mecharm {tag} (train config)\n")


NAME, REPO_ID, OLD_TAGS, MARK_A, MARK_B = variant("v3")      # 既定の版 (今までと同じ)
# 学習率の予定 (依頼 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, v="v3"):
    name, repo_id, _, mark_a, mark_b = variant(v)
    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("--variant", choices=sorted(VARIANTS), default="v3",
                    help="v3 = pi05_mecharm_trim_corr_lora (..._trim2_corr_v3. 既定) / "
                         "v4 = pi05_mecharm_trim_corr_v4_lora (..._trim2_corr_v4b. 依頼 57 / 58)")
    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"
    name, repo_id, old_tags, mark_a, mark_b = variant(args.variant)
    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
    if f'name="{name}"' in s:
        sys.exit(f"[ERROR] {name} がもう別の形で入っている. openpi の config.py を git checkout してから入れ直す")
    block = make_block(s, args.steps, args.init, args.variant)
    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())
