"""Fixed generic neural runtime: recurrent signed-add threshold network (task-neutral).

Imports only numpy/json/hashlib. It never imports the task contract, interpreter, renderer, datasets or any
expert/constructor code. Architecture id: rsat-v1.

Symbols: bytes 0..255, BOS=256, EOS=257, SEP=258, PAD=259 (V = 260).
Per-tick input x (270 channels): one-hot symbol (0..259), mode (260 = prefix, 261 = output), one-hot clock
phase (262..269). Every external symbol occupies 8 ticks (phase 0..7) with its symbol and mode held.

State: H bits, reset to 0 for every request. Synchronous update:
    a_i = sum_{j: R[i,j]=+1} s_j - sum_{j: R[i,j]=-1} s_j + sum_{k: I[i,k]=+1} x_k - sum_{k: I[i,k]=-1} x_k
    s'_i = 1 if a_i >= theta_i else 0
Output (after the 8th tick of every OUTPUT-mode symbol):
    score_o = sum_{i: O[o,i]=+1} s_i - sum_{i: O[o,i]=-1} s_i + bias_o ;  emit smallest o with maximal score.
Protocol: feed BOS then the request bytes in PREFIX mode; then feed SEP in OUTPUT mode (ending the prefix and
starting generation); after each output-mode symbol emit one symbol and feed it back (OUTPUT mode).
EOS ends generation. Emitting BOS/SEP/PAD, or more than 2,048 source bytes, is a failure.

Integer bounds (proved, enforced by validate()): connections in {-1,0,+1}; |a_i| <= H + 3 (one symbol, one
mode, one phase channel active) <= 2,883 for H = 2,880; |score| <= H + 32,767 < 2^31. All accumulation is in
int32; thresholds and biases are int16. No multiplication, matrix product, BLAS, einsum or attention is used:
contributions are added/subtracted column by column.
"""
from __future__ import annotations

import hashlib
import json

import numpy as np

V, BOS, EOS, SEP, PAD = 260, 256, 257, 258, 259
MODE_PREFIX, MODE_OUTPUT = 260, 261
PHASE0 = 262
D_IN = 270
TICKS = 8
MAX_OUT_BYTES = 2048
ARCH_ID = "rsat-v1"
KEYS = ("R", "I", "O", "theta", "bias")


def shapes(H):
    return {"R": (H, H), "I": (H, D_IN), "O": (V, H), "theta": (H,), "bias": (V,)}


def n_slots(H):
    return H * H + 531 * H + 260


def validate(arrs, H):
    sh = shapes(H)
    out = {}
    for k in KEYS:
        a = np.asarray(arrs[k])
        if a.shape != sh[k]:
            raise ValueError(f"{k}: shape {a.shape} != {sh[k]}")
        if not np.issubdtype(a.dtype, np.integer):
            raise ValueError(f"{k}: integer dtype required, got {a.dtype}")
        if k in ("R", "I", "O"):
            if a.size and (a.min() < -1 or a.max() > 1):
                raise ValueError(f"{k}: connections must be -1, 0 or +1")
            out[k] = a.astype(np.int8)
        else:
            if a.size and (a.min() < -32768 or a.max() > 32767):
                raise ValueError(f"{k}: must fit signed 16-bit")
            out[k] = a.astype(np.int16)
    return out


def weights_sha256(arrs):
    h = hashlib.sha256()
    for k in KEYS:
        a = np.ascontiguousarray(arrs[k])
        h.update(k.encode() + str(a.dtype).encode() + str(a.shape).encode() + a.tobytes())
    return h.hexdigest()


def save(path, arrs, config: dict):
    meta = dict(config, architecture_id=ARCH_ID, weights_sha256=weights_sha256(arrs))
    blob = np.frombuffer(json.dumps(meta, sort_keys=True).encode(), dtype=np.uint8)
    np.savez_compressed(path, __meta__=blob, **{k: arrs[k] for k in KEYS})
    return meta


def load(path):
    with np.load(path, allow_pickle=False) as z:
        meta = json.loads(bytes(z["__meta__"]).decode())
        arrs = {k: z[k].copy() for k in KEYS}
    arrs = validate(arrs, int(meta["H"]))
    if weights_sha256(arrs) != meta["weights_sha256"]:
        raise ValueError("weights hash mismatch")
    return arrs, meta


