"""
GPU test: how many Wan 2.1 1.3B weights does the CGLA DiT block load?

Builds a Wan-2.1-T2V-1.3B-shaped DiT whose every ``DiTBlock`` is replaced by
the CGLA block (``diffsynth/models/memory/u_vit_cgla_blocks.py::
CGLATransformerBlock`` — the real DFOT SSE-GLA block, sequential attn+mlp,
carrying Wan-shaped submodules ``cross_attn``/``norm1/2/3``/``ffn``/
``modulation``/``gate``). Loads the Wan 2.1 1.3B checkpoint and reports:

  * which block submodule prefixes loaded successfully (Wan key -> CGLA param)
  * which were NOT loaded, split into:
      - "unexpected" (Wan has them, CGLA block does not — e.g. Wan's softmax
        ``self_attn``), and
      - "missing"    (CGLA block has them, Wan does not — the new SSE-GLA
        attention ``spatial_*``/``temporal_attn``/``noise_write_gate``/``mlp_*``)
  * non-block Wan keys (patch_embedding / text_embedding / time_embedding /
    final_layer / ...) reported separately (not part of the DiTBlock swap)
  * coverage numbers: #keys, #params, % of Wan block params loaded, % of CGLA
    block params that are Wan-initialised.

Wan 2.1 T2V 1.3B config (``diffsynth/models/wan_video_dit.py`` hash
``9269f8db...``): dim=1536, num_heads=12, num_layers=30, ffn_dim=8960,
has_image_input=False, eps=1e-6, patch_size=(1,2,2). These are auto-detected
from the checkpoint where possible (dim/ffn_dim/num_layers/has_image_input);
``num_heads`` is not weight-inferable and uses the known config (12).

NOTE: this is a weight-LOADING test (no forward pass). Run on a GPU node that
has the Wan base model + the echo-memory env:

    PYTHONPATH=. python3 tests/test_cgla_wan.py
    # optional: --ckpt /path/to/diffusion_pytorch_model.safetensors
    # optional: --mechanism prope|ucpe   (loading result is identical across
    #                                     mechanisms; only new params differ)
"""

from __future__ import annotations

import os
import sys
from collections import Counter, defaultdict

_REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if _REPO not in sys.path:
    sys.path.insert(0, _REPO)

import torch
import torch.nn as nn
from safetensors.torch import load_file as safe_load_file

from diffsynth.models.memory.u_vit_cgla_blocks import CGLATransformerBlock, remap_wan_to_cgla

# ── Wan 2.1 T2V 1.3B config (diffsynth/models/wan_video_dit.py, hash 9269f8db…) ──
WAN_T2V_1_3B = dict(
    dim=1536,
    num_heads=12,            # not weight-inferable; from the Wan config
    num_layers=30,
    ffn_dim=8960,
    has_image_input=False,
    eps=1e-6,
    patch_size=(1, 2, 2),
)

# CGLA block forward-only shape params (do not affect weight shapes/loading):
#   640x352 frame -> VAE /8 -> 80x44 -> DiT patchify /2 -> 40x22 = 880 patches;
#   81 frames -> VAE /4 (+1) -> 21 latent frames.
CGLA_NUM_PATCHES = 880
CGLA_TEMPORAL_LENGTH = 21
CGLA_EMB_DIM = 1024       # dfot NormalizeWithCond FiLM emb_dim (CGLA-new param)
CGLA_POSE_DIM = 12        # per-frame RT camera pose
CGLA_HEAD_DIM = 128       # = dim/num_heads = 1536/12 -> SSEGLA key_dim == dim

DEFAULT_CKPT = "/apdcephfs_zwfy/share_303204533/jiakuihu/checkpoints/Wan2.1-T2V-1.3B/diffusion_pytorch_model.safetensors"


def detect_config(sd: dict) -> dict:
    """Auto-detect dim / ffn_dim / num_layers / has_image_input from the ckpt."""
    cfg = dict(WAN_T2V_1_3B)
    # num_layers: distinct blocks.<N>. prefixes.
    blk_ids = sorted({int(k.split(".")[1]) for k in sd if k.startswith("blocks.")})
    if blk_ids:
        cfg["num_layers"] = len(blk_ids)
    # dim / ffn_dim / has_image_input from block 0 weight shapes.
    def _shape(name):
        return tuple(sd[name].shape) if name in sd else None
    q = _shape("blocks.0.self_attn.q.weight") or _shape("blocks.0.cross_attn.q.weight")
    if q is not None:
        cfg["dim"] = q[0]
    ffn0 = _shape("blocks.0.ffn.0.weight")
    if ffn0 is not None:
        cfg["ffn_dim"] = ffn0[0]
    cfg["has_image_input"] = "blocks.0.cross_attn.k_img.weight" in sd
    return cfg


