"""Pure-numpy inference for both Nexus safetensors models. No torch, no sklearn needed.
Usage:
    from inference import predict
    probs, labels = predict("no-edge-detector", X)     # X: (n_features,) or (batch, n_features)
"""
import json, os, numpy as np
from safetensors.numpy import load_file

def _softmax(z):
    z = z - z.max(axis=-1, keepdims=True); e = np.exp(z)
    return e / e.sum(axis=-1, keepdims=True)

def _sigmoid(z):
    return 1.0 / (1.0 + np.exp(-z))

def predict(model_dir, X):
    cfg = json.load(open(os.path.join(model_dir, "config.json"), encoding="utf-8"))
    T = load_file(os.path.join(model_dir, "model.safetensors"))
    X = np.asarray(X, float)
    if X.ndim == 1: X = X[None, :]
    if "scaler_mean" in T:
        sc = np.where(T["scaler_scale"] == 0, 1.0, T["scaler_scale"])
        X = (X - T["scaler_mean"]) / sc
    if cfg["type"] == "mlp":
        h = X; n = cfg["n_layers"]
        for i in range(n):
            h = h @ T[f"W{i}"] + T[f"b{i}"]
            if i < n - 1: h = np.maximum(h, 0.0)      # relu
        p = _softmax(h)
        labels = [cfg["classes"][i] for i in p.argmax(1)]
        return p, labels
    else:  # logistic
        z = X @ T["coef"].reshape(-1) + float(T["intercept"].reshape(-1)[0])
        return _sigmoid(z), None