"""On-device accuracy for NewBERT embedding models (gte-*-en-v1.5, gte-multilingual-base): CLS embeddings of the
patched model vs the stock Hugging Face model on CPU in fp32 (exact GELU), and vs stock PyTorch Native.

Gates: mean per-text cosine >= min(0.9999, stock bf16 mean) (patched vs CPU fp32; stock itself is below 0.9999 at L=1024); min per-text cosine within COS_MIN_SLACK of stock
bf16 PyTorch Native's min (an absolute 0.9999 floor on the minimum is not met by stock bf16 itself);
retrieval top-1 per query equal to the CPU fp32 ranking at least as often as stock Native; padding isolation
bit-exact; fused ops actually traced.

  GTE_TEST_MODEL=Alibaba-NLP/gte-large-en-v1.5 NEURON_LOGICAL_NC_CONFIG=2 python -m pytest -m neuron -s tests/test_device_embed.py
"""
import json
import os
import sys

import pytest

from conftest import load_package
from test_device import TOPICS, _pairs

torch = pytest.importorskip("torch")
pytest.importorskip("torch_neuronx")
transformers = pytest.importorskip("transformers")
pytestmark = pytest.mark.neuron

MODEL = os.environ.get("GTE_TEST_MODEL", "Alibaba-NLP/gte-large-en-v1.5")
RESULTS = os.environ.get("GTE_RESULTS", os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                                     "device_results.jsonl"))
SHAPES = [(256, 32), (128, 32), (512, 3), (1024, 8), (768, 3)]
COS_MEAN_MIN = 0.9999
ATTN_LENGTHS = tuple(int(v) for v in os.environ.get("GTE_ATTN_LENGTHS", "128,256,384,512,768,1024").split(","))
COS_MIN_SLACK = 1e-3     # worst single text may trail stock's worst by this much (bf16 noise over 24 layers)


def _load(dtype):
    m = transformers.AutoModel.from_pretrained(MODEL, trust_remote_code=True, attn_implementation="eager").eval()
    if dtype is not None:
        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(dtype)
    return m


class _CLS(torch.nn.Module):
    def __init__(self, m):
        super().__init__()
        self.m = m

    def forward(self, ids, msk):
        return self.m(input_ids=ids, attention_mask=msk).last_hidden_state[:, 0]


@pytest.fixture(scope="module")
def setup():
    import torch._dynamo as dyn
    dyn.config.cache_size_limit = 64
    pkg = load_package()
    to = sys.modules["gte_encoder_neuron_kernels.torch_ops"]
    calls = {}
    for name in ("attn_block", "gated_mlp_ln", "res_ln", "ln"):
        def wrap(f, name=name):
            def g(*a, **k):
                calls[name] = calls.get(name, 0) + 1
                return f(*a, **k)
            return g
        setattr(to, name, wrap(getattr(to, name)))
    tok = transformers.AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
    cpu = _CLS(_load(None))
    dev = torch.device("neuron")
    stock = _CLS(_load(torch.bfloat16)).to(dev)
    p = _load(torch.bfloat16)
    n = pkg.patch_model(p, attn_lengths=ATTN_LENGTHS)
    patched = _CLS(p).to(dev)
    return dict(pkg=pkg, tok=tok, cpu=cpu, dev=dev, stock=stock, patched=patched, n_layers=n, calls=calls)


def _emb(model, ids, msk, bs, dev):
    cm = torch.compile(model, backend="neuron", dynamic=False)
    with torch.no_grad():
        out = [cm(ids[i:i + bs].to(dev), msk[i:i + bs].to(dev)).float().cpu() for i in range(0, ids.shape[0], bs)]
    return torch.cat(out), cm


def _cos(a, b):
    return torch.nn.functional.cosine_similarity(a, b, dim=-1)


@pytest.mark.parametrize("L,bs", SHAPES)
def test_embed_accuracy_and_isolation(setup, L, bs):
    tok, cpu, dev = setup["tok"], setup["cpu"], setup["dev"]
    per = 8
    Q, Ps = _pairs(L, per)
    queries = list(TOPICS)
    texts = queries + Ps                          # 8 queries + 64 passages
    n = len(texts) - len(texts) % bs
    texts = texts[:n]
    e = tok(texts, return_tensors="pt", max_length=L, padding="max_length", truncation=True)
    ids, msk = e["input_ids"], e["attention_mask"]
    with torch.no_grad():
        ref = torch.cat([cpu(ids[i:i + 8], msk[i:i + 8]).float() for i in range(0, n, 8)])
    e_stock, _ = _emb(setup["stock"], ids, msk, bs, dev)
    setup["calls"].clear()
    e_nki, cm = _emb(setup["patched"], ids, msk, bs, dev)
    traced = dict(setup["calls"])

    cos_nki = _cos(e_nki, ref)
    cos_stock = _cos(e_stock, ref)

    def top1(E):
        En = torch.nn.functional.normalize(E, dim=-1)
        q = En[:len(queries)]
        p = En[len(queries):]
        return (q @ p.T).argmax(-1)

    t_ref, t_stock, t_nki = top1(ref), top1(e_stock), top1(e_nki)
    agree_stock = float((t_stock == t_ref).float().mean())
    agree_nki = float((t_nki == t_ref).float().mean())

    b1 = ids[:bs].clone(), msk[:bs].clone()
    if bs > 1:
        perm = torch.arange(n - 1, n - bs, -1)
        b1[0][1:], b1[1][1:] = ids[perm], msk[perm]
    with torch.no_grad():
        y0 = cm(ids[:bs].to(dev), msk[:bs].to(dev))[0].float().cpu()
        y1 = cm(b1[0].to(dev), b1[1].to(dev))[0].float().cpu()
    isolated = bool(torch.equal(y0, y1))
    lens = msk.sum(-1).numpy()
    worst = torch.argsort(cos_nki)[:4].tolist()
    rec_worst = [dict(i=int(i), len=int(lens[i]), cos_nki=float(cos_nki[i]), cos_stock=float(cos_stock[i]))
                 for i in worst]
    rec = dict(model=MODEL, kind="embed", lnc=setup["pkg"].LNC, L=L, bs=bs, texts=n,
               real_len=[int(lens.min()), int(lens.max())], layers_patched=setup["n_layers"],
               cos_min_stock=float(cos_stock.min()), cos_mean_stock=float(cos_stock.mean()),
               cos_min_nki=float(cos_nki.min()), cos_mean_nki=float(cos_nki.mean()),
               top1_agree_stock=agree_stock, top1_agree_nki=agree_nki, padding_isolated=isolated,
               traced_ops=traced, worst=rec_worst)
    print(json.dumps(rec), flush=True)
    with open(RESULTS, "a") as f:
        f.write(json.dumps(rec) + "\n")
    nl = setup["n_layers"]
    if L in ATTN_LENGTHS:
        assert traced.get("attn_block", 0) >= nl, rec
    assert traced.get("gated_mlp_ln", 0) >= nl, rec
    assert float(cos_nki.mean()) >= min(COS_MEAN_MIN, float(cos_stock.mean())), rec
    assert float(cos_nki.min()) >= float(cos_stock.min()) - COS_MIN_SLACK, rec
    assert agree_nki >= agree_stock, rec
    assert isolated, rec