"""CPU integration checks against real tiny Transformers Qwen3.5 layers."""
import importlib
import sys
from pathlib import Path
import torch
from torch import nn
from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig
from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5TextModel

root = Path(__file__).parent
repo = root / sys.argv[1] if len(sys.argv) > 1 else root.parent
sys.path.insert(0, str(repo))
from dqwen.core import DiffusionLM, patch_bidirectional
from dqwen.kvcache import BlockCache, _segment, _patch, enable_kv_cache, disable_kv_cache, stream_kv
from dqwen import hub

torch.manual_seed(123)
torch.set_num_threads(1)
cfg = Qwen3_5TextConfig(vocab_size=32, hidden_size=16, intermediate_size=32,
    num_hidden_layers=2, num_attention_heads=2, num_key_value_heads=1, head_dim=8,
    linear_num_key_heads=1, linear_num_value_heads=2, linear_key_head_dim=4,
    linear_value_head_dim=4, linear_conv_kernel_dim=4,
    layer_types=['linear_attention', 'full_attention'],
    rope_parameters={'rope_type': 'default', 'rope_theta': 10000.0,
                     'mrope_section': [1, 1, 2]})
cfg._attn_implementation = 'sdpa'
lm = nn.Module()
lm.text = Qwen3_5TextModel(cfg).eval()
lm.lm_head = nn.Linear(16, 32, bias=False)
patch_bidirectional(lm)
for m in lm.modules():
    if hasattr(m, 'rev_gate'):
        m.rev_gate.data.fill_(0.25)
dlm = DiffusionLM(lm, 31).eval()
dlm.eos_id = 30
dlm.bart = None
ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]])

with torch.no_grad():
    baseline = dlm.hidden(ids)
    _patch(dlm)
    # Prefix-empty and suffix-empty cases, plus a genuine two-sided block.
    for lo, hi in [(1, 4), (3, 6), (3, 8)]:
        cache = BlockCache(lo, hi, ids.shape[1])
        pos = torch.arange(ids.shape[1])
        record = dlm.text(inputs_embeds=dlm.embed(ids),
            attention_mask={'full_attention': None, 'linear_attention': None},
            position_ids=pos[None], past_key_values=cache, use_cache=True).last_hidden_state
        torch.testing.assert_close(record, baseline, atol=2e-5, rtol=2e-4)
        saved = {i: tuple(t.clone() if t is not None else None for t in s)
                 for i, s in cache.f_state.items()}
        cache.mode = 'step'
        pos = torch.arange(lo - 1, hi)
        for _ in range(2):
            active = dlm.text(inputs_embeds=dlm.embed(ids[:, lo-1:hi]),
                attention_mask={'full_attention': None, 'linear_attention': None},
                position_ids=pos[None], past_key_values=cache, use_cache=True).last_hidden_state
            torch.testing.assert_close(active, record[:, lo-1:hi], atol=2e-5, rtol=2e-4)
        for i, state in saved.items():
            for a, b in zip(state, cache.f_state[i]):
                if a is not None:
                    torch.testing.assert_close(a, b, atol=0, rtol=0)

class Tok:
    def apply_chat_template(self, *args, **kwargs): return 'prompt'
    def __call__(self, *args, **kwargs): return {'input_ids': [1, 2, 3]}
    def decode(self, ids, **kwargs): return str(ids)

tok = Tok()
enable_kv_cache(dlm)
results = list(hub.stream(dlm, tok, 'test', 4, 1.0))
assert results and results[-1][2] <= 4
# 4B/9B generate passes block_len positionally; 27B decodes lists.
assert isinstance(hub.generate(dlm, tok, 'test', gen_len=4), str)
if not hub._kv_legacy_list:
    assert list(hub.stream(dlm, tok, 'test', 4, 1.0, 2))
# Fully refreshed runs should match the uncached block generator (MASK is banned).
refreshed = list(stream_kv(dlm, tok, 'test', 4, 1.0, block_len=2, refresh=1))
assert refreshed
# Check refresh=1 against an independently denoised full-sequence reference.
reference = torch.tensor([[1, 2, 3, 31, 31, 31, 31]])
expected = []
with torch.no_grad():
    for lo in (3, 5):
        hi = lo + 2
        while (reference[0, lo:hi] == dlm.mask_id).any():
            h = dlm.hidden(reference)
            pos = torch.arange(lo, hi)[reference[0, lo:hi] == dlm.mask_id]
            logits = dlm.head(h[0, pos-1]).float()
            logits[:, dlm.mask_id] = -torch.inf
            conf, cand = logits.softmax(-1).max(-1)
            select = conf >= 1.0
            if not select.any(): select[conf.argmax()] = True
            reference[0, pos[select]] = cand[select]
            expected.append(reference[:, 3:].clone())
        if (reference[0, 3:hi] == dlm.eos_id).any():
            reference[0, hi:] = dlm.eos_id
            expected.append(reference[:, 3:].clone())
            break
assert len(expected) == len(refreshed)
for want, (got, _, _) in zip(expected, refreshed):
    torch.testing.assert_close(want, got)
assert list(stream_kv(dlm, tok, 'test', 4, 1.0, max_steps=1))[-1][2] == 1
disable_kv_cache(dlm)
assert not any(hasattr(m, '_dq_kv_old') for m in lm.modules())
torch.testing.assert_close(dlm.hidden(ids), baseline, atol=2e-5, rtol=2e-4)
enable_kv_cache(dlm, refresh=1)
assert list(hub.stream(dlm, tok, 'test', 2, 1.0))
disable_kv_cache()
print('PASS:', repo.name, 'record/active equivalence, empty segments, immutable state, stream/generate, restore')