"""Device timing: stock PyTorch Native vs patched, one model on one logical NeuronCore (single process).

  GTE_TEST_MODEL=Alibaba-NLP/gte-large-en-v1.5 NEURON_LOGICAL_NC_CONFIG=2 python tests/bench_model.py 256x32,128x32,512x8
Prints and appends JSON lines (ms per batch median of N_ITERS, seq/s on one logical core) to $GTE_RESULTS.
Both variants use the tanh GELU so the comparison isolates the kernels. Hardware numbers only.
"""
import json
import os
import sys
import time

import numpy as np
import torch
import torch._dynamo as dyn
import torch_neuronx
import transformers

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from conftest import load_package  # noqa: E402

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"))
N_ITERS = int(os.environ.get("N_ITERS", "50"))
SHAPES = [tuple(int(v) for v in s.split("x")) for s in (sys.argv[1] if len(sys.argv) > 1 else "256x32").split(",")]
VARIANTS = (sys.argv[2] if len(sys.argv) > 2 else "stock,nki").split(",")
dyn.config.cache_size_limit = 64
pkg = load_package()
dev = torch.device("neuron")


AUTO = transformers.AutoModelForSequenceClassification if "reranker" in MODEL else transformers.AutoModel


def build(v):
    m = AUTO.from_pretrained(MODEL, trust_remote_code=True, attn_implementation="eager").eval()
    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)
    if v == "nki":
        pkg.patch_model(m)
    elif v == "nki_cte":           # v1.0 behaviour: fused attention block only up to 512, nkilib cte above
        pkg.patch_model(m, attn_lengths=(128, 256, 384, 512))
    return m.to(dev)


tok = transformers.AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
for v in VARIANTS:
    model = build(v)
    for L, bs in SHAPES:
        cm = torch.compile(model, backend="neuron", dynamic=False)
        text = " ".join(["solar wind storage grid battery policy"] * L)
        e = tok([text] * bs, return_tensors="pt", max_length=L, padding="max_length", truncation=True)
        ids, msk = e["input_ids"].to(dev), e["attention_mask"].to(dev)
        with torch.no_grad():
            cm(input_ids=ids, attention_mask=msk)              # compile + warm-up
            torch_neuronx.synchronize()
            for _ in range(5):
                cm(input_ids=ids, attention_mask=msk)
            torch_neuronx.synchronize()
            bar = os.environ.get("BARRIER_DIR")          # multi-process DP: start timing together
            if bar:
                n_w = int(os.environ["BARRIER_N"])
                open(os.path.join(bar, f"ready_{os.environ.get('NEURON_RT_VISIBLE_CORES', os.getpid())}_{L}x{bs}"), "w").close()
                while len([f for f in os.listdir(bar) if f.endswith(f"_{L}x{bs}")]) < n_w:
                    time.sleep(0.05)
            t_start = time.time()
            lat = []
            for _ in range(N_ITERS):
                s = time.perf_counter()
                cm(input_ids=ids, attention_mask=msk)
                torch_neuronx.synchronize()
                lat.append((time.perf_counter() - s) * 1e3)
        wall = time.time() - t_start
        ms = float(np.median(lat))
        rec = dict(model=MODEL, kind="bench", variant=v, lnc=pkg.LNC, L=L, bs=bs, ms_median=round(ms, 3),
                   ms_p90=round(float(np.percentile(lat, 90)), 3), seq_per_s_one_core=round(bs * 1e3 / ms, 1),
                   iters=N_ITERS,
                   t_start=round(t_start, 2), wall_s=round(wall, 2))
        print(json.dumps(rec), flush=True)
        with open(RESULTS, "a") as f:
            f.write(json.dumps(rec) + "\n")
    del model