"""nki.simulate parity of the three kernels vs NumPy fp32 references. Simulator only: NOT device evidence
(the simulator does not check SBUF/PSUM capacity and has known DMA-pattern fidelity gaps)."""
import numpy as np
import pytest

nki = pytest.importorskip("nki")
ml_dtypes = pytest.importorskip("ml_dtypes")

from conftest import EPS, attn_block_ref, gated_mlp_ln_ref, layer_norm, rel_l2, rope_table  # noqa: E402

from attn_block_kernel import attn_block_kernel  # noqa: E402
from attn_block_pack import pack  # noqa: E402
from gated_mlp_kernel import gated_mlp_ln_kernel, gated_mlp_ln_stream_kernel, mlp_weights_fit  # noqa: E402
from res_ln_kernel import res_ln_kernel  # noqa: E402

bf16 = ml_dtypes.bfloat16
TOL = 1e-2          # bf16 in/out; observed ~3-5e-3
pytestmark = pytest.mark.sim


def _bc(v):
    return np.tile(np.asarray(v, np.float32), (128, 1))


# ---------------------------------------------------------------- residual + LayerNorm
@pytest.mark.parametrize("R,lnc,has_res,H", [(256, 1, True, 768), (512, 2, True, 768), (256, 1, False, 768),
                                              (1024, 2, False, 768), (256, 1, True, 1024), (512, 2, True, 1024)])
def test_res_ln(R, lnc, has_res, H):
    rng = np.random.default_rng(R + H)
    x = (rng.standard_normal((R, H)) * 3 + 0.5).astype(bf16)
    r = rng.standard_normal((R, H)).astype(bf16)
    g = 1 + 0.1 * rng.standard_normal(H)
    b = 0.1 * rng.standard_normal(H)
    y = nki.simulate(res_ln_kernel[lnc])(x, r, _bc(g), _bc(b), EPS, has_res)
    s = x.astype(np.float32) + (r.astype(np.float32) if has_res else 0)
    assert rel_l2(y, layer_norm(s, g, b)) < TOL


# ---------------------------------------------------------------- gated MLP + residual + LayerNorm
def _mlp_case(R, K, I, seed):
    rng = np.random.default_rng(seed)
    x = rng.standard_normal((R, K)).astype(bf16)
    wu = (rng.standard_normal((I, K)) / np.sqrt(K)).astype(bf16)
    wg = (rng.standard_normal((I, K)) / np.sqrt(K)).astype(bf16)
    wd = (rng.standard_normal((K, I)) / np.sqrt(I)).astype(bf16)
    bd, g, b = 0.1 * rng.standard_normal(K), 1 + 0.1 * rng.standard_normal(K), 0.1 * rng.standard_normal(K)
    args = (x, np.ascontiguousarray(wu.T), np.ascontiguousarray(wg.T), np.ascontiguousarray(wd.T),
            _bc(bd), _bc(g), _bc(b), EPS)
    ref = gated_mlp_ln_ref(x, wu.astype(np.float32), wg.astype(np.float32), wd.astype(np.float32), bd, g, b)
    return args, ref


@pytest.mark.parametrize("R,lnc", [(256, 1), (640, 1), (512, 2)])
def test_gated_mlp_ln(R, lnc):
    assert mlp_weights_fit(768, 3072)
    args, ref = _mlp_case(R, 768, 3072, R)
    assert rel_l2(nki.simulate(gated_mlp_ln_kernel[lnc])(*args), ref) < TOL


# streaming variant: 1024/4096 (gte-large) does not fit resident; also 768/3072 for a direct comparison,
# and R > STREAM_SB (two super-blocks) with a partial last super-block.
@pytest.mark.parametrize("R,lnc,K,I", [(256, 1, 1024, 4096), (512, 2, 1024, 4096), (1280, 1, 1024, 4096),
                                       (384, 1, 768, 3072)])
def test_gated_mlp_ln_stream(R, lnc, K, I):
    if K == 1024:
        assert not mlp_weights_fit(K, I)
    args, ref = _mlp_case(R, K, I, R + K)
    assert rel_l2(nki.simulate(gated_mlp_ln_stream_kernel[lnc])(*args), ref) < TOL


# ---------------------------------------------------------------- fused attention block
H = 768


def _attn_case(B, L, seed, lengths=None, H=H):
    rng = np.random.default_rng(seed)
    x = rng.standard_normal((B, L, H)).astype(bf16)
    wqkv = (rng.standard_normal((3 * H, H)) * 0.06).astype(bf16).astype(np.float32)
    bqkv = (0.1 * rng.standard_normal(3 * H)).astype(np.float32)
    wo = (rng.standard_normal((H, H)) / np.sqrt(H)).astype(bf16).astype(np.float32)
    bo = (0.1 * rng.standard_normal(H)).astype(np.float32)
    g = (1 + 0.1 * rng.standard_normal(H)).astype(np.float32)
    b = (0.1 * rng.standard_normal(H)).astype(np.float32)
    cos, sin = rope_table(L)
    if lengths is None:                                   # ragged right padding, every row different
        lengths = [max(1, L - 29 * (i + 1)) for i in range(B)]
    mask = np.zeros((B, L), np.float32)
    for i, n in enumerate(lengths):
        mask[i, :n] = 1
    kb = ((1 - mask) * -30000.0).astype(bf16)
    pk = pack(wqkv, bqkv, wo, bo, g, b, cos, sin, H, 1.0 / np.sqrt(64))
    args = [x.reshape(B * L, H), pk["wqk"].astype(bf16), pk["bqk"], pk["wv"].astype(bf16), pk["bv"],
            pk["wo"].astype(bf16), pk["bo"], pk["g"], pk["b"], pk["cs"], kb]
    ref = (x.astype(np.float32), wqkv, bqkv, wo, bo, g, b, cos, sin, mask)
    return args, ref


