"""CPU check of the host-side packing used by NkiAttnBlock (lane permutation, 1/sqrt(d) fold, NTK RoPE
table, key bias) against the real Hugging Face NewLayer in fp32 with ragged right padding.

The packed operands go through a NumPy emulation of attn_block_kernel's math (no NKI, no device), so
this test isolates packing errors from kernel errors. The RoPE table is built the way NkiAttnBlock builds
it: once at max(ATTN_LENGTHS), then sliced per length.
Needs torch + transformers; downloads the model on first run.
"""
import os

import numpy as np
import pytest

from conftest import layer_norm

from attn_block_pack import pack

torch = pytest.importorskip("torch")
transformers = pytest.importorskip("transformers")
pytestmark = pytest.mark.cpu

MODELS = [os.environ.get("GTE_TEST_MODEL", "Alibaba-NLP/gte-multilingual-reranker-base")]
LMAX = 1024                  # == max(torch_ops.ATTN_LENGTHS)


@pytest.fixture(scope="module", params=MODELS)
def new_model(request):
    m = transformers.AutoModel.from_pretrained(request.param, trust_remote_code=True,
                                               attn_implementation="eager").eval()
    return m


def _emulate(pk, xn, kb, NH, D, L):
    """attn_block_kernel math on packed operands, fp32."""
    B = xn.shape[0]
    H = NH * D
    qk = xn @ pk["wqk"] + pk["bqk"].T.reshape(-1)              # [B, L, 2H], permuted lanes, Q pre-scaled
    c, s = pk["cs"][0][:, :L].T, pk["cs"][1][:, :L].T            # [L, 128]

    def rope(t):
        out = np.zeros((B, L, NH, D), np.float32)
        for j in range(NH // 4):
            lo = t[..., (2 * j) * 128:(2 * j + 1) * 128]
            hi = t[..., (2 * j + 1) * 128:(2 * j + 2) * 128]
            rlo, rhi = lo * c - hi * s, hi * c + lo * s
            for bq in range(4):
                out[:, :, j * 4 + bq, :32] = rlo[..., bq * 32:(bq + 1) * 32]
                out[:, :, j * 4 + bq, 32:] = rhi[..., bq * 32:(bq + 1) * 32]
        return out

    q, k = rope(qk[..., :H]), rope(qk[..., H:])
    v = (xn @ pk["wv"] + pk["bv"][0]).reshape(B, L, NH, D)
    sc = np.einsum("blhd,bmhd->bhlm", q, k) + kb[:, None, None, :]
    sc -= sc.max(-1, keepdims=True)
    p = np.exp(sc)
    p /= p.sum(-1, keepdims=True)
    o = np.einsum("bhlm,bmhd->blhd", p, v).reshape(B, L, H) @ pk["wo"] + pk["bo"][0]
    return layer_norm(xn + o, pk["g"][0], pk["b"][0])


@pytest.mark.parametrize("L", [128, 256, 384, 512, 768, 1024])
def test_pack_matches_hf_newlayer(new_model, L):
    enc = new_model
    cfg = enc.config
    H, NH = cfg.hidden_size, cfg.num_attention_heads
    D = H // NH
    lay = enc.encoder.layer[min(3, cfg.num_hidden_layers - 1)]
    rot = enc.embeddings.rotary_emb
    B = 2
    torch.manual_seed(0)
    x = torch.randn(B, L, H)
    mask = torch.ones(B, L, dtype=torch.long)
    mask[0, L - 40:] = 0
    mask[1, L // 3:] = 0
    bias = enc.get_extended_attention_mask(mask, (B, L))
    cos, sin = rot(x, seq_len=L)
    pos = torch.arange(L)
    with torch.no_grad():
        a = lay.attention(x, bias, (cos[pos][None, :, None, :], sin[pos][None, :, None, :]))[0]
        ref = lay.attn_ln(x + a).numpy()
    at = lay.attention
    f = lambda t: t.detach().float().numpy()
    cmax, smax = rot(torch.zeros(1, dtype=torch.float32), seq_len=LMAX)
    pk = pack(f(at.qkv_proj.weight), f(at.qkv_proj.bias), f(at.o_proj.weight), f(at.o_proj.bias),
              f(lay.attn_ln.weight), f(lay.attn_ln.bias), f(cmax), f(smax), H, 1.0 / np.sqrt(D))
    kb = np.clip(bias.reshape(B, L).numpy(), -30000.0, 0)
    y = _emulate(pk, x.numpy(), kb, NH, D, L)
    valid = mask.bool().numpy()
    rel = np.linalg.norm((y - ref)[valid]) / np.linalg.norm(ref[valid])
    assert rel < 1e-4, rel