"""End-to-end CPU check of the torch glue (patch_model, buffer layouts, mask clamp, fallbacks) with the NKI
kernels executed by nki.simulate. torch_neuronx is stubbed; the model is truncated to 2 layers to keep the
simulator time down. Simulator only: NOT device evidence.

Needs torch + transformers + nki + ml_dtypes in one venv.
"""
import os
import sys
import types

import numpy as np
import pytest

torch = pytest.importorskip("torch")
transformers = pytest.importorskip("transformers")
nki = pytest.importorskip("nki")
ml_dtypes = pytest.importorskip("ml_dtypes")
pytestmark = [pytest.mark.sim, pytest.mark.cpu]

MODEL = os.environ.get("GTE_TEST_MODEL", "Alibaba-NLP/gte-multilingual-reranker-base")


class _Sim:
    """kernel[lnc](*torch_args) -> nki.simulate on NumPy copies, result back as a torch tensor."""

    def __init__(self, k):
        self.k = k

    def __getitem__(self, lnc):
        def run(*args):
            np_args = []
            for a in args:
                if isinstance(a, torch.Tensor):
                    a = a.detach().contiguous()
                    if a.dtype == torch.bfloat16:
                        a = a.view(torch.int16).numpy().view(ml_dtypes.bfloat16)
                    else:
                        a = a.numpy()
                np_args.append(a)
            y = np.asarray(nki.simulate(self.k[lnc])(*np_args))
            if y.dtype == ml_dtypes.bfloat16:
                return torch.from_numpy(y.view(np.int16).copy()).view(torch.bfloat16)
            return torch.from_numpy(y.copy())
        return run


@pytest.fixture(scope="module")
def pkg():
    if "torch_neuronx" not in sys.modules or not hasattr(sys.modules["torch_neuronx"], "synchronize"):
        stub = types.ModuleType("torch_neuronx")
        stub.nki_op = lambda name, mutates_args=None: (lambda f: f)
        sys.modules["torch_neuronx"] = stub
    else:
        pytest.skip("real torch_neuronx present; use test_device.py")
    from conftest import load_package
    p = load_package()
    to = sys.modules["gte_encoder_neuron_kernels.torch_ops"]
    for n in ("res_ln_kernel", "gated_mlp_ln_kernel", "attn_block_kernel"):
        setattr(to, n, _Sim(getattr(to, n)))
    return p


def _model(n_layers):
    m = transformers.AutoModelForSequenceClassification.from_pretrained(
        MODEL, trust_remote_code=True, attn_implementation="eager").eval()
    m.new.encoder.layer = m.new.encoder.layer[:n_layers]
    return m


@pytest.mark.parametrize("L,lens", [(128, [128, 37])])
def test_patched_model_matches_fp32(pkg, L, lens):
    tok = transformers.AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
    ref_m = _model(2)
    torch.manual_seed(0)
    B = len(lens)
    ids = torch.randint(1000, 30000, (B, L))
    msk = torch.zeros(B, L, dtype=torch.long)
    for i, n in enumerate(lens):
        msk[i, :n] = 1
    ids[:, 0] = tok.cls_token_id
    with torch.no_grad():
        ref = ref_m(ids, msk).logits.float()
        hs_ref = ref_m.new(ids, msk).last_hidden_state.float()
    m = _model(2)
    for mod in m.modules():
        for cn, c in mod.named_children():
            if "GELU" in type(c).__name__:
                setattr(mod, cn, torch.nn.GELU(approximate="tanh"))
    m = m.to(torch.bfloat16)
    assert pkg.patch_model(m) == 2
    with torch.no_grad():
        hs = m.new(ids, msk).last_hidden_state.float()
        out = m(ids, msk).logits.float()
    valid = msk.bool()
    rel = (hs - hs_ref)[valid].norm() / hs_ref[valid].norm()
    print(f"hidden rel_l2 {rel:.2e}, logits {out.flatten().tolist()} vs fp32 {ref.flatten().tolist()}")
    assert rel < 2e-2
    assert (out - ref).abs().max() < 0.1