# (B, L, lnc): packing factor SP = 4 / 2 / 1 at L = 128 / 256 / >=384; B not divisible by SP falls back to
# SP = 1; LNC2 with odd B uses the query split (needs (L/128) % 2 == 0).
ATTN_CASES = [
    (4, 128, 1),    # SP=4
    (2, 128, 1),    # SP=4 does not divide -> SP=1
    (3, 128, 1),    # odd B, SP=1
    (2, 256, 1),    # SP=2
    (1, 384, 1),    # L=384, SP=1
    (1, 512, 1),
    (8, 128, 2),    # LNC2, SP=4 per core
    (4, 256, 2),    # LNC2, SP=2 per core
    (2, 384, 2),    # LNC2, L=384
    (1, 256, 2),    # LNC2 query split, B=1
    (1, 512, 2),    # LNC2 query split, B=1
    (3, 256, 2),    # LNC2 query split, odd B
]


@pytest.mark.parametrize("B,L,lnc", ATTN_CASES)
def test_attn_block(B, L, lnc):
    args, ref = _attn_case(B, L, seed=B * 1000 + L + lnc)
    y = np.asarray(nki.simulate(attn_block_kernel[lnc])(*args, EPS)).astype(np.float32).reshape(B, L, H)
    assert rel_l2(y, attn_block_ref(*ref)) < TOL


# hidden 1024 / 16 heads (gte-large-en-v1.5)
@pytest.mark.parametrize("B,L,lnc", [(4, 128, 1), (2, 256, 1), (1, 512, 1), (4, 256, 2), (3, 256, 2)])
def test_attn_block_h1024(B, L, lnc):
    args, ref = _attn_case(B, L, seed=B * 1000 + L + lnc + 7, H=1024)
    y = np.asarray(nki.simulate(attn_block_kernel[lnc])(*args, EPS)).astype(np.float32).reshape(B, L, 1024)
    assert rel_l2(y, attn_block_ref(*ref)) < TOL


# long sequences (key chunks of 512, grouped LayerNorm epilogue), H 768 and 1024, ragged padding incl. inside the
# last key chunk; LNC2 with odd B uses the query split.
@pytest.mark.parametrize("B,L,lnc,Hm", [(1, 768, 1, 768), (2, 1024, 1, 768), (1, 1024, 2, 768), (3, 768, 2, 768),
                                        (1, 1024, 1, 1024), (2, 768, 2, 1024)])
def test_attn_block_long(B, L, lnc, Hm):
    lengths = [L - 29 * (i + 1) for i in range(B)]
    lengths[0] = L - 600 if L == 1024 else lengths[0]    # padding reaches into the first key chunk
    args, ref = _attn_case(B, L, seed=B * 1000 + L + lnc + Hm, lengths=lengths, H=Hm)
    y = np.asarray(nki.simulate(attn_block_kernel[lnc])(*args, EPS)).astype(np.float32).reshape(B, L, Hm)
    assert rel_l2(y, attn_block_ref(*ref)) < TOL


def test_attn_block_short_sequence():
    """One valid token (all other keys padded) and one full-length sequence in the same pass."""
    args, ref = _attn_case(2, 256, seed=7, lengths=[1, 256])
    y = np.asarray(nki.simulate(attn_block_kernel[1])(*args, EPS)).astype(np.float32).reshape(2, 256, H)
    assert rel_l2(y, attn_block_ref(*ref)) < TOL


@pytest.mark.parametrize("B,L,lnc", [(4, 128, 1), (3, 256, 2)])
def test_attn_block_padding_isolation(B, L, lnc):
    """Sequence 0's output must be bit-identical when every other sequence in the batch changes."""
    args, _ = _attn_case(B, L, seed=11)
    y0 = np.asarray(nki.simulate(attn_block_kernel[lnc])(*args, EPS)).view(np.uint16).reshape(B, L, H)
    rng = np.random.default_rng(12)
    x2 = args[0].reshape(B, L, H).copy()
    x2[1:] = rng.standard_normal(x2[1:].shape).astype(bf16)
    kb2 = args[-1].copy()
    kb2[1:] = 0
    kb2[1:, L // 2:] = -30000.0
    args2 = [x2.reshape(B * L, H)] + args[1:-1] + [kb2]
    y1 = np.asarray(nki.simulate(attn_block_kernel[lnc])(*args2, EPS)).view(np.uint16).reshape(B, L, H)
    assert np.array_equal(y0[0], y1[0])