"""Spiking readout: weights trained off-chip (ridge regression), inference by AltAI-like integer neurons.

Input neurons fire with the same incrementing raster as mnist-learn (value/255 per step, 10 active steps
of 15). One output neuron per class integrates integer weights of the spikes it receives, fires at
`thr` and resets to 0; the potential is clipped at 0 like the AltAI negative threshold. A bias input
fires on every active step. The answer is the class whose neuron fired most (earliest on ties).

usage: python spiking_readout.py <data.bin> <labels> --width 784 --train 30000 --eval-first 30000 --eval 10000
                                  [--save weights.npz] [--thr 50,100,200]
"""
import argparse

import numpy as np

STEPS, ACTIVE = 15, 10


def spike_counts(x):
    """x: [n, width] uint8 → spikes per input over ACTIVE steps (raster from rest)."""
    st = np.zeros(x.shape, dtype=np.float32)
    inc = x.astype(np.float32) * np.float32(1 / 255.0)
    cnt = np.zeros(x.shape, dtype=np.int16)
    for _ in range(ACTIVE):
        st += inc
        f = st >= 1.0
        cnt += f
        st[f] -= 1.0
    return cnt


def raster(x):
    """x: [n, width] uint8 → [ACTIVE, n, width] bool."""
    st = np.zeros(x.shape, dtype=np.float32)
    inc = x.astype(np.float32) * np.float32(1 / 255.0)
    out = np.zeros((ACTIVE,) + x.shape, dtype=bool)
    for t in range(ACTIVE):
        st += inc
        out[t] = st >= 1.0
        st[out[t]] -= 1.0
    return out


def train_ridge(x, y, classes, lam):
    f = np.hstack([spike_counts(x).astype(np.float32), np.full((len(x), 1), ACTIVE, np.float32)])
    t = np.eye(classes, dtype=np.float32)[y] * 2 - 1
    return np.linalg.solve(f.T @ f + lam * np.eye(f.shape[1], dtype=np.float32), f.T @ t)


def quantize(w, wmax=255):
    scale = wmax / np.abs(w).max()
    return np.round(w * scale).astype(np.int32), scale


def spiking_predict(x, wq, thr, batch=2000):
    preds = []
    for i in range(0, len(x), batch):
        r = raster(x[i:i + batch])
        n = r.shape[1]
        pot = np.zeros((n, wq.shape[1]), dtype=np.int64)
        count = np.zeros_like(pot)
        first = np.full(pot.shape, STEPS, dtype=np.int64)
        bias = wq[-1]
        for t in range(STEPS):
            if t < ACTIVE:
                pot += r[t].astype(np.int64) @ wq[:-1] + bias
            fired = pot >= thr
            count += fired
            first = np.where(fired & (first == STEPS), t, first)
            pot[fired] = 0
            pot = np.maximum(pot, 0)
        # most spikes; ties broken by earliest first spike; no spikes → -1
        score = count * (STEPS + 1) - first
        p = score.argmax(1)
        p[count.max(1) == 0] = -1
        preds.append(p)
    return np.concatenate(preds)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("data")
    ap.add_argument("labels")
    ap.add_argument("--width", type=int, default=784)
    ap.add_argument("--classes", type=int, default=10)
    ap.add_argument("--train", type=int, required=True)
    ap.add_argument("--eval-first", type=int, required=True)
    ap.add_argument("--eval", type=int, required=True)
    ap.add_argument("--lam", type=float, default=10.0)
    ap.add_argument("--thr", default="")
    ap.add_argument("--save")
    a = ap.parse_args()
    x = np.fromfile(a.data, dtype=np.uint8).reshape(-1, a.width)
    y = np.array([int(v) for v in open(a.labels).read().split()])
    w = train_ridge(x[:a.train], y[:a.train], a.classes, a.lam)
    xe, ye = x[a.eval_first:a.eval_first + a.eval], y[a.eval_first:a.eval_first + a.eval]
    fe = np.hstack([spike_counts(xe).astype(np.float32), np.full((len(xe), 1), ACTIVE, np.float32)])
    print(f"float linear on spike counts: {100 * ((fe @ w).argmax(1) == ye).mean():.2f}%")
    wq, scale = quantize(w)
    ref = np.abs(wq).max()
    thrs = [float(t) for t in a.thr.split(",")] if a.thr else [ref * k for k in (0.5, 1, 2, 4, 8)]
    best = None
    for thr in thrs:
        p = spiking_predict(xe, wq, thr)
        acc = 100 * (p == ye).mean()
        print(f"spiking, thr {thr:g}: correct {acc:.2f}%  silent {100 * (p < 0).mean():.2f}%")
        if best is None or acc > best[0]:
            best = (acc, thr)
    if a.save:
        np.savez(a.save, w=wq, thr=best[1], scale=scale, width=a.width)
        print("saved", a.save, "thr", best[1])


if __name__ == "__main__":
    main()