"""xattn 臂（Matrix-Game 3 动作模块）的自测：baseline/SPEC.md 的测试 1–4（一张 H200，DiT bf16 ~10 GB）。

    CUDA_VISIBLE_DEVICES=5 PYTHONPATH=/opt/dlami/nvme/zhiyangdeng/ActionRoPE DIFFSYNTH_SKIP_DOWNLOAD=True \
        /opt/dlami/nvme/zhiyangdeng/ActionRoPE/.venv/bin/python -m pytest \
        /opt/dlami/nvme/zhiyangdeng/ActionRoPE/tests/test_arm_xattn.py -s -v

每个用例都能单独跑（-k），模型只在 module 级 fixture 里加载一次；每个用例自己建一个新的 arm。
实测数字追加写到 $AROPE_TEST_NUMBERS（默认 outputs/test_artifacts/arm_xattn_numbers.json）。
"""
from __future__ import annotations

import json
import os
import time

import pytest
import torch

os.environ.setdefault("DIFFSYNTH_SKIP_DOWNLOAD", "True")

ROOT = "/opt/dlami/nvme/zhiyangdeng/ActionRoPE"
MODEL_DIR = f"{ROOT}/models/Wan2.2-TI2V-5B"
DIT_FILES = [f"{MODEL_DIR}/diffusion_pytorch_model-0000{i}-of-00003.safetensors" for i in (1, 2, 3)]
TEST_DIR = os.environ.get("AROPE_TEST_DIR", os.path.join(ROOT, "outputs", "test_artifacts"))
NUMBERS_PATH = os.environ.get("AROPE_TEST_NUMBERS", os.path.join(TEST_DIR, "arm_xattn_numbers.json"))

C, F, LAT_H, LAT_W = 48, 21, 30, 52
TOK_H, TOK_W = 15, 26
L_TXT, D_TXT = 512, 4096
N_BLOCKS = 15  # 上游 blocks=[0..14]

# 解析式参数量（见 baseline/xattn.py 顶部 docstring），dim=3072
DIM = 3072
KB_PER_BLOCK = lambda win: (4 * 128 + 128 + 128 * 128 + 128) + DIM * 1024 + (128 * win) * 2048 + 1024 * DIM  # noqa: E731
MOUSE_PER_BLOCK = lambda win: ((2 * win + DIM) * 1024 + 1024) + (1024 * 1024 + 1024) + 2 * 1024 + 1024 * 3072 + 1024 * DIM  # noqa: E731


def record(**kv):
    data = {}
    try:
        with open(NUMBERS_PATH) as fp:
            data = json.load(fp)
    except (FileNotFoundError, json.JSONDecodeError):
        pass
    data.update({k: (float(v) if isinstance(v, (int, float)) and not isinstance(v, bool) else v) for k, v in kv.items()})
    os.makedirs(os.path.dirname(NUMBERS_PATH), exist_ok=True)
    with open(NUMBERS_PATH, "w") as fp:
        json.dump(data, fp, indent=2, ensure_ascii=False)
    print("\n[numbers]", json.dumps(kv, ensure_ascii=False))


@pytest.fixture(scope="module")
def dit():
    from diffsynth.pipelines.wan_video import ModelConfig, WanVideoPipeline
    from actionrope.arope import install_arope

    pipe = WanVideoPipeline.from_pretrained(
        torch_dtype=torch.bfloat16, device="cuda",
        model_configs=[ModelConfig(path=DIT_FILES)],
        tokenizer_config=None, redirect_common_files=False,
    )
    model = pipe.dit
    model.eval().requires_grad_(False)
    install_arope(model, mask_channel=False)  # 新臂的接入方式：普通 RoPE、无 mask 通道
    model._baseline_keys = set(model.state_dict().keys())
    return model


