"""Exact step-by-step spiking simulation of the object pipeline: pixels → layer 1 → readout neuron.

Pixels (3×32×32, value/255 per step, 10 active steps of 15) drive layer-1 neurons: one neuron per learned
prototype k and patch position (27×27), integer weights over its 6×6×3 patch. Patch centring, ZCA and the
prototype fold into one weight vector; the prototype offset and alpha become a per-step bias (the bias of a
neuron that also receives a constant input). A neuron adds its input each step, fires at `thr1`, resets to 0
and never goes below 0 (AltAI). Its spikes go, with delay 1, to the readout neuron "dog − cat" with a weight
shared by prototype k and image quadrant (this replaces pooling neurons, which could fire only once a step).
The readout weights are trained off-chip (ridge) on layer-1 spike counts of the training images.

usage: ~/anvenv/bin/python spiking_pipeline.py --model /home/int/run/cats-dogs/features-v3/kmeans_model_k800.npz
"""
import argparse
import json
import sys
import time
from pathlib import Path

import numpy as np

sys.path.insert(0, "/mnt/c/knp/objects")
sys.path.insert(0, "/mnt/c/knp/readout")
import kmeans_features as kf  # noqa: E402

STEPS, ACTIVE, PATCH, SIDE = 15, 10, 6, 27
WSCALE = 127  # layer-1 weights are integers in ±WSCALE


def fold(model, alpha):
    """Float layer-1 weights [108, K] and bias [K] for raw patches in [0, 1]: act = p @ W + b."""
    n = PATCH * PATCH * 3
    centre = (np.eye(n) - 1 / n) * 4  # kf.normalise with CONTRAST_DIVIDE=False
    w = centre @ model["zca"] @ model["d"].T
    b = -(model["mean"] @ model["zca"] @ model["d"].T) - alpha
    return w.astype(np.float32), b.astype(np.float32)


def quant_layer1(w, b):
    s = WSCALE / np.abs(w).max()
    # Over ACTIVE steps the pixel spikes sum to ~ACTIVE * p, so the bias is added on each active step.
    return np.rint(w * s).astype(np.int32), np.rint(b * s).astype(np.int32), s


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


def layer1_counts(x, wq, bq, thr1, batch=40, per_step=False):
    """Spike counts of layer-1 neurons pooled per (quadrant, k): [n, 4 * K] (or [n, STEPS, 4 * K] with
    per_step); also total spikes per image."""
    pos = range(SIDE)
    edges = [0, 14, 27]
    res, total = [], []
    for i in range(0, len(x), batch):
        r = raster(x[i:i + batch])
        n = r.shape[1]
        pot = np.zeros((n, SIDE * SIDE, wq.shape[1]), np.int32)
        cnt = np.zeros_like(pot, dtype=np.int16)
        steps = []
        for t in range(STEPS):
            if t < ACTIVE:
                p = kf.patches_at(r[t].astype(np.float32), pos, pos).astype(np.int8)
                pot += (p.reshape(-1, p.shape[-1]).astype(np.float32) @ wq.astype(np.float32)).astype(np.int32).reshape(pot.shape) + bq
            fired = pot >= thr1
            cnt += fired
            pot[fired] = 0
            np.maximum(pot, 0, out=pot)
            if per_step:
                steps.append(pool(fired, edges))
        res.append(np.stack(steps, axis=1) if per_step else pool(cnt, edges))
        total.append(cnt.reshape(n, -1).sum(1))
    return np.concatenate(res), np.concatenate(total)


def pool(c, edges):
    c = c.reshape(c.shape[0], SIDE, SIDE, -1)
    return np.concatenate([c[:, edges[gy]:edges[gy + 1], edges[gx]:edges[gx + 1]].sum((1, 2))
                           for gy in range(2) for gx in range(2)], axis=1).astype(np.int32)


