#!/usr/bin/env python3
"""openpi config.py に deep_even10_v1 (位置 1〜5 を各 10 本にそろえた 50 本) 用の学習設定を足す
(2026-11-09 依頼 133. lab_train.sh --model deep1. 依頼 135 で --variant b (lab_train.sh --model deep1b) を足した).

π0.5 base から新しく学習する (続きの学習ではない). 中身は trim の pi05_mecharm_trim_lora と同じ LoRA・batch 32・
学習率の予定 (warmup 300 / peak 5e-5 / decay_lr 5e-6)・weight_loader (π0.5 base) で, 次の所だけ変える
(add_trim_config.py と同じ方式):
  name             pi05_mecharm_trim_lora                     -> --variant v1: pi05_mecharm_deep_even10_lora
                                                                 --variant b:  pi05_mecharm_deep_even10_b_lora
  repo_id          local/mecharm_redblock_box_teach_v1_trim2  -> local/mecharm_redblock_box_deep_even10_v1
  num_train_steps  4_000                                      -> --steps (依頼 133 は 3000)
  decay_steps      5_000                                      -> --variant v1: 5_000 のまま (依頼 133 の 2999 はこれで学習した)
                                                                 --variant b:  --steps と同じ (依頼 135. 学習の終わりで decay_lr まで下げきる)
正規化の値は写さず, Pod で compute_norm_stats.py をこの設定で回して測り直す (pod_prep_deep.sh).

install_addon.py と add_trim_config.py (lab/Log11/pod/) を実行した後に 1 回実行:
  python3 add_deep_config.py /workspace/openpi --steps 3000 --variant b
何度実行しても同じ結果 (二重に入れない). 入っている設定と --steps が違うときは止まる. v1 と b は名前が違うので両方入れられる.
"""
import argparse
import re
import sys
from pathlib import Path

REPO_ID = "local/mecharm_redblock_box_deep_even10_v1"
TRIM_A = "    # >>> r2r2r mecharm teach_v1_trim2 (train config)\n"
TRIM_B = "    # <<< r2r2r mecharm teach_v1_trim2 (train config)\n"
BASE_LOADER = 'weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),'
TRIM_DECAY = "decay_steps=5_000,"
WARMUP = 300                    # trim と同じ warmup_steps. decay_steps はこれより大きくないと下げる区間がない
# trim と同じであることを確かめる値 (写したあとも変わっていないこと). decay_steps は --variant ごとに見る
SAME_AS_TRIM = ("batch_size=32,", f"warmup_steps={WARMUP},", "peak_lr=5e-5,", "decay_lr=5e-6,", BASE_LOADER)
# --variant ごとの設定の名前・印・decay_steps (None は --steps と同じ)
VARIANTS = {
    "v1": {"name": "pi05_mecharm_deep_even10_lora", "mark": "deep_even10_v1", "decay": 5000},     # 依頼 133 (今回の 2999)
    "b": {"name": "pi05_mecharm_deep_even10_b_lora", "mark": "deep_even10_v1 b", "decay": None},   # 依頼 135
}
# 依頼 133 の形を残す (テスト・ほかから読む名前)
NAME = VARIANTS["v1"]["name"]


def marks(variant):
    m = VARIANTS[variant]["mark"]
    return (f"    # >>> r2r2r mecharm {m} (train config)\n", f"    # <<< r2r2r mecharm {m} (train config)\n")


MARK_A, MARK_B = marks("v1")


def decay_of(variant, steps):
    d = VARIANTS[variant]["decay"]
    return steps if d is None else d


def decay_text(n):
    """openpi の書き方に合わせる (5000 -> 5_000)"""
    return f"decay_steps={n:_},"


def make_block(s, steps, variant="v1"):
    assert s.count(TRIM_A) == 1 and s.count(TRIM_B) == 1, "先に install_addon.py と add_trim_config.py を実行する"
    block = s[s.index(TRIM_A):s.index(TRIM_B) + len(TRIM_B)]
    ma, mb = marks(variant)
    reps = [
        (TRIM_A, ma),
        (TRIM_B, mb),
        ('name="pi05_mecharm_trim_lora"', f'name="{VARIANTS[variant]["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},"),
        (TRIM_DECAY, decay_text(decay_of(variant, steps))),
    ]
    for a, b in reps:
        assert block.count(a) == 1, f"見つからない/重複: {a!r}"
        block = block.replace(a, b)
    for x in SAME_AS_TRIM:
        assert block.count(x) == 1, f"trim の設定に見つからない: {x!r}"
    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. 依頼 133 は 3000)")
    ap.add_argument("--variant", choices=sorted(VARIANTS), required=True,
                    help="v1 = decay_steps 5_000 (依頼 133 の 2999) / b = decay_steps を --steps と同じに (依頼 135)")
    args = ap.parse_args(argv)
    if args.steps <= 0:
        sys.exit("[ERROR] --steps は 1 以上")
    decay = decay_of(args.variant, args.steps)
    if decay <= WARMUP:
        sys.exit(f"[ERROR] decay_steps ({decay}) が warmup_steps ({WARMUP}) 以下. --steps を大きくする")
    name = VARIANTS[args.variant]["name"]
    ma, mb = marks(args.variant)
    cfg = Path(args.openpi_root).resolve() / "src" / "openpi" / "training" / "config.py"
    s = cfg.read_text()
    if ma in s:
        have = s[s.index(ma):]
        have = have[:have.index(mb)]
        m = re.search(r"num_train_steps=(\d+),", have)
        if not (f'name="{name}"' in have and f'repo_id="{REPO_ID}"' in have and all(x in have for x in SAME_AS_TRIM)):
            sys.exit("[ERROR] 入っている設定の名前・データ名・学習の値が違う. openpi の config.py を git checkout してから入れ直す")
        if not (m and int(m.group(1)) == args.steps):
            sys.exit("[ERROR] 別の --steps で入っている. openpi の config.py を git checkout してから入れ直す")
        if have.count(decay_text(decay)) != 1:
            sys.exit(f"[ERROR] 入っている設定の decay_steps が {decay} ではない. 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.variant)
    j = s.index(TRIM_B) + len(TRIM_B)
    cfg.write_text(s[:j] + block + s[j:])
    print(f"installed: {name} (repo_id={REPO_ID}, {args.steps} steps, decay_steps {decay}, init=π0.5 base)")
    return 0


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