def make_inputs(seed=0, batch=1, device="cuda"):
    g = torch.Generator(device="cpu").manual_seed(seed)
    latents = torch.randn(batch, C, F, LAT_H, LAT_W, generator=g).to(device=device, dtype=torch.bfloat16)
    context = torch.randn(batch, L_TXT, D_TXT, generator=g).to(device=device, dtype=torch.bfloat16)
    timestep = torch.tensor([500.0], device=device, dtype=torch.bfloat16)
    return latents, context, timestep


def make_actions(seed=0, batch=1, device="cuda"):
    """dataset 契约的 action_inputs：向右走为主、中途换向，delta 由 offset 差分得到。"""
    g = torch.Generator(device="cpu").manual_seed(seed)
    offset_tok = torch.zeros(batch, F, 2)
    offset_tok[:, :, 0] = torch.linspace(0, 3.0, F)  # 全程向右 3 token
    offset_tok[:, 11:, 1] = torch.linspace(0, -1.5, F - 11)  # 后半段同时向上
    offset_tok += 0.05 * torch.randn(batch, F, 2, generator=g)
    offset_tok[:, 0] = 0
    delta_tok = torch.zeros_like(offset_tok)
    delta_tok[:, 1:] = offset_tok[:, 1:] - offset_tok[:, :-1]
    action_idx = torch.full((batch, F), 7, dtype=torch.long)  # moving right
    action_idx[:, 11:] = 8  # moving up-right
    action_idx[:, 0] = 0
    return {
        "offset_px": (offset_tok * 32).to(device),
        "offset_tok": offset_tok.to(device),
        "delta_tok": delta_tok.to(device),
        "action_idx": action_idx.to(device),
    }


def err_stats(out, ref):
    out, ref = out.float(), ref.float()
    diff = (out - ref).abs()
    return {"max_abs": diff.max().item(), "rel_l2": (diff.norm() / ref.norm().clamp_min(1e-12)).item()}


def build(dit, **kw):
    from baseline.xattn import XAttnArm

    arm = XAttnArm(**kw)
    arm.install(dit)
    return arm


# ---------------------------------------------------------------- 1

@torch.no_grad()
def test_1_zero_init_equivalence(dit):
    from actionrope.arope import arope_forward

    arm = build(dit)
    assert arm.zero_init_check()
    assert next(arm.parameters()).dtype == torch.bfloat16 and next(arm.parameters()).device.type == "cuda"
    latents, context, timestep = make_inputs()
    acts = make_actions()
    ref = arope_forward(dit, latents, timestep, context)
    out = arope_forward(dit, latents, timestep, context, arm=arm, action_inputs=acts)
    s = err_stats(out, ref)
    # 装了 arm 但没给动作（action_inputs=None）也必须恒等
    out_none = arope_forward(dit, latents, timestep, context, arm=arm, action_inputs=None)
    record(t1_equiv_max_abs=s["max_abs"], t1_equiv_rel_l2=s["rel_l2"], t1_bitwise_equal=bool(torch.equal(out, ref)),
           t1_none_bitwise_equal=bool(torch.equal(out_none, ref)))
    assert torch.isfinite(out).all()
    assert torch.equal(out, ref) or s["rel_l2"] <= 1e-6
    assert torch.equal(out_none, ref)
    del arm
    torch.cuda.empty_cache()


# ---------------------------------------------------------------- 2

