"""Shared helpers: NumPy fp32 references and import paths.

The raw kernel modules are imported by file (not through the package __init__, which needs torch), so
the simulator tests run in a venv with only nki + numpy + ml_dtypes.
"""
import os
import sys

os.environ.setdefault("NEURON_PLATFORM_TARGET_OVERRIDE", "trn2")

import numpy as np  # noqa: E402

KDIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "build", "torch-neuron")
if KDIR not in sys.path:
    sys.path.insert(0, KDIR)

EPS = 1e-12
D = 64


def layer_norm(s, g, b, eps=EPS):
    return (s - s.mean(-1, keepdims=True)) / np.sqrt(s.var(-1, keepdims=True) + eps) * g + b


def gelu_tanh(v):
    return 0.5 * v * (1 + np.tanh(0.7978845608028654 * (v + 0.044715 * v ** 3)))


def gated_mlp_ln_ref(x, wu, wg, wd, bd, g, b, eps=EPS):
    """wu, wg [I, K]; wd [K, I] (nn.Linear layout)."""
    x = x.astype(np.float32)
    m = (gelu_tanh(x @ wg.T) * (x @ wu.T)) @ wd.T + bd
    return layer_norm(x + m, g, b, eps)


def rope_table(L, base=20000.0):
    inv = 1.0 / (base ** (np.arange(0, D, 2) / D))
    e = np.outer(np.arange(L), inv)
    e = np.concatenate([e, e], -1)
    return np.cos(e), np.sin(e)


def _rot_half(x):
    return np.concatenate([-x[..., D // 2:], x[..., :D // 2]], -1)


def attn_block_ref(x, wqkv, bqkv, wo, bo, g, b, cos, sin, mask, eps=EPS):
    """NewLayer attention half in fp32: qkv_proj -> rotate-half RoPE -> masked softmax -> o_proj -> +x -> LN.
    x [B, L, H], mask [B, L] (1 valid, 0 padded)."""
    B, L, H = x.shape
    NH = H // D
    qkv = x @ wqkv.T + bqkv
    q, k, v = [t.reshape(B, L, NH, D) for t in np.split(qkv, 3, -1)]
    c, s = cos[None, :L, None, :], sin[None, :L, None, :]
    q = q * c + _rot_half(q) * s
    k = k * c + _rot_half(k) * s
    sc = np.einsum("blhd,bmhd->bhlm", q, k) / np.sqrt(D) + np.where(mask > 0, 0, -1e9)[:, 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) @ wo.T + bo
    return layer_norm(x + o, g, b, eps)


def load_package():
    """Import build/torch-neuron as the package `gte_encoder_neuron_kernels` (needs torch + torch_neuronx).
    Equivalent to kernels.get_kernel(...) on a published repo."""
    import importlib.util
    name = "gte_encoder_neuron_kernels"
    if name in sys.modules:
        return sys.modules[name]
    spec = importlib.util.spec_from_file_location(name, os.path.join(KDIR, "__init__.py"),
                                                  submodule_search_locations=[KDIR])
    mod = importlib.util.module_from_spec(spec)
    sys.modules[name] = mod
    spec.loader.exec_module(mod)
    return mod


def rel_l2(y, e, rows=None):
    y = np.asarray(y, np.float32)
    e = np.asarray(e, np.float32)
    if rows is not None:
        y, e = y[rows], e[rows]
    return float(np.linalg.norm(y - e) / np.linalg.norm(e))