"""Numpy-only neural inference; no semantic phrase rules. Fixed feature hashing."""

from importlib.resources import files

from core import Diagnostic, assemble, compile_js, frontend

# isort: split
# core must set CPU thread limits before NumPy initializes.
import numpy as np


def default_checkpoint():
    return files("tiny_codemode_weights").joinpath("model.npz")


class Model:
    def __init__(self, path=None):
        with np.load(path if path is not None else default_checkpoint()) as z:
            self.w = {k: z[k] for k in z.files}

    def logits(self, item):
        _, ids, f, _ = frontend(item)
        return self.forward(ids, f)

    def forward(self, ids, f):
        w = self.w
        dims = w["convs.0.weight"].shape[1] - w["embed.weight"].shape[1]
        h = np.concatenate(
            [w["embed.weight"][ids % len(w["embed.weight"])], f[:, :dims]], axis=1
        )
        for i in range(3):
            weight = w[f"convs.{i}.weight"]
            bias = w[f"convs.{i}.bias"]
            pad = weight.shape[2] // 2
            p = np.pad(h, ((pad, pad), (0, 0)))
            h = np.maximum(
                sum(
                    p[k : k + len(h)] @ weight[:, :, k].T
                    for k in range(weight.shape[2])
                )
                + bias,
                0,
            )
        pooled = np.concatenate([h.mean(0), h.max(0)])
        heads = (
            h @ w["roles.weight"].T + w["roles.bias"],
            pooled @ w["mode.weight"].T + w["mode.bias"],
            pooled @ w["output.weight"].T + w["output.bias"],
        )
        return (
            heads + (pooled @ w["reject.weight"].T + w["reject.bias"],)
            if "reject.weight" in w
            else heads
        )

    def predict(self, item):
        heads = self.logits(item)
        if len(heads) == 4 and int(heads[3].argmax()) == 1:
            raise Diagnostic("learned rejection: unsupported/invalid request")
        r, m, o = heads[:3]
        return assemble(item, r.argmax(1), int(m.argmax()), int(o.argmax()))

    def program(self, item):
        return compile_js(self.predict(item), item["tools"])