def build_cgla_dit(cfg: dict, mechanism: str) -> nn.Module:
    """A Wan-shaped DiT whose blocks are CGLATransformerBlock (no VAE/text).

    CGLATransformerBlock is a Wan DiTBlock with CGLA (SSE-GLA) replacing the
    self-attention (linear attention on the full flattened token sequence), then
    cross-attention to text, then FFN. Wan submodules (cross_attn/norm1/2/3/
    ffn/modulation/gate) + the SSE-GLA q/k/v/o projections load from the Wan
    checkpoint; the rest of the SSE-GLA params are new.
    """
    use_pose_rope = mechanism in ("prope", "ucpe")
    # head_dim = dim // num_heads so SSE-GLA key_dim == dim (Wan's q/k/v/o are
    # Linear(dim, dim), loadable into the linear-attention projections). For
    # real Wan 1.3B this is 1536 // 12 = 128 (== CGLA_HEAD_DIM).
    head_dim = cfg["dim"] // cfg["num_heads"]

    class CGLADiT(nn.Module):
        def __init__(self):
            super().__init__()
            self.blocks = nn.ModuleList([
                CGLATransformerBlock(
                    has_image_input=cfg["has_image_input"],
                    dim=cfg["dim"],
                    num_heads=cfg["num_heads"],
                    ffn_dim=cfg["ffn_dim"],
                    eps=cfg["eps"],
                    head_dim=head_dim,
                    num_sparse_partition=4,
                    num_writer=1,
                    num_reader=1,
                    pose_dim=CGLA_POSE_DIM,
                    pose_bottleneck=64,
                    gate_logit_normalizer=16,
                    gate_low_rank_dim=16,
                    use_pose_rope=use_pose_rope,
                    use_pose_gate_mod=False,
                    layer_idx=i,
                    emb_dim=CGLA_EMB_DIM,
                )
                for i in range(cfg["num_layers"])
            ])

    return CGLADiT()


def _top(prefix_key: str) -> str:
    """Block submodule prefix: 'blocks.N.<attr>...' -> '<attr>';
    non-block key -> its first component ('patch_embedding.weight' -> 'patch_embedding')."""
    if prefix_key.startswith("blocks."):
        parts = prefix_key.split(".")  # ["blocks", "<N>", "<attr>", ...]
        return parts[2] if len(parts) > 2 else parts[-1]
    return prefix_key.split(".")[0]


def _numel(sd, keys):
    return int(sum(sd[k].numel() for k in keys if k in sd))


def _prefix_summary(keys):
    """Group keys by their first component; return {prefix: (n_keys, n_params)}."""
    # keys here are full ckpt keys (blocks.N.<rest>) — group by <rest>[0].
    g = defaultdict(list)
    for k in keys:
        rest = k.split(".", 2)[2] if k.startswith("blocks.") else k
        g[rest.split(".")[0]].append(k)
    return g