def readout_neuron(steps, wq, bq, thr):
    """Exact readout neuron: steps [n, STEPS, F] layer-1 spikes arrive with delay 1; integer weights wq [F] and a
    bias input bq on every active step; fires at thr, resets to 0, never below 0. Returns 1 (dog) if it fired."""
    pot = np.zeros(len(steps), np.int64)
    fired_any = np.zeros(len(steps), bool)
    for t in range(1, STEPS):
        pot += steps[:, t - 1] @ wq + (bq if t - 1 < ACTIVE else 0)
        f = pot >= thr
        fired_any |= f
        pot[f] = 0
        np.maximum(pot, 0, out=pot)
    return fired_any.astype(int)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--data", default="/home/int/run/cats-dogs/data-v3")
    ap.add_argument("--alpha", type=float, default=0.25)
    ap.add_argument("--thr1", type=float, default=0, help="layer-1 threshold in quantised units (0: auto)")
    ap.add_argument("--train-limit", type=int, default=0)
    ap.add_argument("--out", required=True)
    a = ap.parse_args()
    model = dict(np.load(a.model))
    w, b = fold(model, a.alpha)
    wq, bq, s = quant_layer1(w, b)
    tr, va = np.load(f"{a.data}/train.npz"), np.load(f"{a.data}/val.npz")
    xt = tr["rgb"] if not a.train_limit else tr["rgb"][:a.train_limit]
    yt = tr["labels"][:len(xt)].astype(int)
    if not a.thr1:
        # Typical positive per-step drive of a layer-1 neuron on training patches.
        p = kf.patches_at(xt[:200].astype(np.float32) / 255, range(SIDE), range(SIDE)).reshape(-1, 108)
        act = p @ w + b
        a.thr1 = float(np.percentile(act[act > 0], 99) * s)
    print("layer-1 threshold", a.thr1, "weight scale", round(s, 2), flush=True)
    t0 = time.time()
    ct, tt = layer1_counts(np.concatenate([xt, xt[:, :, :, ::-1]]), wq, bq, a.thr1)
    sv, tv = layer1_counts(va["rgb"], wq, bq, a.thr1, per_step=True)
    cv = sv.sum(1)
    print("simulated", len(ct) + len(cv), "images in", round(time.time() - t0), "s; layer-1 spikes per image",
          int(tv.mean()), flush=True)
    yy = np.r_[yt, yt]
    np.savez(a.out, ct=ct, sv=sv, yt=yy, yv=va["labels"], thr1=a.thr1, wq=wq, bq=bq)
    f = lambda c: np.hstack([c.astype(np.float32), np.ones((len(c), 1), np.float32)])
    res = {}
    for lam in [1e5, 1e6, 1e7]:
        sol = np.linalg.solve(f(ct).T @ f(ct) + lam * np.eye(ct.shape[1] + 1), f(ct).T @ (yy * 2 - 1.0))
        acc = 100 * (((f(cv) @ sol) > 0) == va["labels"]).mean()
        rq = np.rint(sol / np.abs(sol).max() * 255).astype(np.int64)
        # The readout bias input fires on each of the ACTIVE steps, so its weight is spread over them.
        bias = int(np.rint(rq[-1] / ACTIVE))
        # Readout threshold relative to the typical size of its per-step input.
        unit = float(np.median(np.abs(sv[:, :ACTIVE].reshape(-1, sv.shape[-1]) @ rq[:-1] + bias))) + 1
        spk = {q: 100 * (readout_neuron(sv, rq[:-1], bias, q * unit) == va["labels"]).mean() for q in (2, 3, 4, 5, 6)}
        res[lam] = {"linear": acc, "spiking": spk, "unit": unit}
        print(f"lam {lam:g}: linear on exact layer-1 counts {acc:.1f}%; spiking readout by threshold "
              + ", ".join(f"{q}×: {v:.1f}%" for q, v in spk.items()) + f" (unit {unit:.0f})", flush=True)
    Path(a.out).with_suffix(".json").write_text(json.dumps({"thr1": a.thr1, "acc": res}, indent=1))


if __name__ == "__main__":
    main()