"""
core.py  —  AllayRT CPU Execution Engine
==================================================
Implements the op-dispatch table and forward-pass executor
for AllayRT's CPU backend.

Supported ops (ONNX-compatible):
  Arithmetic  : Add, Sub, Mul, Div, Pow, Abs, Neg, Sqrt, Exp, Log, Ceil, Floor
  Activations : Relu, LeakyRelu, Sigmoid, Tanh, Softmax, HardSwish, Gelu, Silu
  Reduction   : ReduceMean, ReduceSum, ReduceMax, ReduceMin, ReduceProd
  Shape ops   : Reshape, Flatten, Squeeze, Unsqueeze, Transpose, Expand
  Tensor ops  : Concat, Split, Slice, Gather, ScatterND, Tile, Pad
  Linear      : Gemm, MatMul
  Conv family : Conv, ConvTranspose, DepthwiseConv (via grouped Conv)
  Pooling     : MaxPool, AveragePool, GlobalMaxPool, GlobalAveragePool
  Norm        : BatchNormalization, LayerNorm, InstanceNorm, GroupNorm
  Logic       : Equal, Greater, Less, And, Or, Not, Where
  Type        : Cast, Constant, Identity, Shape, Size, Range
  Misc        : Clip, Dropout (inference passthrough), Resize, TopK, ArgMax, ArgMin
"""

import numpy as np
from typing import Dict, List, Any, Optional
from math import erf as _erf


# ─────────────────────────────────────────────
#  Execution context
# ─────────────────────────────────────────────

class ExecContext:
    """Holds all live tensors during a forward pass."""

    def __init__(self, weights: Dict[str, np.ndarray]):
        self.tensors: Dict[str, np.ndarray] = dict(weights)

    def get(self, name: str) -> Optional[np.ndarray]:
        return self.tensors.get(name)

    def set(self, name: str, val: np.ndarray):
        self.tensors[name] = val

    def get_req(self, name: str) -> np.ndarray:
        if name not in self.tensors:
            raise KeyError(f"[AllayRT] Tensor '{name}' not found in context")
        return self.tensors[name]


# ─────────────────────────────────────────────
#  Individual op implementations
# ─────────────────────────────────────────────

def _op_add(inputs, attrs, ctx): return [inputs[0] + inputs[1]]
def _op_sub(inputs, attrs, ctx): return [inputs[0] - inputs[1]]
def _op_mul(inputs, attrs, ctx): return [inputs[0] * inputs[1]]
def _op_div(inputs, attrs, ctx): return [inputs[0] / inputs[1]]
def _op_pow(inputs, attrs, ctx): return [np.power(inputs[0], inputs[1])]
def _op_abs(inputs, attrs, ctx): return [np.abs(inputs[0])]
def _op_neg(inputs, attrs, ctx): return [-inputs[0]]
def _op_sqrt(inputs, attrs, ctx): return [np.sqrt(inputs[0])]
def _op_exp(inputs, attrs, ctx): return [np.exp(inputs[0])]
def _op_log(inputs, attrs, ctx): return [np.log(inputs[0])]
def _op_ceil(inputs, attrs, ctx): return [np.ceil(inputs[0])]
def _op_floor(inputs, attrs, ctx): return [np.floor(inputs[0])]
def _op_relu(inputs, attrs, ctx): return [np.maximum(0, inputs[0])]

def _op_leaky_relu(inputs, attrs, ctx):
    alpha = attrs.get("alpha", 0.01)
    x = inputs[0]
    return [np.where(x >= 0, x, alpha * x)]

def _op_sigmoid(inputs, attrs, ctx):
    x = inputs[0]
    return [np.where(x >= 0,
                     1 / (1 + np.exp(-x)),
                     np.exp(x) / (1 + np.exp(x)))]

def _op_tanh(inputs, attrs, ctx): return [np.tanh(inputs[0])]

def _op_softmax(inputs, attrs, ctx):
    axis = attrs.get("axis", -1)
    x = inputs[0]
    x = x - np.max(x, axis=axis, keepdims=True)
    e = np.exp(x)
    return [e / np.sum(e, axis=axis, keepdims=True)]

def _op_hardswish(inputs, attrs, ctx):
    x = inputs[0]
    return [x * np.clip(x + 3, 0, 6) / 6]

def _op_gelu(inputs, attrs, ctx):
    x = inputs[0]
    # Approximate GELU
    return [0.5 * x * (1 + np.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * x**3)))]

def _op_silu(inputs, attrs, ctx):
    x = inputs[0]
    return [x * (np.where(x >= 0, 1/(1+np.exp(-x)), np.exp(x)/(1+np.exp(x))))]

def _op_gemm(inputs, attrs, ctx):
    A, B = inputs[0], inputs[1]
    C = inputs[2] if len(inputs) > 2 else None
    alpha = attrs.get("alpha", 1.0)
    beta  = attrs.get("beta",  1.0)
    transA = attrs.get("transA", 0)
    transB = attrs.get("transB", 0)
    if transA: A = A.T
    if transB: B = B.T
    Y = alpha * (A @ B)
    if C is not None:
        Y = Y + beta * C
    return [Y]


def _op_matmul(inputs, attrs, ctx):
    return [inputs[0] @ inputs[1]]


def _im2col_dilated(Xp, kH, kW, H_out, W_out, strides, dilations):
    """im2col via index grid. Only used for dilation > 1."""
    N, C, _, _ = Xp.shape
    sH, sW = strides; dH, dW = dilations
    i0 = np.arange(kH)*dH; i1 = np.arange(H_out)*sH
    j0 = np.arange(kW)*dW; j1 = np.arange(W_out)*sW
    i = i0[:,None] + i1[None,:]
    j = j0[:,None] + j1[None,:]
    cols = Xp[:,:, i[:,None,:,None], j[None,:,None,:]]
    return cols.reshape(N, C*kH*kW, H_out*W_out)