def report(sd, model, cfg, mechanism):
    n_blk = cfg["num_layers"]

    # Remap Wan self-attn keys (self_attn.{q,k,v,o}) to CGLA SSE-GLA names
    # (self_attn.{q,k,v,o}_proj) so Wan's softmax-attn weights initialise the
    # linear-attention projections. After remap, "loaded" reflects this.
    sd = remap_wan_to_cgla(sd)

    model_sd = model.state_dict()
    ckpt_keys = set(sd.keys())
    model_keys = set(model_sd.keys())

    # Block-only keys (the DiTBlock swap domain).
    ckpt_blk = {k for k in ckpt_keys if k.startswith("blocks.")}
    model_blk = {k for k in model_keys if k.startswith("blocks.")}
    # Non-block Wan keys (patch_embedding / text_embedding / time_embedding / ...).
    ckpt_nonblk = ckpt_keys - ckpt_blk

    missing, unexpected = model.load_state_dict(sd, strict=False)
    # Restrict to block keys for the swap analysis (non-block keys are reported
    # separately as "not part of the DiTBlock replacement").
    missing_blk = [k for k in missing if k.startswith("blocks.")]
    unexpected_blk = [k for k in unexpected if k.startswith("blocks.")]
    unexpected_nonblk = [k for k in unexpected if not k.startswith("blocks.") and k in ckpt_nonblk]

    loaded_blk = sorted(model_blk - set(missing_blk))  # in both -> loaded

    # ── param tallies ──────────────────────────────────────────────────────
    wan_blk_params = _numel(sd, ckpt_blk)
    loaded_params = _numel(sd, loaded_blk)
    cgla_blk_params = _numel(model_sd, model_blk)
    wan_init_params = loaded_params

    # ── per-prefix grouping (over one block, × num_layers) ─────────────────
    def _per_prefix(keys, src):
        g = defaultdict(lambda: [0, 0])  # prefix -> [n_keys, n_params]
        for k in keys:
            pref = _top(k)
            g[pref][0] += 1
            g[pref][1] += int(src[k].numel()) if k in src else 0
        return g

    loaded_g = _per_prefix(loaded_blk, sd)
    miss_g = _per_prefix(missing_blk, model_sd)
    unexp_g = _per_prefix(unexpected_blk, sd)
    nonblk_g = _per_prefix(ckpt_nonblk, sd)

    # ── print ──────────────────────────────────────────────────────────────
    print("=" * 78)
    print("WAN 2.1 1.3B -> CGLA DiT  weight-loading report")
    print("=" * 78)
    print(f"checkpoint          : {CKPT}")
    print(f"detected Wan config : {cfg}")
    print(f"CGLA mechanism      : {mechanism}")
    print(f"  (use_pose_rope={mechanism in ('prope','ucpe')}; loading coverage is")
    print(f"   identical across cgla/prope/ucpe — only *new* (non-Wan) params differ)")
    print(f"num blocks          : {n_blk}")
    print(f"head_dim (SSEGLA)   : {CGLA_HEAD_DIM}  (= dim/num_heads -> SSEGLA key_dim == dim)")
    print()

    print("-" * 78)
    print("1) SUCCESSFULLY LOADED  (Wan block key -> CGLA param, name+shape match)")
    print("-" * 78)
    if loaded_g:
        print(f"  {'prefix':<22}{'keys (×N blk)':>16}{'params':>16}")
        for pref in sorted(loaded_g):
            n, p = loaded_g[pref]
            print(f"  {pref:<22}{n:>16}{p:>16,}")
    print(f"  {'(total loaded)':<22}{'':>16}{loaded_params:>16,}")
    print()

    print("-" * 78)
    print("2a) NOT LOADED — unexpected  (in Wan ckpt, NOT in CGLA block)")
    print("-" * 78)
    if unexp_g:
        print(f"  {'prefix':<22}{'keys (×N blk)':>16}{'params':>16}")
        for pref in sorted(unexp_g):
            n, p = unexp_g[pref]
            print(f"  {pref:<22}{n:>16}{p:>16,}")
        print(f"  {'(total unexpected)':<22}{'':>16}{_numel(sd, unexpected_blk):>16,}")
    else:
        print("  (none)")
    print()

    print("-" * 78)
    print("2b) NOT LOADED — missing  (CGLA block has, NOT in Wan ckpt = new CGLA params)")
    print("-" * 78)
    if miss_g:
        print(f"  {'prefix':<22}{'keys (×N blk)':>16}{'params':>16}")
        for pref in sorted(miss_g):
            n, p = miss_g[pref]
            print(f"  {pref:<22}{n:>16}{p:>16,}")
        print(f"  {'(total missing/new)':<22}{'':>16}{_numel(model_sd, missing_blk):>16,}")
    else:
        print("  (none)")
    print()

    print("-" * 78)
    print("3) NON-BLOCK Wan keys  (not part of the DiTBlock swap; e.g. patch/text")
    print("   /time/final embeddings — the CGLA model has no such submodules)")
    print("-" * 78)
    if nonblk_g:
        print(f"  {'prefix':<22}{'keys':>10}{'params':>16}")
        for pref in sorted(nonblk_g):
            n, p = nonblk_g[pref]
            print(f"  {pref:<22}{n:>10}{p:>16,}")
    else:
        print("  (none)")
    print()

    # ── coverage summary ───────────────────────────────────────────────────
    wan_total = sum(v.numel() for v in sd.values())
    pct_wan_loaded = 100.0 * loaded_params / wan_blk_params if wan_blk_params else 0.0
    pct_cgla_waninit = 100.0 * wan_init_params / cgla_blk_params if cgla_blk_params else 0.0
    print("=" * 78)
    print("COVERAGE SUMMARY")
    print("=" * 78)
    print(f"  Wan ckpt total params              : {wan_total:>14,}")
    print(f"  Wan BLOCK params (blocks.*.*       ) : {wan_blk_params:>14,}")
    print(f"  -> loaded into CGLA block          : {loaded_params:>14,}  ({pct_wan_loaded:5.2f}% of Wan block params)")
    print(f"  CGLA block params (new + Wan-shape): {cgla_blk_params:>14,}")
    print(f"  -> Wan-initialised (loaded)         : {wan_init_params:>14,}  ({pct_cgla_waninit:5.2f}% of CGLA block params)")
    print(f"  -> CGLA-new (missing, needs training): {cgla_blk_params - wan_init_params:>14,}  "
          f"({100.0*(cgla_blk_params-wan_init_params)/cgla_blk_params:5.2f}% of CGLA block params)")
    print("=" * 78)

    _report_ssegla_param_audit(model, sd, mechanism)