def test_2_perturbed_changes_output_and_grads(dit):
    from actionrope.arope import arope_forward

    arm = build(dit)
    # 只扰动零初始化的输出投影：其余参数已是 xavier；这样动作信号有了进入残差流的通路
    g = torch.Generator(device="cpu").manual_seed(0)
    with torch.no_grad():
        for mod in arm.action_modules.values():
            for name in ("proj_keyboard", "proj_mouse"):
                lin = getattr(mod, name)
                lin.weight.copy_((0.02 * torch.randn(lin.weight.shape, generator=g)).to(lin.weight))
    assert not arm.zero_init_check()

    latents, context, timestep = make_inputs(seed=1)
    acts = make_actions()
    with torch.no_grad():
        ref = arope_forward(dit, latents, timestep, context)
        out = arope_forward(dit, latents, timestep, context, arm=arm, action_inputs=acts)
        # 换一组动作（反向走），输出也要不同：说明变化来自动作而不只是权重
        acts2 = make_actions()
        acts2["action_idx"] = torch.where(acts2["action_idx"] > 0, 3, 0)
        acts2["delta_tok"] = -acts2["delta_tok"]
        out2 = arope_forward(dit, latents, timestep, context, arm=arm, action_inputs=acts2)
    s = err_stats(out, ref)
    s_act = err_stats(out2, out)
    record(t2_perturbed_rel_l2=s["rel_l2"], t2_perturbed_max_abs=s["max_abs"], t2_action_change_rel_l2=s_act["rel_l2"],
           t2_out_std_ref=ref.float().std().item(), t2_out_std_arm=out.float().std().item())
    assert s["rel_l2"] > 1e-3
    assert s_act["rel_l2"] > 1e-4
    assert abs(out.float().std().item() / ref.float().std().item() - 1) < 0.5

    # 前向 + 反向（梯度检查点开，DiT 与 arm 都要梯度）
    try:
        dit.train().requires_grad_(True)
        arm.train().requires_grad_(True)
        ctx_cells = context.unsqueeze(1).expand(1, F, L_TXT, D_TXT).contiguous()
        torch.cuda.synchronize()
        torch.cuda.reset_peak_memory_stats()
        t0 = time.time()
        pred = arope_forward(dit, latents, timestep, ctx_cells, arm=arm, action_inputs=acts,
                             use_gradient_checkpointing=True, context_ids=torch.zeros(1, F, dtype=torch.long))
        target = torch.randn_like(pred)
        loss = ((pred.float() - target.float()) ** 2)[:, :, 1:].mean()  # 帧 0 是条件帧，不计
        loss.backward()
        torch.cuda.synchronize()
        dt = time.time() - t0
        peak_gb = torch.cuda.max_memory_allocated() / 1024 ** 3
        assert torch.isfinite(loss)

        arm_named = list(arm.named_parameters())
        arm_no_grad = [n for n, p in arm_named if p.grad is None or not torch.isfinite(p.grad).all() or p.grad.abs().sum() == 0]
        dit_no_grad = [n for n, p in dit.named_parameters() if p.grad is None or p.grad.abs().sum() == 0]
        grad_norms = {
            "arm.action_modules.0.proj_keyboard.weight": arm.action_modules["0"].proj_keyboard.weight.grad.float().norm().item(),
            "arm.action_modules.0.keyboard_embed.0.weight": arm.action_modules["0"].keyboard_embed[0].weight.grad.float().norm().item(),
            "arm.action_modules.14.proj_mouse.weight": arm.action_modules["14"].proj_mouse.weight.grad.float().norm().item(),
            "arm.action_modules.14.mouse_mlp.0.weight": arm.action_modules["14"].mouse_mlp[0].weight.grad.float().norm().item(),
            "blocks.0.self_attn.q.weight": dit.blocks[0].self_attn.q.weight.grad.float().norm().item(),
            "blocks.29.ffn.2.weight": dit.blocks[29].ffn[2].weight.grad.float().norm().item(),
            "patch_embedding.weight": dit.patch_embedding.weight.grad.float().norm().item(),
        }
        grad_ok = not arm_no_grad and not dit_no_grad
        record(t2_loss=loss.item(), t2_peak_mem_gb=peak_gb, t2_fwd_bwd_sec=dt, t2_grad_norms=grad_norms,
               t2_arm_params_with_grad=f"{len(arm_named) - len(arm_no_grad)}/{len(arm_named)}",
               t2_dit_params_with_grad=f"{sum(1 for _ in dit.parameters()) - len(dit_no_grad)}/{sum(1 for _ in dit.parameters())}",
               t2_grad_ok=grad_ok)
        assert not arm_no_grad, f"arm 参数无梯度: {arm_no_grad[:5]}"
        assert not dit_no_grad, f"DiT 参数无梯度: {dit_no_grad[:5]}"
    finally:
        for p in list(dit.parameters()) + list(arm.parameters()):
            p.grad = None
        dit.eval().requires_grad_(False)
        del arm
        torch.cuda.empty_cache()


