"""End-to-end verification of the converted RWKV-7 HF model."""
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

PATH = "/workspace/rwkv7-g1d-olmo"

print("=" * 70)
print("1) Loading with AutoModelForCausalLM(trust_remote_code=True)")
tok = AutoTokenizer.from_pretrained(PATH, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    PATH, trust_remote_code=True, dtype=torch.bfloat16
).cuda()
model.eval()
print("   model class:", type(model).__name__)
print("   emb:", tuple(model.get_input_embeddings().weight.shape),
      "head:", tuple(model.get_output_embeddings().weight.shape))
nparams = sum(p.numel() for p in model.parameters())
print(f"   params: {nparams/1e6:.1f}M")

ids = tok("The Eiffel tower is in the city of", return_tensors="pt").input_ids.cuda()
print("   input ids:", ids.tolist())

print("=" * 70)
print("2) Forward with CUDA kernel")
with torch.no_grad():
    out_k = model(ids).logits
print("   logits:", tuple(out_k.shape), out_k.dtype)

print("=" * 70)
print("3) Forward with PyTorch fallback + kernel-vs-fallback parity")
model.config.use_cuda_kernel = False
with torch.no_grad():
    out_f = model(ids).logits
model.config.use_cuda_kernel = True
diff = (out_k.float() - out_f.float()).abs()
print(f"   max abs diff kernel vs fallback: {diff.max().item():.4e}")
print(f"   mean abs diff:                   {diff.mean().item():.4e}")
# argmax agreement on last token
print("   kernel top-1 next id:", out_k[0, -1].argmax().item(),
      "| fallback:", out_f[0, -1].argmax().item())

print("=" * 70)
print("4) Cached vs uncached parity (prefill + stateful decode vs full recompute)")
# The stateful path always runs on the PyTorch implementation; compare against a
# full PyTorch-fallback forward so both sides share the same numerics (fp32
# internal, sequential over time) -> should match (near-)exactly.
SPLIT = 4  # decode the last SPLIT tokens one at a time
assert ids.shape[1] > SPLIT
model.config.use_cuda_kernel = False
with torch.no_grad():
    full = model(ids).logits
    pre = model(ids[:, :-SPLIT], use_cache=True)      # prefill -> state
    st, chunks = pre.state, [pre.logits]
    for t in range(ids.shape[1] - SPLIT, ids.shape[1]):
        o = model(ids[:, t:t + 1], state=st)          # passing state implies use_cache
        st, chunks = o.state, chunks + [o.logits]
    inc = torch.cat(chunks, dim=1)
model.config.use_cuda_kernel = True
diff_c = (inc.float() - full.float()).abs()
print(f"   max abs diff cached vs uncached: {diff_c.max().item():.4e}")
print(f"   mean abs diff:                   {diff_c.mean().item():.4e}")
assert torch.allclose(inc.float(), full.float(), atol=1e-3, rtol=0), \
    "cached (prefill+decode) logits diverge from full recompute"
# also check against the kernel forward with a looser bf16 tolerance
diff_ck = (inc.float() - out_k.float()).abs()
print(f"   max abs diff cached vs kernel:   {diff_ck.max().item():.4e}")
ax, wkv, fx = st[0]
print("   state[0] shapes:",
      "att_x_prev", tuple(ax.shape), ax.dtype, "|",
      "wkv", tuple(wkv.shape), wkv.dtype, "|",
      "ffn_x_prev", tuple(fx.shape), fx.dtype)

print("=" * 70)
print("5) Backward pass (gradients flow through kernel)")
model.train()
ids2 = tok("Backward test sentence for gradient check.", return_tensors="pt").input_ids.cuda()
out = model(ids2, labels=ids2)
loss = out.loss
loss.backward()
gnorm_emb = model.get_input_embeddings().weight.grad
g_att = model.rwkv.blocks[0].att.receptance.weight.grad
g_w1 = model.rwkv.blocks[6].att.w1.grad
print(f"   loss: {loss.item():.4f}")
print(f"   emb.grad is not None: {gnorm_emb is not None}, norm={gnorm_emb.float().norm().item():.4e}")
print(f"   blocks.0.att.receptance.grad norm: {g_att.float().norm().item():.4e}")
print(f"   blocks.6.att.w1 (decay-lora) grad norm: {g_w1.float().norm().item():.4e}")
n_with_grad = sum(1 for p in model.parameters() if p.grad is not None and p.grad.abs().sum() > 0)
n_total = sum(1 for _ in model.parameters())
print(f"   params with non-zero grad: {n_with_grad}/{n_total}")
model.zero_grad(set_to_none=True)

print("=" * 70)
print("6) generate() through the state cache")
model.eval()
with torch.no_grad():
    gen = model.generate(ids, max_new_tokens=10, do_sample=False)
print("   generated:", tok.decode(gen[0]))

# generate() must equal a manual RNN-mode greedy loop (same state-cache path)
with torch.no_grad():
    o = model(ids, use_cache=True)
    cur = o.logits[:, -1].argmax(-1, keepdim=True)
    st, toks = o.state, [cur]
    for _ in range(9):
        o = model(cur, state=st)
        st = o.state
        cur = o.logits[:, -1].argmax(-1, keepdim=True)
        toks.append(cur)
manual = torch.cat(toks, dim=1)
match = bool((gen[:, ids.shape[1]:] == manual).all())
print("   generate == manual stateful greedy:", match)
assert match, "generate() output diverges from manual stateful greedy decode"

print("=" * 70)
print("ALL CHECKS DONE")