import sys, tempfile, torch
from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig
from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5ForCausalLM
from looped_gdn import build_from_source

torch.manual_seed(0)
L = 8
cfg = Qwen3_5TextConfig(
    vocab_size=500, hidden_size=128, intermediate_size=256, num_hidden_layers=L,
    layer_types=["linear_attention"] * L, linear_num_key_heads=4, linear_num_value_heads=4,
    linear_key_head_dim=32, linear_value_head_dim=32, num_attention_heads=2, num_key_value_heads=1,
    head_dim=64, tie_word_embeddings=True,
    rope_parameters={"rope_type": "default", "rope_theta": 1e7, "partial_rotary_factor": 0.25,
                     "mrope_section": [3, 3, 2], "mrope_interleaved": True},
)
src = Qwen3_5ForCausalLM(cfg).eval()
with torch.no_grad():  # make per-layer small params distinct so merging is non-trivial
    for l in src.model.layers:
        l.input_layernorm.weight.normal_(0, 0.1); l.post_attention_layernorm.weight.normal_(0, 0.1)
d = tempfile.mkdtemp(); src.save_pretrained(d)
ids = torch.randint(0, 500, (2, 12))
ref = src(ids, use_cache=False).logits

ok = True
# 1) r = 1 (G = L): must reproduce the source exactly (gate ~1, RoPE-injection alpha=0)
m1 = build_from_source(d, num_groups=L, dtype=torch.float32, gate_init_bias=30.0).eval()
out = m1(ids)["logits"]
err = (out - ref).abs().max().item(); print("r=1 vs source max|diff| =", err); ok &= err < 1e-4

# 2) spatial merge G=2, r=4 ; 3) cached decode == full forward
m = build_from_source(d, num_groups=2, dtype=torch.float32).eval()
print("G=2 r=4 params:", sum(p.numel() for p in m.parameters()), "vs source", sum(p.numel() for p in src.parameters()))
full = m(ids)["logits"]
o = m(ids[:, :8], use_cache=True); c = o["past_key_values"]; steps = [o["logits"]]
for t in range(8, 12):
    o = m(ids[:, t:t + 1], past_key_values=c, use_cache=True); steps.append(o["logits"])
inc = torch.cat(steps, 1)
err2 = (inc - full).abs().max().item(); print("cached decode vs full max|diff| =", err2); ok &= err2 < 1e-3

# 4) gradients reach gates / RoPE-injection / per-step params; non-trivial RoPE alters output
m.train(); m(ids)["logits"].float().pow(2).mean().backward()
for n in ["rope_alpha", "step_embed", "groups.0.gate_w", "groups.0.gate_b", "groups.0.ln1", "groups.0.A_log"]:
    g = dict(m.named_parameters())[n].grad; print(f"grad {n:18s}", None if g is None else f"{g.norm():.3e}"); ok &= g is not None
with torch.no_grad(): m.rope_alpha.fill_(0.5)
m.eval(); print("RoPE active changes output:", (m(ids)["logits"] - full).abs().max().item() > 1e-3)

# 5) early exit (no cache)
m.halt_threshold = 0.5
with torch.no_grad(): m.groups[1].gate_b[2:] = -10
r = m(ids, return_gates=True); print("executed inner steps:", len(r["gates"]), "/", m.G * m.r)
print("generate:", m.generate(ids[:1, :5], 5).tolist())
print("ALL OK" if ok else "FAIL"); sys.exit(0 if ok else 1)