# ---------------------------------------------------------------- 3

def test_3_param_count(dit):
    arm = build(dit)
    n = arm.n_new_params()
    expected = N_BLOCKS * (KB_PER_BLOCK(12) + MOUSE_PER_BLOCK(12))
    n_dit = sum(p.numel() for p in dit.parameters())
    # 旧报告口径（+98.6M）= 只有键盘通路、每 cell 一条嵌入（window_frames=1）；上游默认 12 帧窗口的键盘通路 = 141.8M
    kb_only = build(dit, enable_mouse=False)
    kb_only_w1 = build(dit, enable_mouse=False, window_frames=1)
    n_kb, n_kb_w1 = kb_only.n_new_params(), kb_only_w1.n_new_params()
    print(f"\n[xattn] 新增参数 {n:,}（{n / 1e6:.1f}M）= 键盘 {N_BLOCKS * KB_PER_BLOCK(12):,} + 鼠标 {N_BLOCKS * MOUSE_PER_BLOCK(12):,}；"
          f"DiT {n_dit:,}；键盘-only(12 帧窗口) {n_kb:,}；键盘-only(window_frames=1，旧报告口径) {n_kb_w1:,}")
    record(t3_n_new_params=n, t3_n_new_params_M=n / 1e6, t3_n_keyboard_only=n_kb, t3_n_keyboard_only_window1=n_kb_w1,
           t3_n_dit=n_dit, t3_old_report_M=98.6)
    assert n == expected
    assert n_kb == N_BLOCKS * KB_PER_BLOCK(12)
    assert n_kb_w1 == N_BLOCKS * KB_PER_BLOCK(1)
    assert abs(n_kb_w1 / 1e6 - 98.6) < 0.1  # 旧报告的 98.6M 正好是这个配置
    del arm, kb_only, kb_only_w1
    torch.cuda.empty_cache()


# ---------------------------------------------------------------- 4

def test_4_state_dict_keys(dit):
    from baseline.xattn import XAttnArm

    arm = build(dit)
    keys = set(arm.state_dict().keys())
    assert len(keys) == len(list(arm.named_parameters()))  # 没有 buffer，全是参数
    assert all(k.startswith("action_modules.") for k in keys)
    assert {k.split(".")[1] for k in keys} == {str(i) for i in range(N_BLOCKS)}
    # 与 DiT 键集无冲突（裸键与 arm. 前缀两种都查），且 install 不改 DiT 的键集
    dit_keys = set(dit.state_dict().keys())
    assert not (keys & dit_keys) and not ({"arm." + k for k in keys} & dit_keys)
    assert dit_keys == dit._baseline_keys

    # 可 strict 加载到同配置的新 arm；配置不同（如 window 1）则 strict 失败
    sd = {k: v.clone() for k, v in arm.state_dict().items()}
    fresh = XAttnArm()
    fresh.install(dit)
    fresh.load_state_dict(sd, strict=True)
    assert all(torch.equal(a, b) for a, b in zip(fresh.state_dict().values(), sd.values()))
    other = XAttnArm(enable_mouse=False, window_frames=1)
    other.install(dit)
    with pytest.raises(RuntimeError):
        other.load_state_dict(sd, strict=True)
    record(t4_n_keys=len(keys), t4_sample_keys=sorted(keys)[:4])
    del arm, fresh, other
    torch.cuda.empty_cache()