"""Train the class neurons offline on detector spike counts produced by KNP (altai_infer), write integer weights.

Ridge regression (one-vs-rest, targets ±1) on per-(quadrant, detector type) spike counts of the training images
(+ mirrored copies). Standardisation is folded into the weights and bias; weights are scaled to integers in ±255,
the bias is spread over the active steps. The class-neuron threshold is chosen on the training counts with a rate
approximation (spikes ≈ positive drive / threshold) so that the most-spiking class matches the ridge decision.
usage: ~/anvenv/bin/python train_readout.py <model_dir> <counts_dir> [--lam 1000]
"""
import argparse
import re
from pathlib import Path

import numpy as np

ACTIVE = 10
ap = argparse.ArgumentParser()
ap.add_argument("model")
ap.add_argument("counts")
ap.add_argument("--lam", type=float, default=1000)
ap.add_argument("--data", default=str(Path.home() / "run/cifar/vehicles"))
a = ap.parse_args()
model, cdir = Path(a.model), Path(a.counts)
k, stride, det_thr, classes, _ = (int(v) for v in (model / "meta.txt").read_text().split())


def load(tag, n):
    pattern = re.compile(re.escape(tag) + r"_(\d+)_counts\.bin")
    parts = sorted((p for p in cdir.glob('*_counts.bin') if pattern.fullmatch(p.name)),
                   key=lambda p: int(pattern.fullmatch(p.name).group(1)))
    arrays, next_index = [], 0
    for path in parts:
        first = int(pattern.fullmatch(path.name).group(1))
        if first != next_index or not (cdir / f'{tag}_{first}.done').exists():
            raise ValueError(f'{tag}: missing checkpoint or noncontiguous range at {first}')
        raw = np.fromfile(path, np.int32)
        if len(raw) % (4 * k):
            raise ValueError(f'Truncated counts: {path}')
        block = raw.reshape(-1, 4 * k)
        arrays.append(block)
        next_index += len(block)
    if next_index != n:
        raise ValueError(f'{tag}: expected {n} rows, got {next_index}')
    return np.concatenate(arrays).astype(np.float32)


yt = np.array(open(f"{a.data}/train_labels.txt").read().split(), int)
yv = np.array(open(f"{a.data}/val_labels.txt").read().split(), int)
xt = np.vstack([load("train", len(yt)), load("train_flip", len(yt))])
yt2 = np.r_[yt, yt]
xv = load("val", len(yv))
mu, sd = xt.mean(0), xt.std(0) + 1e-3
f = lambda z: np.hstack([(z - mu) / sd, np.ones((len(z), 1), np.float32)])
T = np.eye(classes)[yt2] * 2 - 1
for lam in sorted({a.lam, 100.0, 10000.0}):
    sol = np.linalg.solve(f(xt).T @ f(xt) + lam * np.eye(4 * k + 1), f(xt).T @ T)
    print(f"lam {lam:g}: ridge on KNP counts, validation {100 * ((f(xv) @ sol).argmax(1) == yv).mean():.2f}%")
sol = np.linalg.solve(f(xt).T @ f(xt) + a.lam * np.eye(4 * k + 1), f(xt).T @ T)
w = sol[:-1] / sd[:, None]  # on raw counts
b = sol[-1] - (mu / sd) @ sol[:-1]
s = 255 / np.abs(w).max()
wq = np.rint(w * s).astype(np.int32)  # [4K, classes]
bq_total = b * s
bq = np.rint(bq_total / ACTIVE).astype(np.int32)
drive_t = xt @ wq + bq * ACTIVE  # total drive per image and class
best = None
# A class neuron fires at most once a step (about ACTIVE + 1 spikes per image), so the threshold is a fraction
# of the typical positive per-step drive; spikes are capped accordingly.
unit = float(np.percentile(drive_t[drive_t > 0], 50)) / ACTIVE
for q in (0.1, 0.2, 0.3, 0.5, 0.8):
    thr = int(min(max(unit * q, 1), 65535))
    spikes = np.minimum(np.floor(np.maximum(drive_t, 0) / thr), ACTIVE + 1)
    acc = 100 * (spikes.argmax(1) == yt2).mean()
    print(f"threshold {thr}: train accuracy (rate approximation) {acc:.2f}%, max spikes {spikes.max():.0f}")
    if best is None or acc > best[0]:
        best = (acc, thr)
wq.T.copy().tofile(model / "readout_w.bin")
bq.tofile(model / "readout_b.bin")
(model / "meta.txt").write_text(f"{k} {stride} {det_thr} {classes} {best[1]}\n")
print("written", model, "threshold", best[1])