def _self_attn_subparam(key: str):
    """For 'blocks.N.self_attn.<sub>.[...]' -> '<sub>'; else None."""
    parts = key.split(".")
    if len(parts) >= 3 and parts[0] == "blocks" and parts[2] == "self_attn":
        return parts[3] if len(parts) > 3 else None
    return None


def _report_ssegla_param_audit(model, sd, mechanism):
    """Per-param audit of the SSE-GLA ``self_attn`` submodules.

    Answers: for each SSEGLA param (q_proj / k_proj / v_proj / o_proj /
    lora_q_proj / lora_k_proj / gk_proj / e_proj / g_proj / o_norm) — LOADED from
    Wan (name+shape match via the remap) or REINIT (no Wan key)? And for the
    PRoPE/UCPE pose modules (pose_encoder / pose_q_proj / pose_k_proj /
    pose_gk_proj / pose_e_proj / pose_rope / pose_gate_mod) — all REINIT, with
    their dfot zero-init policy noted (the 'up' projections are zero-init; the
    'down' pose_encoder is default-init).
    """
    model_sd = model.state_dict()
    # Group by self_attn subparam over all blocks.
    def _group(keys, src):
        g = defaultdict(lambda: [0, 0])  # subparam -> [n_keys, n_params]
        for k in keys:
            sp = _self_attn_subparam(k)
            if sp is None:
                continue
            g[sp][0] += 1
            g[sp][1] += int(src[k].numel()) if k in src else 0
        return g

    model_sub = _group(model_sd.keys(), model_sd)      # CGLA block params
    ckpt_sub = _group(sd.keys(), sd)                   # Wan (remapped) params

    # Wan self_attn projection names that the remap targets.
    WAN_REMAP_TARGETS = {"q_proj", "k_proj", "v_proj", "o_proj"}
    # dfot pose 'up'-projection zero-init (see fla/layers/sse.py:352-355, 359-360;
    # PoseRoPE net[-1] zero-init at sse.py:98-99).
    POSE_ZERO_UP = {
        "pose_q_proj", "pose_k_proj", "pose_gk_proj", "pose_e_proj",
        "pose_gate_mod", "pose_rope",
    }

    print()
    print("=" * 78)
    print("4) SSE-GLA self_attn  PARAM-BY-PARAM  (loaded vs reinitialised)")
    print("=" * 78)
    print(f"  {'SSEGLA param':<16}{'Wan source':<20}{'keys (×N blk)':>14}{'params':>14}  status")
    print("  " + "-" * 74)
    total_loaded_sa = 0
    total_reinit_sa = 0
    for sp in sorted(model_sub):
        n, p = model_sub[sp]
        wan_src = "—"
        status = "REINIT (no Wan key)"
        if sp in WAN_REMAP_TARGETS:
            wan_name = sp.replace("_proj", "")   # q_proj -> q
            wan_src = f"self_attn.{wan_name}"
            if sp in ckpt_sub:
                status = "LOADED"
                total_loaded_sa += p
            else:
                status = "REINIT (Wan key absent)"
                total_reinit_sa += p
        else:
            total_reinit_sa += p
        note = ""
        if sp in POSE_ZERO_UP and mechanism in ("prope", "ucpe"):
            note = " [dfot zero-init 'up' proj]"
        elif sp == "pose_encoder" and mechanism in ("prope", "ucpe"):
            note = " [dfot default-init 'down' proj]"
        print(f"  {sp:<16}{wan_src:<20}{n:>14}{p:>14,}  {status}{note}")
    print("  " + "-" * 74)
    print(f"  self_attn loaded   (q/k/v/o_proj from Wan): {total_loaded_sa:>14,}")
    print(f"  self_attn reinit   (GLA gates/LoRA/norm):    {total_reinit_sa:>14,}")
    print()

    # Pose modules exist only for PRoPE / UCPE.
    print("-" * 78)
    if mechanism in ("prope", "ucpe"):
        print(f"PRoPE/UCPE pose modules (use_pose_rope=True) — ALL reinitialised:")
        print(f"  {'pose module':<16}{'init policy (dfot)':<40}{'params':>14}")
        print("  " + "-" * 70)
        pose_rows = [
            ("pose_encoder",  "default-init (down proj; NOT zeroed)"),
            ("pose_q_proj",   "zero-init 'up' proj  => pose contributes 0 at step 0"),
            ("pose_k_proj",   "zero-init 'up' proj"),
            ("pose_gk_proj",  "zero-init 'up' proj"),
            ("pose_e_proj",   "zero-init (direct)"),
            ("pose_rope",     "PoseRoPE net[-1] zero-init (angle MLP)"),
        ]
        if any(_self_attn_subparam(k) == "pose_gate_mod" for k in model_sd):
            pose_rows.append(("pose_gate_mod", "zero-init (weight+bias)"))
        for name, policy in pose_rows:
            p = model_sub.get(name, [0, 0])[1]
            present = "present" if name in model_sub else "ABSENT"
            print(f"  {name:<16}{policy:<40}{p:>14,}  {present}")
        print("  => at step 0, pose contributes exactly 0 to q2/k2/gk2/eta (LoRA-style),")
        print("     so PRoPE/UCPE is numerically identical to vanilla CGLA until trained.")
    else:
        print(f"mechanism={mechanism!r}: no pose modules (use_pose_rope=False).")
    print("=" * 78)


