"""scripts/train.py を包む。Trainer が「DiT 以外を凍結」した直後に、環境変数 FLEX_FREEZE_REGEX に
合うパラメータを追加で凍結する（F2 = 動画専門家 mot.mixtures.video を凍結し行動側だけ学習）。
upstream は編集しない。学習可能パラメータ数を印字して、凍結が実際に効いたことを目で確かめる。"""
import os, re, sys, runpy
sys.path.insert(0, "/workspace/flexpi_repo/src")
import flexpi.trainer as T

_orig = T.Wan22Trainer.__dict__["_apply_dit_only_train_mode"].__func__

def _patched(model):
    _orig(model)
    pat = os.environ.get("FLEX_FREEZE_REGEX", "")
    frozen = 0
    if pat:
        rx = re.compile(pat)
        for n, p in model.named_parameters():
            if p.requires_grad and rx.search(n):
                p.requires_grad_(False); frozen += p.numel()
    tr = sum(p.numel() for p in model.parameters() if p.requires_grad)
    tot = sum(p.numel() for p in model.parameters())
    heads = {}
    for n, p in model.named_parameters():
        if p.requires_grad:
            k = ".".join(n.split(".")[:3]); heads[k] = heads.get(k, 0) + p.numel()
    top = sorted(heads.items(), key=lambda x: -x[1])[:8]
    print(f"[wrap] 学習可能 {tr/1e9:.3f}B / 全体 {tot/1e9:.3f}B  regex 凍結 {frozen/1e9:.3f}B  pattern={pat!r}", flush=True)
    print("[wrap] 学習可能の内訳（上位）: " + "  ".join(f"{k}={v/1e6:.0f}M" for k, v in top), flush=True)

T.Wan22Trainer._apply_dit_only_train_mode = staticmethod(_patched)
os.chdir("/workspace/flexpi_repo")
sys.argv = ["scripts/train.py"] + sys.argv[1:]
runpy.run_path("scripts/train.py", run_name="__main__")