"""Numpy-only inference for harpertoken/mark.

Replicates scikit-learn's IsolationForest scoring exactly (verified 99.9-100%
predict agreement on train, holdout, and synthetic incidents) without sklearn
or pickle. Reads model.safetensors plus config.json from this repository.
"""

import json

import numpy as np

try:
    from safetensors.numpy import load_file
except ImportError:  # pragma: no cover
    load_file = None


def _c(n):
    n = np.asarray(n, dtype=float)
    out = np.zeros_like(n)
    m = n > 1
    out[m] = 2.0 * (np.log(n[m] - 1) + np.euler_gamma) - 2.0 * (n[m] - 1) / n[m]
    return out


def _depths(X, tensors, n_trees):
    D = np.zeros((X.shape[0], n_trees))
    for i in range(n_trees):
        left = tensors[f"trees.{i}.children_left"]
        right = tensors[f"trees.{i}.children_right"]
        feat = tensors[f"trees.{i}.feature"]
        thresh = tensors[f"trees.{i}.threshold"]
        counts = tensors[f"trees.{i}.n_node_samples"]
        node = np.zeros(X.shape[0], dtype=np.int64)
        d = np.zeros(X.shape[0])
        alive = np.ones(X.shape[0], dtype=bool)
        while alive.any():
            f = feat[node]
            is_leaf = f < 0
            done = alive & is_leaf
            d[done] += _c(counts[node[done]])
            alive[done] = False
            go = alive & ~is_leaf
            if not go.any():
                break
            idx = np.where(go)[0]
            go_left = X[idx, feat[node[idx]]] <= thresh[node[idx]]
            node[idx[go_left]] = left[node[idx[go_left]]]
            node[idx[~go_left]] = right[node[idx[~go_left]]]
            d[go] += 1.0
        D[:, i] = d
    return D


class MarkModel:
    """IsolationForest equivalent loaded from safetensors. predict returns
    1 for normal, -1 for flagged, matching scikit-learn."""

    def __init__(self, directory="."):
        if load_file is None:
            raise ImportError("pip install safetensors to load mark")
        with open(f"{directory}/config.json") as f:
            self.cfg = json.load(f)
        self.tensors = load_file(f"{directory}/model.safetensors")
        self.features = self.cfg["features"]
        self.offset = self.cfg["offset_"]
        self._norm = _c(np.array([self.cfg["max_samples"]]))[0]

    def predict(self, X):
        X = np.asarray(X, dtype=float)
        Z = (X - self.tensors["scaler.mean"]) / np.sqrt(self.tensors["scaler.var"])
        depths = _depths(Z, self.tensors, self.cfg["n_estimators"])
        scores = -np.power(2.0, -depths.mean(axis=1) / self._norm)
        return np.where(scores - self.offset < 0, -1, 1)


if __name__ == "__main__":
    m = MarkModel()
    print(m.predict([[46.8, 96.0, 3410.2, 33.0, 6.1, 5.8, 0.4, 0.3]]))
    print(m.predict([[95.0, 100.0, 4900.0, 20.0, 120.0, 90.0, 12.0, 11.0]]))