def parse_args():
    import argparse
    p = argparse.ArgumentParser(description="CGLA <-> Wan 2.1 1.3B weight-loading test")
    p.add_argument("--ckpt", default=DEFAULT_CKPT,
                   help="path to diffusion_pytorch_model.safetensors")
    p.add_argument("--mechanism", default="prope", choices=["cgla", "prope", "ucpe"],
                   help="CGLA variant (loading coverage is identical across them)")
    return p.parse_args()


CKPT = DEFAULT_CKPT  # set in main() from args


def main():
    global CKPT
    args = parse_args()
    CKPT = args.ckpt
    if not os.path.isfile(CKPT):
        print(f"ERROR: checkpoint not found: {CKPT}", file=sys.stderr)
        print("  set --ckpt /path/to/Wan2.1-T2V-1.3B/diffusion_pytorch_model.safetensors",
              file=sys.stderr)
        sys.exit(1)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"[cgla_wan] loading checkpoint (CPU): {CKPT}")
    sd = safe_load_file(CKPT)
    cfg = detect_config(sd)
    print(f"[cgla_wan] detected config: {cfg}")

    # Build the CGLA DiT (CPU — loading does not need GPU) and load Wan weights.
    model = build_cgla_dit(cfg, args.mechanism)
    report(sd, model, cfg, args.mechanism)

    # Confirm the CGLA block constructs on GPU (fla/triton compile at forward,
    # not at construction; this just moves params to cuda).
    if device == "cuda":
        try:
            model = model.to(device=device, dtype=torch.bfloat16)
            n = sum(p.numel() for p in model.parameters())
            print(f"[cgla_wan] CGLA DiT moved to {device} (bf16); {n:,} params construct OK")
        except Exception as e:
            print(f"[cgla_wan] WARN: GPU move failed: {e}")
    else:
        print("[cgla_wan] no CUDA — skipping GPU construct check")


if __name__ == "__main__":
    main()