class Net:
    """Optimized CPU implementation (exactly equal to `reference_step`; see verify_runtime)."""

    def __init__(self, arrs, H):
        a = validate(arrs, H)
        self.H = H
        self.RT = np.ascontiguousarray(a["R"].T)        # row j = outgoing connections of source unit j
        self.IT = np.ascontiguousarray(a["I"].T)        # row k = connections of input channel k
        self.OT = np.ascontiguousarray(a["O"].T)        # row i = output connections of unit i
        self.theta = a["theta"].astype(np.int32)
        self.bias = a["bias"].astype(np.int32)
        self.reset()

    def reset(self):
        self.s = np.zeros(self.H, dtype=bool)
        self.rec = np.zeros(self.H, dtype=np.int32)    # sum_j R[:, j] s_j, maintained incrementally
        self.ticks = 0

    def step(self, sym, mode, phase):
        a = self.rec + self.IT[sym] + self.IT[mode] + self.IT[PHASE0 + phase]   # int32 + int8 adds
        new = a >= self.theta
        on = np.flatnonzero(new & ~self.s)
        off = np.flatnonzero(self.s & ~new)
        if on.size:
            self.rec += self.RT[on].sum(axis=0, dtype=np.int32)
        if off.size:
            self.rec -= self.RT[off].sum(axis=0, dtype=np.int32)
        self.s = new
        self.ticks += 1

    def feed(self, sym, mode):
        for p in range(TICKS):
            self.step(sym, mode, p)

    def scores(self):
        act = np.flatnonzero(self.s)
        sc = self.bias.copy()
        if act.size:
            sc += self.OT[act].sum(axis=0, dtype=np.int32)
        return sc

    def emit(self):
        return int(np.argmax(self.scores()))   # argmax returns the smallest index among maxima


def generate(net: Net, request: bytes, max_bytes=MAX_OUT_BYTES, detect_cycles=True):
    """Returns dict(status, source_bytes, ticks). status: ok | invalid_symbol | too_long.
    Exact cycle detection: generation is a deterministic function of (state before feeding, fed symbol);
    if that pair repeats, the emitted sequence is periodic forever without EOS, so the outcome (too_long)
    is decided early. The result equals running to the byte limit."""
    net.reset()
    net.feed(BOS, MODE_PREFIX)
    for b in request:
        net.feed(b, MODE_PREFIX)
    cur = SEP
    out = bytearray()
    seen = set()
    while True:
        if detect_cycles:
            key = (cur, np.packbits(net.s).tobytes())
            if key in seen:
                return {"status": "too_long", "cycle_detected": True, "source": bytes(out[:max_bytes]), "ticks": net.ticks}
            seen.add(key)
        net.feed(cur, MODE_OUTPUT)
        y = net.emit()
        if y == EOS:
            return {"status": "ok", "source": bytes(out), "ticks": net.ticks}
        if y >= 256:
            return {"status": "invalid_symbol", "symbol": int(y), "source": bytes(out), "ticks": net.ticks}
        out.append(y)
        if len(out) > max_bytes:
            return {"status": "too_long", "source": bytes(out[:max_bytes]), "ticks": net.ticks}
        cur = y


# ---------------- slow reference (definitional; small networks / short sequences) ----------------
def reference_step(R, I, theta, s, x):
    H = len(s)
    new = [0] * H
    for i in range(H):
        acc = 0
        for j in range(H):
            if R[i][j] == 1 and s[j]:
                acc += 1
            elif R[i][j] == -1 and s[j]:
                acc -= 1
        for k in range(D_IN):
            if I[i][k] == 1 and x[k]:
                acc += 1
            elif I[i][k] == -1 and x[k]:
                acc -= 1
        new[i] = 1 if acc >= theta[i] else 0
    return new


def reference_scores(O, bias, s):
    out = []
    for o in range(V):
        acc = int(bias[o])
        for i in range(len(s)):
            if O[o][i] == 1 and s[i]:
                acc += 1
            elif O[o][i] == -1 and s[i]:
                acc -= 1
        out.append(acc)
    return out


def reference_generate(arrs, H, request: bytes, max_bytes=MAX_OUT_BYTES):
    R, I, O = arrs["R"].tolist(), arrs["I"].tolist(), arrs["O"].tolist()
    theta, bias = arrs["theta"].tolist(), arrs["bias"].tolist()
    s = [0] * H
    ticks = 0

    def feed(sym, mode, s):
        nonlocal ticks
        for p in range(TICKS):
            x = [0] * D_IN
            x[sym] = 1
            x[mode] = 1
            x[PHASE0 + p] = 1
            s = reference_step(R, I, theta, s, x)
            ticks += 1
        return s
    s = feed(BOS, MODE_PREFIX, s)
    for b in request:
        s = feed(b, MODE_PREFIX, s)
    cur, out = SEP, bytearray()
    while True:
        s = feed(cur, MODE_OUTPUT, s)
        sc = reference_scores(O, bias, s)
        y = max(range(V), key=lambda o: (sc[o], -o))
        if y == EOS:
            return {"status": "ok", "source": bytes(out), "ticks": ticks}
        if y >= 256:
            return {"status": "invalid_symbol", "symbol": y, "source": bytes(out), "ticks": ticks}
        out.append(y)
        if len(out) > max_bytes:
            return {"status": "too_long", "source": bytes(out[:max_bytes]), "ticks": ticks}
        cur = y