"""Offline simulator of the quantized AltAI mnist-learn network at inference.

Reproduces KNP inference: raster -> INPUT (delay 1), optional INPUT -> INPUT inhibition (delay 1),
INPUT -> OUTPUT column vote (delay 1), 15 steps per image, potentials carried across images.

usage: python3 altai_sim.py <dump_dir> <images.bin> <labels.target> [--inhib 255] [--thr 40,60,80] [--first N] [--test 10000]
       [--pix 1568]  (ON/OFF input, KNP_MNIST_INPUT_SIZE)
"""
import argparse
import csv
import numpy as np

STEPS, ACTIVE, CLASSES, PIX = 15, 10, 10, 784


def load_dump(d, tag="quant"):
    with open(f"{d}/{tag}_neurons.csv") as f:
        rows = list(csv.DictReader(f))
    thr = np.array([float(r["activation_threshold"]) + float(r["additional_threshold"]) for r in rows])
    w = np.zeros((PIX, len(rows)))
    with open(f"{d}/{tag}_weights.csv") as f:
        next(f)
        for line in f:
            pre, post, wt = line.split(",")
            w[int(pre), int(post)] = float(wt)
    return w, thr


def load_images(path, label_path, first, count):
    imgs = np.fromfile(path, dtype=np.uint8).reshape(-1, PIX)[first:first + count]
    labels = np.array([int(x) for x in open(label_path).read().split()])[first:first + count]
    return imgs, labels


def raster(img, states):
    """Incrementing image-to-spikes converter; `states` persists between images like in KNP."""
    out = np.zeros((STEPS, PIX), dtype=bool)
    inc = img.astype(np.float32) * np.float32(1 / 255.0)
    for t in range(ACTIVE):
        states += inc
        fired = states >= 1.0
        out[t] = fired
        states[fired] -= 1.0
    return out


class Sim:
    def __init__(self, w, thr, inhib=0.0):
        self.w, self.thr, self.inhib = w, thr, inhib
        n = len(thr)
        self.col = n // CLASSES
        self.pot = np.zeros(n)
        self.states = np.zeros(PIX, dtype=np.float32)
        self.prev_raster = np.zeros(PIX, dtype=bool)
        self.prev_input = np.zeros(n, dtype=bool)

    def run_image(self, img, trace=False):
        """Returns (votes per class, per-step trace if requested)."""
        spikes = raster(img, self.states)
        votes = np.zeros(CLASSES, dtype=int)
        steps = []
        for t in range(STEPS):
            # OUTPUT: column fires if any of its INPUT neurons fired on the previous step.
            col_fired = self.prev_input.reshape(CLASSES, self.col).any(axis=1)
            votes += col_fired
            # INPUT: impacts from previous-step raster and inhibition.
            self.pot = np.round(self.pot)
            self.pot += self.prev_raster @ self.w
            if self.inhib and self.prev_input.any():
                k = self.prev_input.sum()
                self.pot -= self.inhib * (k - self.prev_input)  # no self-inhibition
            fired = self.pot >= self.thr
            self.pot[fired] = 0
            self.pot = np.maximum(self.pot, 0)  # negative threshold 0 with saturation
            if trace:
                steps.append({"raster": spikes[t].copy(), "pot": self.pot.copy(), "fired": fired.copy(),
                              "votes": votes.copy()})
            self.prev_raster, self.prev_input = spikes[t], fired
        return votes, steps


def predict(votes):
    return int(np.argmax(votes)) if votes.max() > 0 else -1


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("dump")
    ap.add_argument("images")
    ap.add_argument("labels")
    ap.add_argument("--inhib", type=float, default=0)
    ap.add_argument("--first", type=int, default=None, help="index of first test image (default: last --test)")
    ap.add_argument("--test", type=int, default=10000)
    ap.add_argument("--thr", default="", help="comma-separated thresholds to sweep instead of the dumped ones")
    ap.add_argument("--pix", type=int, default=784, help="input size per image (1568 for ON/OFF)")
    a = ap.parse_args()
    global PIX
    PIX = a.pix
    w, thr = load_dump(a.dump)
    total = len(open(a.labels).read().split())
    first = total - a.test if a.first is None else a.first
    imgs, labels = load_images(a.images, a.labels, first, a.test)
    for t in [float(x) for x in a.thr.split(",")] if a.thr else [None]:
        sim = Sim(w, thr if t is None else np.full_like(thr, t), a.inhib)
        tp = fn = fp = 0
        for img, lab in zip(imgs, labels):
            p = predict(sim.run_image(img)[0])
            if p < 0:
                fn += 1
            elif p == lab:
                tp += 1
            else:
                fp += 1
        tag = "" if t is None else f"thr {t:g}  "
        print(f"{tag}correct {tp}  unanswered {fn}  wrong {fp}  accuracy {tp / len(labels):.4f}")

if __name__ == "__main__":
    main()