"""r0b_probe.py — ANIMA R0b: the fork-integration gate (runs INSIDE the pod
clone of diffusion-pipe @ feat/aleph-adapter). Certifies, without a dataset:
  G1 anima weights load through CosmosPredict2Pipeline (dit_config printed);
  G2 aleph relays attach post-materialization (site count, param count);
  G3 DTYPE LAW: relay dtype == block dtype (bf16);
  G4 param groups: aleph_relay bucket present, trunk freezable by lr=0;
  G5 zero-init: all-off/enabled-at-init forward delta == 0 through one block.
Full exp004 training goes to the evening slot per the day-3 plan.
Run: cd /workspace/geolip2/diffusion-pipe && python3 ../pod2/r0b_probe.py
"""
from __future__ import annotations

import sys

import torch

sys.path.insert(0, ".")
import utils.common  # bind the fork utils BEFORE ComfyUI (regular-pkg shadowing)
sys.path.append("submodules/ComfyUI")

CFG = {
    "model": {
        "type": "anima",
        "transformer_path":
            "/workspace/models/anima/split_files/diffusion_models/"
            "anima-base-v1.0.safetensors",
        "llm_path": "/workspace/models/anima/split_files/text_encoders/"
                    "qwen_3_06b_base.safetensors",
        "vae_path": "/workspace/models/anima/split_files/vae/"
                    "qwen_image_vae.safetensors",
        "dtype": torch.bfloat16,
        "aleph_relay": True,
        "aleph_relay_every": 1,
        "aleph_relay_lr": 1e-3,
        "self_attn_lr": 0, "cross_attn_lr": 0, "mlp_lr": 0, "mod_lr": 0,
        "llm_adapter_lr": 0,
    },
    "optimizer": {"lr": 0},
    "reentrant_activation_checkpointing": False,
}


def main():
    import os
    for k, v in (("MASTER_ADDR", "127.0.0.1"), ("MASTER_PORT", "29571"),
                 ("RANK", "0"), ("WORLD_SIZE", "1"), ("LOCAL_RANK", "0")):
        os.environ.setdefault(k, v)
    import deepspeed
    deepspeed.init_distributed()
    from models.cosmos_predict2 import CosmosPredict2Pipeline
    pipe = CosmosPredict2Pipeline(CFG)
    print("[G1] pipeline constructed (text encoder loaded, name="
          f"{pipe.name})", flush=True)
    pipe.load_diffusion_model()
    tr = pipe.transformer
    relays = [getattr(b, "aleph_relay", None) for b in tr.blocks]
    n_sites = sum(r is not None for r in relays)
    n_params = sum(p.numel() for r in relays if r is not None
                   for p in r.parameters())
    print(f"[G2] {n_sites}/{len(tr.blocks)} blocks carry relays, "
          f"{n_params:,} adapter params", flush=True)
    assert n_sites == len(tr.blocks) > 0

    r0 = next(r for r in relays if r is not None)
    b0 = tr.blocks[0]
    bdt = next(p for n, p in b0.named_parameters()
               if "aleph_relay" not in n).dtype
    rdt = next(r0.parameters()).dtype
    print(f"[G3] block dtype {bdt} | relay dtype {rdt}", flush=True)
    assert rdt == bdt == torch.bfloat16, "dtype law violated"

    params = [p for p in tr.parameters()]
    for p, (n, _) in zip(params, tr.named_parameters()):
        pass
    groups = pipe.get_param_groups(
        [p for p in tr.parameters() if hasattr(p, "original_name")])
    n_trainable = sum(p.numel() for g in groups for p in g["params"])
    print(f"[G4] {len(groups)} param groups, trainable {n_trainable:,} "
          f"(should == adapter count {n_params:,})", flush=True)
    assert n_trainable == n_params, "freeze-by-lr-0 leaked trunk params"

    r0.assert_zero_init()
    x = torch.randn(1, 2, 4, 4, tr.model_channels, dtype=bdt)
    with torch.no_grad():
        out_on = r0(x)
        r0.enabled = False
        out_off = r0(x)
        r0.enabled = True
    assert torch.equal(out_on, x) and out_off is x, "zero-init/toggle broken"
    print("[G5] zero-init + toggle exact on the DiT relay", flush=True)
    print("R0B ALL GATES GREEN — fork integration certified; exp004 train "
          "is evening-slot ready", flush=True)


if __name__ == "__main__":
    main()