def _resolve_pads(auto_pad, pads, H, W_in, kH, kW, strides, dilations):
    if auto_pad in ("SAME_UPPER", "SAME_LOWER"):
        pad_h = max((H   - 1)*strides[0] + dilations[0]*(kH-1) + 1 - H,    0)
        pad_w = max((W_in- 1)*strides[1] + dilations[1]*(kW-1) + 1 - W_in, 0)
        if auto_pad == "SAME_UPPER":
            return [pad_h//2, pad_w//2, pad_h-pad_h//2, pad_w-pad_w//2]
        else:
            return [pad_h-pad_h//2, pad_w-pad_w//2, pad_h//2, pad_w//2]
    return pads


def _op_conv(inputs, attrs, ctx):
    """
    Fast Conv — three specialized paths, zero Python loops over spatial dims:

    Path 1: 1x1 pointwise  → reshape + matmul only (no sliding window at all)
    Path 2: depthwise       → stride_tricks windows + per-channel broadcast sum
    Path 3: regular/grouped → stride_tricks windows + batched matmul
    Fallback: dilation>1    → explicit index grid + matmul
    """
    X, W = inputs[0], inputs[1]
    B         = inputs[2] if len(inputs) > 2 else None
    group     = int(attrs.get("group",    1))
    dilations = list(attrs.get("dilations", [1, 1]))
    strides   = list(attrs.get("strides",   [1, 1]))
    pads      = list(attrs.get("pads",      [0, 0, 0, 0]))
    auto_pad  = attrs.get("auto_pad",  "NOTSET")

    N, C, H, Wi = X.shape
    F, C_g, kH, kW = W.shape
    F_g = F // group

    pads = _resolve_pads(auto_pad, pads, H, Wi, kH, kW, strides, dilations)

    Xp    = np.pad(X, ((0,0),(0,0),(pads[0],pads[2]),(pads[1],pads[3])))
    H_out = (Xp.shape[2] - dilations[0]*(kH-1) - 1) // strides[0] + 1
    W_out = (Xp.shape[3] - dilations[1]*(kW-1) - 1) // strides[1] + 1

    # ── Path 1: 1x1 conv, no dilation, no groups ─────────────────────────
    if kH == 1 and kW == 1 and dilations == [1,1] and group == 1:
        if strides == [1, 1]:
            x_flat = X.reshape(N, C, H*Wi).transpose(0, 2, 1)   # (N, HW, C)
        else:
            x_flat = Xp[:,:,::strides[0],::strides[1]]
            x_flat = x_flat.reshape(N, C, H_out*W_out).transpose(0,2,1)
        w_mat = W.reshape(F, C)
        Y = (x_flat @ w_mat.T).transpose(0,2,1).reshape(N, F, H_out, W_out)
        if B is not None: Y += B[None,:,None,None]
        return [Y.astype(X.dtype)]

    # ── Path 2: depthwise (group == C_in, C_g == 1) ──────────────────────
    if group == C and C_g == 1 and dilations == [1,1]:
        shape_st   = (N, C, H_out, W_out, kH, kW)
        strides_st = (Xp.strides[0], Xp.strides[1],
                      Xp.strides[2]*strides[0], Xp.strides[3]*strides[1],
                      Xp.strides[2],             Xp.strides[3])
        wins   = np.lib.stride_tricks.as_strided(Xp, shape=shape_st, strides=strides_st)
        w_flat = W[:,0]                             # (C, kH, kW)
        Y = (wins * w_flat[None,:,None,None,:,:]).sum(axis=(-2,-1))
        if B is not None: Y += B[None,:,None,None]
        return [Y.astype(X.dtype)]

    # ── Path 3: general — stride_tricks + batched matmul ─────────────────
    Y = np.empty((N, F, H_out, W_out), dtype=np.float32)

    if dilations == [1,1]:
        for g in range(group):
            Xg = Xp[:, g*C_g:(g+1)*C_g]
            shape_st   = (N, C_g, H_out, W_out, kH, kW)
            strides_st = (Xg.strides[0], Xg.strides[1],
                          Xg.strides[2]*strides[0], Xg.strides[3]*strides[1],
                          Xg.strides[2],             Xg.strides[3])
            wins  = np.lib.stride_tricks.as_strided(Xg, shape=shape_st, strides=strides_st)
            # transpose before reshape: (N,C_g,Ho,Wo,kH,kW) -> (N,Ho,Wo,C_g,kH,kW)
            cols  = wins.transpose(0,2,3,1,4,5).reshape(N, H_out*W_out, C_g*kH*kW)
            w_mat = W[g*F_g:(g+1)*F_g].reshape(F_g, C_g*kH*kW)
            out   = np.matmul(cols, w_mat.T)        # (N, HoWo, F_g)
            Y[:, g*F_g:(g+1)*F_g] = out.transpose(0,2,1).reshape(N, F_g, H_out, W_out)
    else:
        for g in range(group):
            cols  = _im2col_dilated(Xp[:, g*C_g:(g+1)*C_g],
                                    kH, kW, H_out, W_out, strides, dilations)
            w_mat = W[g*F_g:(g+1)*F_g].reshape(F_g, -1)
            Y[:, g*F_g:(g+1)*F_g] = np.matmul(w_mat[None], cols).reshape(N,F_g,H_out,W_out)

    if B is not None: Y += B[None,:,None,None]
    return [Y.astype(X.dtype)]


def _pool_windows(Xp, kH, kW, H_out, W_out, strides):
    """Zero-copy sliding window view via stride_tricks → (N,C,Ho,Wo,kH,kW)."""
    shape   = (Xp.shape[0], Xp.shape[1], H_out, W_out, kH, kW)
    strides_ = (Xp.strides[0], Xp.strides[1],
                Xp.strides[2]*strides[0], Xp.strides[3]*strides[1],
                Xp.strides[2],            Xp.strides[3])
    return np.lib.stride_tricks.as_strided(Xp, shape=shape, strides=strides_)


def _op_maxpool(inputs, attrs, ctx):
    X = inputs[0]
    kH, kW  = attrs.get("kernel_shape", [2, 2])
    strides = attrs.get("strides", [kH, kW])
    pads    = attrs.get("pads",    [0, 0, 0, 0])
    Xp = np.pad(X, ((0,0),(0,0),(pads[0],pads[2]),(pads[1],pads[3])),
                mode="constant", constant_values=-np.inf)
    H_out = (Xp.shape[2] - kH) // strides[0] + 1
    W_out = (Xp.shape[3] - kW) // strides[1] + 1
    return [_pool_windows(Xp, kH, kW, H_out, W_out, strides).max(axis=(4,5))]


def _op_avgpool(inputs, attrs, ctx):
    X = inputs[0]
    kH, kW  = attrs.get("kernel_shape", [2, 2])
    strides = attrs.get("strides", [kH, kW])
    pads    = attrs.get("pads",    [0, 0, 0, 0])
    Xp = np.pad(X, ((0,0),(0,0),(pads[0],pads[2]),(pads[1],pads[3])))
    H_out = (Xp.shape[2] - kH) // strides[0] + 1
    W_out = (Xp.shape[3] - kW) // strides[1] + 1
    return [_pool_windows(Xp, kH, kW, H_out, W_out, strides).mean(axis=(4,5))]

def _op_global_maxpool(inputs, attrs, ctx):
    X = inputs[0]
    return [X.max(axis=(2,3), keepdims=True)]

def _op_global_avgpool(inputs, attrs, ctx):
    X = inputs[0]
    return [X.mean(axis=(2,3), keepdims=True)]

def _op_batchnorm(inputs, attrs, ctx):
    X, scale, bias, mean, var = inputs[:5]
    eps = attrs.get("epsilon", 1e-5)
    # Inference mode only
    X_norm = (X - mean[np.newaxis,:,np.newaxis,np.newaxis]) / \
             np.sqrt(var[np.newaxis,:,np.newaxis,np.newaxis] + eps)
    Y = scale[np.newaxis,:,np.newaxis,np.newaxis] * X_norm + \
        bias[np.newaxis,:,np.newaxis,np.newaxis]
    return [Y]

def _op_layernorm(inputs, attrs, ctx):
    X = inputs[0]
    scale = inputs[1] if len(inputs) > 1 else None
    bias  = inputs[2] if len(inputs) > 2 else None
    axis  = attrs.get("axis", -1)
    eps   = attrs.get("epsilon", 1e-5)
    mean  = X.mean(axis=axis, keepdims=True)
    var   = X.var(axis=axis, keepdims=True)
    Y = (X - mean) / np.sqrt(var + eps)
    if scale is not None: Y = Y * scale
    if bias  is not None: Y = Y + bias
    return [Y]

def _op_reshape(inputs, attrs, ctx):
    data, shape = inputs[0], inputs[1].astype(np.int64)
    # 0 means keep dim, -1 means infer
    new_shape = []
    for i, s in enumerate(shape):
        if s == 0:
            new_shape.append(data.shape[i])
        else:
            new_shape.append(s)
    return [data.reshape(new_shape)]

def _op_flatten(inputs, attrs, ctx):
    axis = attrs.get("axis", 1)
    x = inputs[0]
    outer = int(np.prod(x.shape[:axis])) if axis > 0 else 1
    inner = int(np.prod(x.shape[axis:]))
    return [x.reshape(outer, inner)]

def _op_transpose(inputs, attrs, ctx):
    perm = attrs.get("perm", None)
    return [inputs[0].transpose(perm)]

def _op_squeeze(inputs, attrs, ctx):
    x = inputs[0]
    axes = attrs.get("axes", None)
    if axes is None and len(inputs) > 1:
        axes = inputs[1].tolist()
    if axes is None:
        return [x.squeeze()]
    axes = [int(a) for a in axes]
    for ax in sorted(axes, reverse=True):
        x = np.squeeze(x, axis=ax)
    return [x]

def _op_unsqueeze(inputs, attrs, ctx):
    x = inputs[0]
    axes = attrs.get("axes", None)
    if axes is None and len(inputs) > 1:
        axes = inputs[1].tolist()
    axes = [int(a) for a in axes]
    for ax in sorted(axes):
        x = np.expand_dims(x, axis=ax)
    return [x]

def _op_concat(inputs, attrs, ctx):
    axis = attrs.get("axis", 0)
    return [np.concatenate(inputs, axis=axis)]

def _op_split(inputs, attrs, ctx):
    x = inputs[0]
    split = attrs.get("split", None)
    if split is None and len(inputs) > 1:
        split = inputs[1].tolist()
    axis = attrs.get("axis", 0)
    if split is not None:
        chunks = []
        pos = 0
        for s in split:
            chunks.append(np.take(x, range(pos, pos+int(s)), axis=axis))
            pos += int(s)
        return chunks
    num = attrs.get("num_outputs", 2)
    return np.split(x, num, axis=axis)

def _op_slice(inputs, attrs, ctx):
    data   = inputs[0]
    starts = inputs[1].tolist() if len(inputs) > 1 else attrs.get("starts", [])
    ends   = inputs[2].tolist() if len(inputs) > 2 else attrs.get("ends",   [])
    axes   = inputs[3].tolist() if len(inputs) > 3 else attrs.get("axes", list(range(len(starts))))
    steps  = inputs[4].tolist() if len(inputs) > 4 else attrs.get("steps", [1]*len(starts))
    slices = [slice(None)] * data.ndim
    for ax, st, en, step in zip(axes, starts, ends, steps):
        slices[int(ax)] = slice(int(st), int(en), int(step))
    return [data[tuple(slices)]]

def _op_gather(inputs, attrs, ctx):
    data    = inputs[0]
    indices = inputs[1].astype(np.int64)
    axis    = attrs.get("axis", 0)
    return [np.take(data, indices, axis=axis)]

def _op_tile(inputs, attrs, ctx):
    return [np.tile(inputs[0], inputs[1].astype(int).tolist())]

def _op_pad(inputs, attrs, ctx):
    data = inputs[0]
    pads = inputs[1].astype(int).tolist() if len(inputs) > 1 else attrs.get("pads", [])
    mode = attrs.get("mode", "constant")
    value = float(inputs[2]) if len(inputs) > 2 else attrs.get("value", 0.0)
    ndim = data.ndim
    pad_width = [(pads[i], pads[i+ndim]) for i in range(ndim)]
    return [np.pad(data, pad_width, mode=mode, constant_values=value)]

def _op_expand(inputs, attrs, ctx):
    shape = inputs[1].astype(int).tolist()
    return [np.broadcast_to(inputs[0], shape).copy()]

def _op_reduce(fn):
    def op(inputs, attrs, ctx):
        axis = attrs.get("axes", attrs.get("axis", None))
        keepdims = bool(attrs.get("keepdims", 1))
        if isinstance(axis, list):
            axis = tuple(axis)
        return [fn(inputs[0], axis=axis, keepdims=keepdims)]
    return op

_op_reduce_mean = _op_reduce(np.mean)
_op_reduce_sum  = _op_reduce(np.sum)
_op_reduce_max  = _op_reduce(np.max)
_op_reduce_min  = _op_reduce(np.min)
_op_reduce_prod = _op_reduce(np.prod)

def _op_clip(inputs, attrs, ctx):
    x = inputs[0]
    lo = float(inputs[1]) if len(inputs) > 1 else attrs.get("min", -np.inf)
    hi = float(inputs[2]) if len(inputs) > 2 else attrs.get("max",  np.inf)
    return [np.clip(x, lo, hi)]

def _op_cast(inputs, attrs, ctx):
    CAST_MAP = {
        1: np.float32, 2: np.uint8,  3: np.int8,  5: np.int16,
        6: np.int32,   7: np.int64,  9: bool,     10: np.float16,
        11: np.float64, 12: np.uint32, 13: np.uint64,
    }
    to = attrs.get("to", 1)
    dtype = CAST_MAP.get(to, np.float32)
    return [inputs[0].astype(dtype)]

def _op_constant(inputs, attrs, ctx):
    v = attrs.get("value")
    if isinstance(v, list):
        return [np.array(v, dtype=np.float32)]
    return [np.array(v, dtype=np.float32)]

def _op_identity(inputs, attrs, ctx):
    return [inputs[0]]

def _op_shape(inputs, attrs, ctx):
    return [np.array(inputs[0].shape, dtype=np.int64)]

def _op_size(inputs, attrs, ctx):
    return [np.array(inputs[0].size, dtype=np.int64)]

def _op_range(inputs, attrs, ctx):
    start, limit, delta = float(inputs[0]), float(inputs[1]), float(inputs[2])
    return [np.arange(start, limit, delta, dtype=inputs[0].dtype)]

def _op_equal(inputs, attrs, ctx):
    return [(inputs[0] == inputs[1]).astype(np.bool_)]

def _op_greater(inputs, attrs, ctx):
    return [(inputs[0] > inputs[1]).astype(np.bool_)]

def _op_less(inputs, attrs, ctx):
    return [(inputs[0] < inputs[1]).astype(np.bool_)]

def _op_where(inputs, attrs, ctx):
    return [np.where(inputs[0].astype(bool), inputs[1], inputs[2])]

def _op_argmax(inputs, attrs, ctx):
    axis = attrs.get("axis", 0)
    keepdims = bool(attrs.get("keepdims", 1))
    r = np.argmax(inputs[0], axis=axis)
    if keepdims: r = np.expand_dims(r, axis=axis)
    return [r.astype(np.int64)]

def _op_argmin(inputs, attrs, ctx):
    axis = attrs.get("axis", 0)
    keepdims = bool(attrs.get("keepdims", 1))
    r = np.argmin(inputs[0], axis=axis)
    if keepdims: r = np.expand_dims(r, axis=axis)
    return [r.astype(np.int64)]

def _op_topk(inputs, attrs, ctx):
    x = inputs[0]
    k = int(inputs[1]) if len(inputs) > 1 else attrs.get("k", 1)
    axis = attrs.get("axis", -1)
    idx = np.argsort(x, axis=axis)[..., -k:][::-1]
    vals = np.take_along_axis(x, idx, axis=axis)
    return [vals, idx.astype(np.int64)]

def _op_dropout(inputs, attrs, ctx):
    return [inputs[0]]  # inference: no dropout

# ── Activations ────────────────────

def _op_elu(inputs, attrs, ctx):
    alpha = attrs.get("alpha", 1.0)
    x = inputs[0]
    return [np.where(x >= 0, x, alpha * (np.exp(x) - 1))]

def _op_selu(inputs, attrs, ctx):
    alpha  = attrs.get("alpha",  1.6732631921768188)
    gamma  = attrs.get("gamma",  1.0507009873554805)
    x = inputs[0]
    return [gamma * np.where(x >= 0, x, alpha * (np.exp(x) - 1))]

def _op_celu(inputs, attrs, ctx):
    alpha = attrs.get("alpha", 1.0)
    x = inputs[0]
    return [np.maximum(0, x) + np.minimum(0, alpha * (np.exp(x / alpha) - 1))]

def _op_prelu(inputs, attrs, ctx):
    x, slope = inputs[0], inputs[1]
    return [np.where(x >= 0, x, slope * x)]

def _op_thresholded_relu(inputs, attrs, ctx):
    alpha = attrs.get("alpha", 1.0)
    x = inputs[0]
    return [np.where(x > alpha, x, 0.0).astype(x.dtype)]

def _op_shrink(inputs, attrs, ctx):
    lambd = attrs.get("lambd", 0.5)
    bias  = attrs.get("bias",  0.0)
    x = inputs[0]
    return [np.where(x < -lambd, x + bias, np.where(x > lambd, x - bias, 0.0))]

def _op_hardsigmoid(inputs, attrs, ctx):
    alpha = attrs.get("alpha", 0.2)
    beta  = attrs.get("beta",  0.5)
    return [np.clip(alpha * inputs[0] + beta, 0, 1)]

def _op_logsoftmax(inputs, attrs, ctx):
    axis = attrs.get("axis", -1)
    x = inputs[0]
    x = x - np.max(x, axis=axis, keepdims=True)
    log_sum = np.log(np.sum(np.exp(x), axis=axis, keepdims=True))
    return [x - log_sum]

def _op_softplus(inputs, attrs, ctx):
    return [np.log1p(np.exp(inputs[0]))]

def _op_softsign(inputs, attrs, ctx):
    x = inputs[0]
    return [x / (1 + np.abs(x))]

def _op_mish(inputs, attrs, ctx):
    x = inputs[0]
    return [x * np.tanh(np.log1p(np.exp(x)))]

# ── Math ───────────────────────────

def _op_erf(inputs, attrs, ctx): return [np.vectorize(_erf)(inputs[0]).astype(inputs[0].dtype)]
def _op_round(inputs, attrs, ctx): return [np.round(inputs[0])]
def _op_sign(inputs, attrs, ctx): return [np.sign(inputs[0])]
def _op_sin(inputs, attrs, ctx): return [np.sin(inputs[0])]
def _op_cos(inputs, attrs, ctx): return [np.cos(inputs[0])]
def _op_tan(inputs, attrs, ctx): return [np.tan(inputs[0])]
def _op_asin(inputs, attrs, ctx): return [np.arcsin(inputs[0])]
def _op_acos(inputs, attrs, ctx): return [np.arccos(inputs[0])]
def _op_atan(inputs, attrs, ctx): return [np.arctan(inputs[0])]
def _op_sinh(inputs, attrs, ctx): return [np.sinh(inputs[0])]
def _op_cosh(inputs, attrs, ctx): return [np.cosh(inputs[0])]
def _op_asinh(inputs, attrs, ctx): return [np.arcsinh(inputs[0])]
def _op_acosh(inputs, attrs, ctx): return [np.arccosh(inputs[0])]
def _op_atanh(inputs, attrs, ctx): return [np.arctanh(inputs[0])]
def _op_isnan(inputs, attrs, ctx): return [np.isnan(inputs[0])]
def _op_isinf(inputs, attrs, ctx): return [np.isinf(inputs[0])]

def _op_cumsum(inputs, attrs, ctx):
    x    = inputs[0]
    axis = int(inputs[1]) if len(inputs) > 1 else attrs.get("axis", 0)
    exclusive = bool(attrs.get("exclusive", 0))
    reverse   = bool(attrs.get("reverse",   0))
    if reverse:
        x = np.flip(x, axis=axis)
    if exclusive:
        result = np.cumsum(x, axis=axis)
        result = np.roll(result, 1, axis=axis)
        idx = [slice(None)] * x.ndim
        idx[axis] = 0
        result[tuple(idx)] = 0
    else:
        result = np.cumsum(x, axis=axis)
    if reverse:
        result = np.flip(result, axis=axis)
    return [result]

# ── Compare (extended) ────────────────────────

def _op_greater_or_equal(inputs, attrs, ctx):
    return [(inputs[0] >= inputs[1]).astype(np.bool_)]

def _op_less_or_equal(inputs, attrs, ctx):
    return [(inputs[0] <= inputs[1]).astype(np.bool_)]

def _op_and(inputs, attrs, ctx):
    return [(inputs[0].astype(bool) & inputs[1].astype(bool))]

def _op_or(inputs, attrs, ctx):
    return [(inputs[0].astype(bool) | inputs[1].astype(bool))]

def _op_not(inputs, attrs, ctx):
    return [~inputs[0].astype(bool)]

def _op_xor(inputs, attrs, ctx):
    return [(inputs[0].astype(bool) ^ inputs[1].astype(bool))]

# ── Conv (extended) ───────────────────────────

def _op_conv_transpose(inputs, attrs, ctx):
    """
    Fast ConvTranspose via col2im (scatter instead of gather).
    No Python loops over spatial dims.
    """
    X, W = inputs[0], inputs[1]
    B              = inputs[2] if len(inputs) > 2 else None
    strides        = attrs.get("strides",        [1, 1])
    pads           = attrs.get("pads",           [0, 0, 0, 0])
    dilations      = attrs.get("dilations",      [1, 1])
    group          = int(attrs.get("group",      1))
    output_padding = attrs.get("output_padding", [0, 0])

    N, C_in, H_in, W_in = X.shape
    C_in_g, C_out_g, kH, kW = W.shape   # W: [C_in, C_out/group, kH, kW]
    C_out = C_out_g * group
    C_in_pg = C_in // group

    H_out = (H_in-1)*strides[0] - pads[0] - pads[2] + dilations[0]*(kH-1) + 1 + output_padding[0]
    W_out = (W_in-1)*strides[1] - pads[1] - pads[3] + dilations[1]*(kW-1) + 1 + output_padding[1]

    # Padded output canvas (we'll accumulate into it, then crop)
    H_pad = H_out + pads[0] + pads[2]
    W_pad = W_out + pads[1] + pads[3]

    Y = np.zeros((N, C_out, H_pad, W_pad), dtype=X.dtype)

    for g in range(group):
        x_g = X[:, g*C_in_pg:(g+1)*C_in_pg]          # (N, C_in_pg, H_in, W_in)
        w_g = W[g*C_in_pg:(g+1)*C_in_pg]              # (C_in_pg, C_out_g, kH, kW)

        # Reshape W: (C_in_pg, C_out_g*kH*kW)
        w_mat = w_g.reshape(C_in_pg, C_out_g * kH * kW)

        # X as (N, C_in_pg, H_in*W_in) → (N, H_in*W_in, C_in_pg)
        x_flat = x_g.reshape(N, C_in_pg, -1).transpose(0, 2, 1)

        # col = (N, H_in*W_in, C_out_g*kH*kW)
        col = np.matmul(x_flat, w_mat)
        col = col.reshape(N, H_in, W_in, C_out_g, kH, kW)

        # Scatter col into Y via index arrays (no Python loops)
        i_src = np.arange(H_in)[:, np.newaxis] * strides[0] + \
                np.arange(kH)[np.newaxis, :] * dilations[0]    # (H_in, kH)
        j_src = np.arange(W_in)[:, np.newaxis] * strides[1] + \
                np.arange(kW)[np.newaxis, :] * dilations[1]    # (W_in, kW)

        # Use np.add.at for scatter-accumulate
        # col: (N, H_in, W_in, C_out_g, kH, kW)
        # Y:   (N, C_out, H_pad, W_pad)
        # Reorder col → (N, C_out_g, H_in, kH, W_in, kW)
        col_t = col.transpose(0, 3, 1, 4, 2, 5)  # (N, C_out_g, H_in, kH, W_in, kW)

        Y_g = Y[:, g*C_out_g:(g+1)*C_out_g]       # view into Y
        # Vectorized scatter using advanced indexing
        np.add.at(Y_g,
                  (slice(None), slice(None),
                   i_src[:, :, np.newaxis, np.newaxis],
                   j_src[np.newaxis, np.newaxis, :, :]),
                  col_t)

    # Crop padding
    Y = Y[:, :, pads[0]:pads[0]+H_out, pads[1]:pads[1]+W_out]

    if B is not None:
        Y += B[np.newaxis, :, np.newaxis, np.newaxis]
    return [Y]

# ── Norm (extended) ───────────────────────────

def _op_instance_norm(inputs, attrs, ctx):
    X, scale, bias = inputs[0], inputs[1], inputs[2]
    eps = attrs.get("epsilon", 1e-5)
    # Per-channel, per-sample normalization  (NCHW)
    axes = tuple(range(2, X.ndim))
    mean = X.mean(axis=axes, keepdims=True)
    var  = X.var(axis=axes, keepdims=True)
    X_norm = (X - mean) / np.sqrt(var + eps)
    shape = (-1,) + (1,) * (X.ndim - 2)
    return [scale.reshape(shape) * X_norm + bias.reshape(shape)]

def _op_group_norm(inputs, attrs, ctx):
    X = inputs[0]
    scale = inputs[1] if len(inputs) > 1 else None
    bias  = inputs[2] if len(inputs) > 2 else None
    num_groups = attrs.get("num_groups", 1)
    eps = attrs.get("epsilon", 1e-5)
    N, C = X.shape[0], X.shape[1]
    spatial = X.shape[2:]
    G = num_groups
    X_r = X.reshape(N, G, C//G, *spatial)
    axes = tuple(range(2, X_r.ndim))
    mean = X_r.mean(axis=axes, keepdims=True)
    var  = X_r.var(axis=axes, keepdims=True)
    X_norm = ((X_r - mean) / np.sqrt(var + eps)).reshape(N, C, *spatial)
    if scale is not None:
        X_norm = X_norm * scale.reshape(1, C, *([1]*len(spatial)))
    if bias is not None:
        X_norm = X_norm + bias.reshape(1, C, *([1]*len(spatial)))
    return [X_norm]

def _op_lp_normalization(inputs, attrs, ctx):
    x = inputs[0]
    axis = attrs.get("axis", -1)
    p    = attrs.get("p", 2)
    norm = np.linalg.norm(x, ord=p, axis=axis, keepdims=True)
    return [x / np.maximum(norm, 1e-12)]

# ── Gather (extended) ─────────────────────────

def _op_gather_elements(inputs, attrs, ctx):
    data    = inputs[0]
    indices = inputs[1].astype(np.int64)
    axis    = attrs.get("axis", 0)
    return [np.take_along_axis(data, indices, axis=axis)]

def _op_gather_nd(inputs, attrs, ctx):
    data    = inputs[0]
    indices = inputs[1].astype(np.int64)
    batch_dims = attrs.get("batch_dims", 0)
    # Flatten batch dims
    idx = tuple(indices[..., i] for i in range(indices.shape[-1]))
    return [data[idx]]

def _op_scatter_nd(inputs, attrs, ctx):
    data    = inputs[0].copy()
    indices = inputs[1].astype(np.int64)
    updates = inputs[2]
    reduction = attrs.get("reduction", "none")
    idx = tuple(indices[..., i] for i in range(indices.shape[-1]))
    if reduction == "add":
        np.add.at(data, idx, updates)
    elif reduction == "mul":
        np.multiply.at(data, idx, updates)
    else:
        data[idx] = updates
    return [data]

def _op_scatter_elements(inputs, attrs, ctx):
    data    = inputs[0].copy()
    indices = inputs[1].astype(np.int64)
    updates = inputs[2]
    axis    = attrs.get("axis", 0)
    reduction = attrs.get("reduction", "none")
    if reduction == "add":
        np.add.at(data, (slice(None),)*axis + (indices,), updates)
    else:
        np.put_along_axis(data, indices, updates, axis=axis)
    return [data]

# ── Tensor utils (extended) ───────────────────

def _op_constant_of_shape(inputs, attrs, ctx):
    shape = inputs[0].astype(int).tolist()
    value = attrs.get("value", 0.0)
    if isinstance(value, list):
        value = value[0] if value else 0.0
    return [np.full(shape, value, dtype=np.float32)]

def _op_nonzero(inputs, attrs, ctx):
    return [np.array(np.nonzero(inputs[0]), dtype=np.int64)]

def _op_onehot(inputs, attrs, ctx):
    indices = inputs[0].astype(np.int64)
    depth   = int(inputs[1])
    values  = inputs[2]
    off_val, on_val = float(values.flat[0]), float(values.flat[1])
    axis = attrs.get("axis", -1)
    result = np.full((*indices.shape, depth), off_val, dtype=values.dtype)
    result[indices == np.arange(depth)] = on_val
    # Simple 1D/2D implementation
    result2 = np.eye(depth, dtype=values.dtype)[indices.flatten()] * (on_val - off_val) + off_val
    result2 = result2.reshape(*indices.shape, depth)
    if axis != -1 and axis != indices.ndim:
        result2 = np.moveaxis(result2, -1, axis)
    return [result2]

def _op_compress(inputs, attrs, ctx):
    inp       = inputs[0]
    condition = inputs[1].astype(bool)
    axis = attrs.get("axis", None)
    if axis is None:
        return [inp.flatten()[condition.flatten()]]
    return [np.compress(condition, inp, axis=axis)]

def _op_trilu(inputs, attrs, ctx):
    x = inputs[0]
    k = int(inputs[1]) if len(inputs) > 1 else attrs.get("k", 0)
    upper = attrs.get("upper", 1)
    if upper:
        return [np.triu(x, k=k)]
    return [np.tril(x, k=k)]

def _op_reverse_sequence(inputs, attrs, ctx):
    x           = inputs[0]
    seq_lens    = inputs[1].astype(int)
    seq_axis    = attrs.get("seq_axis",   1)
    batch_axis  = attrs.get("batch_axis", 0)
    result = x.copy()
    for i, length in enumerate(seq_lens):
        idx_b = [slice(None)] * x.ndim
        idx_b[batch_axis] = i
        slc = tuple(idx_b)
        idx_s = [slice(None)] * x.ndim
        idx_s[batch_axis] = i
        idx_s[seq_axis]   = slice(None, length)
        t = x[tuple(idx_s)].copy()
        # reverse along seq_axis within this batch
        rev_axes = list(range(t.ndim))
        local_seq_axis = seq_axis if seq_axis < batch_axis else seq_axis - 1
        result[tuple(idx_s)] = np.flip(t, axis=local_seq_axis)
    return [result]

def _op_upsample(inputs, attrs, ctx):
    """Legacy Upsample op — delegate to Resize logic."""
    X = inputs[0]
    scales = inputs[1] if len(inputs) > 1 else None
    if scales is None:
        scales = np.array(attrs.get("scales", [1.0, 1.0, 2.0, 2.0]), dtype=np.float32)
    mode = attrs.get("mode", "nearest")
    sc = scales.tolist()
    H_out = int(X.shape[2] * sc[2])
    W_out = int(X.shape[3] * sc[3])
    Y = np.zeros((X.shape[0], X.shape[1], H_out, W_out), dtype=X.dtype)
    for i in range(H_out):
        for j in range(W_out):
            si = int(i * X.shape[2] / H_out)
            sj = int(j * X.shape[3] / W_out)
            Y[:, :, i, j] = X[:, :, si, sj]
    return [Y]

# ── Reduction (extended) ──────────────────────

def _op_reduce_l1(inputs, attrs, ctx):
    axis = attrs.get("axes", attrs.get("axis", None))
    keepdims = bool(attrs.get("keepdims", 1))
    if isinstance(axis, list): axis = tuple(axis)
    return [np.sum(np.abs(inputs[0]), axis=axis, keepdims=keepdims)]

def _op_reduce_l2(inputs, attrs, ctx):
    axis = attrs.get("axes", attrs.get("axis", None))
    keepdims = bool(attrs.get("keepdims", 1))
    if isinstance(axis, list): axis = tuple(axis)
    return [np.sqrt(np.sum(inputs[0]**2, axis=axis, keepdims=keepdims))]

def _op_reduce_log_sum_exp(inputs, attrs, ctx):
    axis = attrs.get("axes", attrs.get("axis", None))
    keepdims = bool(attrs.get("keepdims", 1))
    if isinstance(axis, list): axis = tuple(axis)
    x = inputs[0]
    m = np.max(x, axis=axis, keepdims=True)
    return [m.squeeze(axis=None if keepdims else axis) +
            np.log(np.sum(np.exp(x - m), axis=axis, keepdims=keepdims))]

def _op_reduce_sum_square(inputs, attrs, ctx):
    axis = attrs.get("axes", attrs.get("axis", None))
    keepdims = bool(attrs.get("keepdims", 1))
    if isinstance(axis, list): axis = tuple(axis)
    return [np.sum(inputs[0]**2, axis=axis, keepdims=keepdims)]

def _op_reduce_log_sum(inputs, attrs, ctx):
    axis = attrs.get("axes", attrs.get("axis", None))
    keepdims = bool(attrs.get("keepdims", 1))
    if isinstance(axis, list): axis = tuple(axis)
    return [np.log(np.sum(inputs[0], axis=axis, keepdims=keepdims))]

# ── Quantization ──────────────────────────────

def _op_quantize_linear(inputs, attrs, ctx):
    x = inputs[0]
    scale = float(inputs[1])
    zero_point = int(inputs[2]) if len(inputs) > 2 else 0
    axis = attrs.get("axis", 1)
    saturate = attrs.get("saturate", 1)
    q = np.round(x / scale).astype(np.int32) + zero_point
    # Determine output dtype from zero_point
    zp_arr = inputs[2] if len(inputs) > 2 else np.array(0, dtype=np.uint8)
    if zp_arr.dtype == np.int8:
        q = np.clip(q, -128, 127).astype(np.int8)
    else:
        q = np.clip(q, 0, 255).astype(np.uint8)
    return [q]

def _op_dequantize_linear(inputs, attrs, ctx):
    x = inputs[0].astype(np.float32)
    scale = float(inputs[1])
    zero_point = float(inputs[2]) if len(inputs) > 2 else 0.0
    return [(x - zero_point) * scale]

def _op_dynamic_quantize_linear(inputs, attrs, ctx):
    x = inputs[0]
    x_min = min(0.0, float(x.min()))
    x_max = max(0.0, float(x.max()))
    scale = np.float32((x_max - x_min) / 255.0) if x_max != x_min else np.float32(1.0)
    zero_point = np.clip(np.round(-x_min / scale), 0, 255).astype(np.uint8)
    q = np.clip(np.round(x / scale + zero_point.astype(np.float32)), 0, 255).astype(np.uint8)
    return [q, np.array(scale, dtype=np.float32), zero_point]

def _op_qlinear_matmul(inputs, attrs, ctx):
    a, a_scale, a_zp, b, b_scale, b_zp, y_scale, y_zp = inputs[:8]
    A = (a.astype(np.float32) - float(a_zp)) * float(a_scale)
    B = (b.astype(np.float32) - float(b_zp)) * float(b_scale)
    Y = A @ B
    q = np.clip(np.round(Y / float(y_scale)) + float(y_zp), 0, 255).astype(np.uint8)
    return [q]

# ── Recurrent (LSTM / GRU / RNN) ─────────────

def _op_lstm(inputs, attrs, ctx):
    """Single-direction LSTM (inference, forward only)."""
    X           = inputs[0]   # [seq, batch, input_size]
    W           = inputs[1]   # [1, 4*hidden, input_size]
    R           = inputs[2]   # [1, 4*hidden, hidden_size]
    B           = inputs[3] if len(inputs) > 3 and inputs[3] is not None else None
    # sequence_lens, initial_h, initial_c ignored for simplicity
    initial_h   = inputs[5] if len(inputs) > 5 and inputs[5] is not None else None
    initial_c   = inputs[6] if len(inputs) > 6 and inputs[6] is not None else None

    hidden_size = attrs.get("hidden_size", R.shape[-1])
    direction   = attrs.get("direction", "forward")

    seq_len, batch, input_size = X.shape
    W0 = W[0]  # [4H, input_size]
    R0 = R[0]  # [4H, hidden_size]
    b  = B[0]  if B  is not None else np.zeros(8 * hidden_size, dtype=X.dtype)

    Wb, Rb = b[:4*hidden_size], b[4*hidden_size:]

    h = initial_h[0] if initial_h is not None else np.zeros((batch, hidden_size), dtype=X.dtype)
    c = initial_c[0] if initial_c is not None else np.zeros((batch, hidden_size), dtype=X.dtype)

    H_seq = []
    def _sig(x): return 1/(1+np.exp(-x))

    for t in range(seq_len):
        xt = X[t]
        gates = xt @ W0.T + h @ R0.T + Wb + Rb  # [batch, 4H]
        i_g = _sig(gates[:, :hidden_size])
        f_g = _sig(gates[:, hidden_size:2*hidden_size])
        g_g = np.tanh(gates[:, 2*hidden_size:3*hidden_size])
        o_g = _sig(gates[:, 3*hidden_size:])
        c   = f_g * c + i_g * g_g
        h   = o_g * np.tanh(c)
        H_seq.append(h)

    Y      = np.stack(H_seq)[np.newaxis]         # [1, seq, batch, hidden]
    Y_h    = h[np.newaxis]                        # [1, batch, hidden]
    Y_c    = c[np.newaxis]                        # [1, batch, hidden]
    return [Y, Y_h, Y_c]

def _op_gru(inputs, attrs, ctx):
    """Single-direction GRU (inference, forward only)."""
    X  = inputs[0]   # [seq, batch, input_size]
    W  = inputs[1]   # [1, 3*hidden, input_size]
    R  = inputs[2]   # [1, 3*hidden, hidden_size]
    B  = inputs[3] if len(inputs) > 3 and inputs[3] is not None else None
    initial_h = inputs[5] if len(inputs) > 5 and inputs[5] is not None else None

    hidden_size = attrs.get("hidden_size", R.shape[-1])
    linear_before_reset = attrs.get("linear_before_reset", 0)

    seq_len, batch, _ = X.shape
    W0, R0 = W[0], R[0]
    b = B[0] if B is not None else np.zeros(6*hidden_size, dtype=X.dtype)
    Wb, Rb = b[:3*hidden_size], b[3*hidden_size:]

    h = initial_h[0] if initial_h is not None else np.zeros((batch, hidden_size), dtype=X.dtype)
    H_seq = []
    def _sig(x): return 1/(1+np.exp(-x))

    for t in range(seq_len):
        xt = X[t]
        gates = xt @ W0[:2*hidden_size].T + h @ R0[:2*hidden_size].T + \
                Wb[:2*hidden_size] + Rb[:2*hidden_size]
        z = _sig(gates[:, :hidden_size])
        r = _sig(gates[:, hidden_size:2*hidden_size])
        if linear_before_reset:
            h_hat = np.tanh(xt @ W0[2*hidden_size:].T + Wb[2*hidden_size:] +
                            r * (h @ R0[2*hidden_size:].T + Rb[2*hidden_size:]))
        else:
            h_hat = np.tanh(xt @ W0[2*hidden_size:].T +
                            (r * h) @ R0[2*hidden_size:].T +
                            Wb[2*hidden_size:] + Rb[2*hidden_size:])
        h = (1 - z) * h_hat + z * h
        H_seq.append(h)

    Y   = np.stack(H_seq)[np.newaxis]
    Y_h = h[np.newaxis]
    return [Y, Y_h]

def _op_rnn(inputs, attrs, ctx):
    """Single-direction vanilla RNN."""
    X  = inputs[0]
    W  = inputs[1]
    R  = inputs[2]
    B  = inputs[3] if len(inputs) > 3 and inputs[3] is not None else None
    initial_h = inputs[5] if len(inputs) > 5 and inputs[5] is not None else None

    hidden_size = attrs.get("hidden_size", R.shape[-1])
    activations = attrs.get("activations", ["Tanh"])
    act_fn = {"Tanh": np.tanh, "Relu": lambda x: np.maximum(0,x), "Sigmoid": lambda x: 1/(1+np.exp(-x))}
    fn = act_fn.get(activations[0], np.tanh)

    seq_len, batch, _ = X.shape
    W0, R0 = W[0], R[0]
    b = B[0] if B is not None else np.zeros(2*hidden_size, dtype=X.dtype)
    Wb, Rb = b[:hidden_size], b[hidden_size:]

    h = initial_h[0] if initial_h is not None else np.zeros((batch, hidden_size), dtype=X.dtype)
    H_seq = []

    for t in range(seq_len):
        h = fn(X[t] @ W0.T + h @ R0.T + Wb + Rb)
        H_seq.append(h)

    return [np.stack(H_seq)[np.newaxis], h[np.newaxis]]

# ── Einsum ────────────────────────────────────

def _op_einsum(inputs, attrs, ctx):
    equation = attrs.get("equation", "")
    return [np.einsum(equation, *inputs)]

# ── NMS / Detection ───────────────────────────

def _op_nms(inputs, attrs, ctx):
    """NonMaxSuppression — returns selected box indices."""
    boxes          = inputs[0]   # [batch, num_boxes, 4]
    scores         = inputs[1]   # [batch, num_classes, num_boxes]
    max_output     = int(inputs[2]) if len(inputs) > 2 else attrs.get("max_output_boxes_per_class", 0)
    iou_threshold  = float(inputs[3]) if len(inputs) > 3 else attrs.get("iou_threshold", 0.0)
    score_threshold= float(inputs[4]) if len(inputs) > 4 else attrs.get("score_threshold", float("-inf"))
    center_point_box = attrs.get("center_point_box", 0)

    selected = []
    for batch_idx in range(boxes.shape[0]):
        for class_idx in range(scores.shape[1]):
            sc  = scores[batch_idx, class_idx]
            bx  = boxes[batch_idx]
            idxs = np.where(sc > score_threshold)[0]
            idxs = idxs[np.argsort(-sc[idxs])]
            if max_output > 0:
                idxs = idxs[:max_output]
            keep = []
            for i in idxs:
                b = bx[i]
                suppress = False
                for k in keep:
                    bk = bx[k]
                    # IoU
                    if center_point_box:
                        b_y1,b_x1 = b[0]-b[2]/2, b[1]-b[3]/2
                        b_y2,b_x2 = b[0]+b[2]/2, b[1]+b[3]/2
                        k_y1,k_x1 = bk[0]-bk[2]/2, bk[1]-bk[3]/2
                        k_y2,k_x2 = bk[0]+bk[2]/2, bk[1]+bk[3]/2
                    else:
                        b_y1,b_x1,b_y2,b_x2 = b[0],b[1],b[2],b[3]
                        k_y1,k_x1,k_y2,k_x2 = bk[0],bk[1],bk[2],bk[3]
                    inter_y1 = max(b_y1,k_y1); inter_x1 = max(b_x1,k_x1)
                    inter_y2 = min(b_y2,k_y2); inter_x2 = min(b_x2,k_x2)
                    inter = max(0,inter_y2-inter_y1)*max(0,inter_x2-inter_x1)
                    area_b = (b_y2-b_y1)*(b_x2-b_x1)
                    area_k = (k_y2-k_y1)*(k_x2-k_x1)
                    iou = inter / (area_b + area_k - inter + 1e-10)
                    if iou > iou_threshold:
                        suppress = True; break
                if not suppress:
                    keep.append(i)
            for k in keep:
                selected.append([batch_idx, class_idx, k])

    if not selected:
        return [np.zeros((0, 3), dtype=np.int64)]
    return [np.array(selected, dtype=np.int64)]

def _op_roi_align(inputs, attrs, ctx):
    """RoiAlign — bilinear interpolation feature extraction."""
    X    = inputs[0]   # [batch, C, H, W]
    rois = inputs[1]   # [num_rois, 4] or [num_rois, 5]
    batch_indices = inputs[2].astype(int) if len(inputs) > 2 else np.zeros(rois.shape[0], dtype=int)

    output_height  = attrs.get("output_height",  1)
    output_width   = attrs.get("output_width",   1)
    sampling_ratio = attrs.get("sampling_ratio", 0)
    spatial_scale  = attrs.get("spatial_scale",  1.0)
    mode           = attrs.get("mode",           "avg")

    N, C, H, W = X.shape
    num_rois = rois.shape[0]
    Y = np.zeros((num_rois, C, output_height, output_width), dtype=X.dtype)

    for r in range(num_rois):
        roi = rois[r]
        if roi.shape[0] == 5:
            b_idx = int(roi[0])
            x1,y1,x2,y2 = roi[1]*spatial_scale, roi[2]*spatial_scale, roi[3]*spatial_scale, roi[4]*spatial_scale
        else:
            b_idx = batch_indices[r]
            x1,y1,x2,y2 = roi[0]*spatial_scale, roi[1]*spatial_scale, roi[2]*spatial_scale, roi[3]*spatial_scale

        roi_w = max(x2 - x1, 1.0)
        roi_h = max(y2 - y1, 1.0)
        bin_w = roi_w / output_width
        bin_h = roi_h / output_height

        for ph in range(output_height):
            for pw in range(output_width):
                cy = y1 + (ph + 0.5) * bin_h
                cx = x1 + (pw + 0.5) * bin_w
                # Bilinear sample
                y0, x0 = int(np.floor(cy)), int(np.floor(cx))
                y1i, x1i = y0 + 1, x0 + 1
                dy, dx = cy - y0, cx - x0
                y0c, y1c = np.clip(y0, 0, H-1), np.clip(y1i, 0, H-1)
                x0c, x1c = np.clip(x0, 0, W-1), np.clip(x1i, 0, W-1)
                val = (X[b_idx, :, y0c, x0c] * (1-dy) * (1-dx) +
                       X[b_idx, :, y1c, x0c] * dy     * (1-dx) +
                       X[b_idx, :, y0c, x1c] * (1-dy) * dx +
                       X[b_idx, :, y1c, x1c] * dy     * dx)
                Y[r, :, ph, pw] = val
    return [Y]

# ── Misc extended ─────────────────────────────

def _op_castlike(inputs, attrs, ctx):
    return [inputs[0].astype(inputs[1].dtype)]

def _op_maxunpool(inputs, attrs, ctx):
    X, indices = inputs[0], inputs[1].astype(int)
    output_shape = inputs[2].astype(int).tolist() if len(inputs) > 2 else None
    kernel_shape = attrs.get("kernel_shape", [2, 2])
    strides = attrs.get("strides", kernel_shape)
    pads = attrs.get("pads", [0,0,0,0])

    N, C, H_in, W_in = X.shape
    if output_shape is None:
        H_out = (H_in-1)*strides[0] - pads[0] - pads[2] + kernel_shape[0]
        W_out = (W_in-1)*strides[1] - pads[1] - pads[3] + kernel_shape[1]
        output_shape = [N, C, H_out, W_out]

    Y = np.zeros(output_shape, dtype=X.dtype)
    for n in range(N):
        for c in range(C):
            flat_out = Y[n, c].flatten()
            flat_x   = X[n, c].flatten()
            flat_idx = indices[n, c].flatten()
            np.add.at(flat_out, flat_idx, flat_x)
            Y[n, c] = flat_out.reshape(Y[n, c].shape)
    return [Y]

def _op_lp_pool(inputs, attrs, ctx):
    X = inputs[0]
    p = attrs.get("p", 2)
    kH, kW = attrs.get("kernel_shape", [2,2])
    strides = attrs.get("strides", [1,1])
    pads    = attrs.get("pads",    [0,0,0,0])
    Xp = np.pad(X, ((0,0),(0,0),(pads[0],pads[2]),(pads[1],pads[3])), mode="constant")
    N, C, H, W = Xp.shape
    H_out = (H - kH) // strides[0] + 1
    W_out = (W - kW) // strides[1] + 1
    Y = np.zeros((N, C, H_out, W_out), dtype=X.dtype)
    for i in range(H_out):
        for j in range(W_out):
            ri, rj = i*strides[0], j*strides[1]
            patch = Xp[:, :, ri:ri+kH, rj:rj+kW]
            Y[:, :, i, j] = np.sum(np.abs(patch)**p, axis=(2,3))**(1/p)
    return [Y]

def _op_mean_variance_norm(inputs, attrs, ctx):
    x    = inputs[0]
    axis = attrs.get("axes", [2, 3])
    if isinstance(axis, int): axis = [axis]
    axis = tuple(axis)
    mean = x.mean(axis=axis, keepdims=True)
    std  = x.std(axis=axis, keepdims=True)
    return [(x - mean) / (std + 1e-9)]

def _op_bernoulli(inputs, attrs, ctx):
    p = inputs[0]
    dtype = attrs.get("dtype", None)
    seed  = attrs.get("seed",  None)
    rng = np.random.default_rng(seed)
    result = rng.random(p.shape) < p
    if dtype:
        CAST_MAP = {1: np.float32, 6: np.int32, 7: np.int64, 9: bool, 11: np.float64}
        result = result.astype(CAST_MAP.get(dtype, np.float32))
    return [result.astype(p.dtype)]

def _op_unique(inputs, attrs, ctx):
    x    = inputs[0]
    axis = attrs.get("axis", None)
    sorted_ = bool(attrs.get("sorted", 1))
    if axis is None:
        u, idx, inv, cnt = np.unique(x.flatten(), return_index=True,
                                     return_inverse=True, return_counts=True)
    else:
        u, idx, inv, cnt = np.unique(x, axis=axis, return_index=True,
                                     return_inverse=True, return_counts=True)
    return [u, idx.astype(np.int64), inv.astype(np.int64), cnt.astype(np.int64)]

def _op_loss_nll(inputs, attrs, ctx):
    """NegativeLogLikelihoodLoss."""
    input_   = inputs[0]  # [N, C] log probabilities
    target   = inputs[1].astype(np.int64)  # [N]
    weight   = inputs[2] if len(inputs) > 2 else None
    reduction = attrs.get("reduction", "mean")
    ignore_index = attrs.get("ignore_index", -100)

    N = input_.shape[0]
    loss = np.array([-input_[i, target[i]] for i in range(N)], dtype=input_.dtype)
    if weight is not None:
        loss *= np.array([weight[target[i]] if target[i] != ignore_index else 0 for i in range(N)])
    mask = (target != ignore_index)
    loss[~mask] = 0
    if reduction == "mean":
        return [loss[mask].mean() if mask.any() else np.array(0.0)]
    elif reduction == "sum":
        return [loss.sum()]
    return [loss]

def _op_softmax_crossentropy(inputs, attrs, ctx):
    """SoftmaxCrossEntropyLoss."""
    logits = inputs[0]
    target = inputs[1].astype(np.int64)
    weight = inputs[2] if len(inputs) > 2 else None
    reduction = attrs.get("reduction", "mean")
    ignore_index = attrs.get("ignore_index", -100)

    # log-softmax
    x = logits - np.max(logits, axis=1, keepdims=True)
    log_sm = x - np.log(np.sum(np.exp(x), axis=1, keepdims=True))
    N = logits.shape[0]
    loss = np.array([-log_sm[i, target[i]] for i in range(N)], dtype=logits.dtype)
    if weight is not None:
        loss *= np.array([weight[target[i]] for i in range(N)])
    mask = (target != ignore_index)
    loss[~mask] = 0
    log_prob = log_sm
    if reduction == "mean":
        return [loss[mask].mean() if mask.any() else np.array(0.0), log_prob]
    elif reduction == "sum":
        return [loss.sum(), log_prob]
    return [loss, log_prob]

def _op_resize(inputs, attrs, ctx):
    """Nearest / bilinear resize (4-D NCHW only)."""
    X = inputs[0]
    # inputs: X, roi (optional), scales (optional), sizes (optional)
    roi    = inputs[1] if len(inputs) > 1 else None
    scales = inputs[2] if len(inputs) > 2 and inputs[2].size > 0 else None
    sizes  = inputs[3] if len(inputs) > 3 and inputs[3].size > 0 else None

    mode = attrs.get("mode", "nearest")

    if sizes is not None and sizes.size > 0:
        target = sizes.astype(int).tolist()
        N, C = target[0], target[1]
        H_out, W_out = target[2], target[3]
    elif scales is not None and scales.size > 0:
        sc = scales.tolist()
        N, C = X.shape[0], X.shape[1]
        H_out = int(X.shape[2] * sc[2])
        W_out = int(X.shape[3] * sc[3])
    else:
        return [X]

    Y = np.zeros((X.shape[0], X.shape[1], H_out, W_out), dtype=X.dtype)
    for i in range(H_out):
        for j in range(W_out):
            si = int(i * X.shape[2] / H_out)
            sj = int(j * X.shape[3] / W_out)
            Y[:, :, i, j] = X[:, :, si, sj]
    return [Y]


# ─────────────────────────────────────────────
#  Op dispatch table
# ─────────────────────────────────────────────

OP_TABLE = {
    # Arithmetic
    "Add":        _op_add,
    "Sub":        _op_sub,
    "Mul":        _op_mul,
    "Div":        _op_div,
    "Pow":        _op_pow,
    "Abs":        _op_abs,
    "Neg":        _op_neg,
    "Sqrt":       _op_sqrt,
    "Exp":        _op_exp,
    "Log":        _op_log,
    "Ceil":       _op_ceil,
    "Floor":      _op_floor,
    # Activations
    "Relu":       _op_relu,
    "LeakyRelu":  _op_leaky_relu,
    "Sigmoid":    _op_sigmoid,
    "Tanh":       _op_tanh,
    "Softmax":    _op_softmax,
    "HardSwish":  _op_hardswish,
    "Gelu":       _op_gelu,
    "Silu":       _op_silu,
    # Linear
    "Gemm":       _op_gemm,
    "MatMul":     _op_matmul,
    # Conv
    "Conv":       _op_conv,
    # Pooling
    "MaxPool":         _op_maxpool,
    "AveragePool":     _op_avgpool,
    "GlobalMaxPool":   _op_global_maxpool,
    "GlobalAveragePool": _op_global_avgpool,
    # Norm
    "BatchNormalization": _op_batchnorm,
    "LayerNormalization": _op_layernorm,
    # Shape
    "Reshape":    _op_reshape,
    "Flatten":    _op_flatten,
    "Transpose":  _op_transpose,
    "Squeeze":    _op_squeeze,
    "Unsqueeze":  _op_unsqueeze,
    "Expand":     _op_expand,
    # Tensor
    "Concat":     _op_concat,
    "Split":      _op_split,
    "Slice":      _op_slice,
    "Gather":     _op_gather,
    "Tile":       _op_tile,
    "Pad":        _op_pad,
    # Reduction
    "ReduceMean": _op_reduce_mean,
    "ReduceSum":  _op_reduce_sum,
    "ReduceMax":  _op_reduce_max,
    "ReduceMin":  _op_reduce_min,
    "ReduceProd": _op_reduce_prod,
    # Misc
    "Clip":       _op_clip,
    "Cast":       _op_cast,
    "Constant":   _op_constant,
    "Identity":   _op_identity,
    "Shape":      _op_shape,
    "Size":       _op_size,
    "Range":      _op_range,
    "Equal":      _op_equal,
    "Greater":    _op_greater,
    "Less":       _op_less,
    "Where":      _op_where,
    "ArgMax":     _op_argmax,
    "ArgMin":     _op_argmin,
    "TopK":       _op_topk,
    "Dropout":    _op_dropout,
    "Resize":     _op_resize,

    # Activations extended
    "Elu":               _op_elu,
    "Selu":              _op_selu,
    "Celu":              _op_celu,
    "PRelu":             _op_prelu,
    "ThresholdedRelu":   _op_thresholded_relu,
    "Shrink":            _op_shrink,
    "HardSigmoid":       _op_hardsigmoid,
    "LogSoftmax":        _op_logsoftmax,
    "Softplus":          _op_softplus,
    "Softsign":          _op_softsign,
    "Mish":              _op_mish,

    # Math extended
    "Erf":     _op_erf,
    "Round":   _op_round,
    "Sign":    _op_sign,
    "Sin":     _op_sin,
    "Cos":     _op_cos,
    "Tan":     _op_tan,
    "Asin":    _op_asin,
    "Acos":    _op_acos,
    "Atan":    _op_atan,
    "Sinh":    _op_sinh,
    "Cosh":    _op_cosh,
    "Asinh":   _op_asinh,
    "Acosh":   _op_acosh,
    "Atanh":   _op_atanh,
    "IsNaN":   _op_isnan,
    "IsInf":   _op_isinf,
    "Isinf":   _op_isinf,
    "CumSum":  _op_cumsum,

    # Compare extended
    "GreaterOrEqual": _op_greater_or_equal,
    "LessOrEqual":    _op_less_or_equal,
    "And":            _op_and,
    "Or":             _op_or,
    "Not":            _op_not,
    "Xor":            _op_xor,

    # Conv extended
    "ConvTranspose":  _op_conv_transpose,
    "Upsample":       _op_upsample,

    # Norm extended
    "InstanceNormalization":     _op_instance_norm,
    "GroupNormalization":        _op_group_norm,
    "LpNormalization":           _op_lp_normalization,
    "MeanVarianceNormalization": _op_mean_variance_norm,

    # Gather extended
    "GatherElements":  _op_gather_elements,
    "GatherND":        _op_gather_nd,
    "ScatterND":       _op_scatter_nd,
    "ScatterElements": _op_scatter_elements,

    # Tensor utils extended
    "ConstantOfShape": _op_constant_of_shape,
    "NonZero":         _op_nonzero,
    "OneHot":          _op_onehot,
    "Compress":        _op_compress,
    "Trilu":           _op_trilu,
    "ReverseSequence": _op_reverse_sequence,

    # Reduction extended
    "ReduceL1":        _op_reduce_l1,
    "ReduceL2":        _op_reduce_l2,
    "ReduceLogSumExp": _op_reduce_log_sum_exp,
    "ReduceSumSquare": _op_reduce_sum_square,
    "ReduceLogSum":    _op_reduce_log_sum,

    # Quantization
    "QuantizeLinear":        _op_quantize_linear,
    "DequantizeLinear":      _op_dequantize_linear,
    "DynamicQuantizeLinear": _op_dynamic_quantize_linear,
    "QLinearMatMul":         _op_qlinear_matmul,

    # Recurrent
    "LSTM": _op_lstm,
    "GRU":  _op_gru,
    "RNN":  _op_rnn,

    # Einsum
    "Einsum": _op_einsum,

    # Detection
    "NonMaxSuppression": _op_nms,
    "RoiAlign":          _op_roi_align,

    # Misc extended
    "CastLike":  _op_castlike,
    "MaxUnpool": _op_maxunpool,
    "LpPool":    _op_lp_pool,
    "Bernoulli": _op_bernoulli,
    "Unique":    _op_unique,
    "NegativeLogLikelihoodLoss": _op_loss_nll,
    "SoftmaxCrossEntropyLoss":   _op_softmax_crossentropy,
}


def register_op(name: str, fn):
    """Register a custom op in the AllayRT dispatch table."""
    OP_TABLE[name] = fn
    print(f"[AllayRT] Registered custom op: {name}")


# ─────────────────────────────────────────────
#  Forward pass runner
# ─────────────────────────────────────────────

def run_graph(
    nodes: List[dict],
    tensors: Dict[str, np.ndarray],
    feed: Dict[str, np.ndarray],
    output_names: List[str],
    device: str = "CPU",
    verbose: bool = False,
) -> Dict[str, np.ndarray]:
    """
    Execute the computational graph node-by-node.

    Parameters
    ----------
    nodes        : list of op dicts (from ARTModel.nodes)
    tensors      : weight/initializer dict
    feed         : input tensors {name: np.ndarray}
    output_names : which outputs to return
    device       : "CPU" (future: "GPU", "NPU")
    verbose      : print per-node debug info

    Returns
    -------
    dict {name: np.ndarray}
    """
    ctx = ExecContext(tensors)

    # Inject inputs
    for name, val in feed.items():
        ctx.set(name, val)

    # Node-by-node execution
    for node in nodes:
        op   = node["op_type"]
        ins  = node["inputs"]
        outs = node["outputs"]
        attrs = node.get("attrs", {})

        # Gather input tensors (empty-string inputs → None, skipped)
        input_tensors = []
        for iname in ins:
            if iname == "":
                continue
            t = ctx.get(iname)
            if t is None:
                raise KeyError(
                    f"[AllayRT] Input '{iname}' for op '{op}' "
                    f"(node '{node.get('name','?')}') not found"
                )
            input_tensors.append(t)

        if op not in OP_TABLE:
            raise NotImplementedError(
                f"[AllayRT] Unsupported op: '{op}'. "
                f"Register it with allay_rt_core.register_op('{op}', fn)"
            )

        results = OP_TABLE[op](input_tensors, attrs, ctx)

        for oname, res in zip(outs, results):
            if oname and res is not None:
                ctx.set(oname, np.asarray(res))
                if verbose:
                    t = ctx.get(oname)
                    print(f"  [{op}]  {oname}  shape={t.shape}  dtype={t.dtype}")

    # Collect outputs
    outputs = {}
    for name in output_names:
        t = ctx.get(name)
        if t is None:
            raise KeyError(f"[AllayRT] Output '{name}' not found after execution")
        outputs[name] = t

    return outputs