"""On-line three-factor learning of the readout neuron (dog − cat) on exact layer-1 spike counts.

For each training image: drive = Σ w·count + b; the neuron answers dog if drive > 0. A signed dopamine
signal carries the error e = target − tanh(drive / scale) and every synapse changes by η · e · count
(pre-synaptic activity × dopamine): a local rule that needs no weight transport. Weights are clipped to
integers in ±255 after each epoch for the spiking check.

usage: ~/anvenv/bin/python dopamine_readout.py /home/int/run/cats-dogs/spiking_v3_k800.npz
"""
import sys

import numpy as np

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

s = np.load(sys.argv[1])
ct, sv, yt, yv = s["ct"].astype(np.float32), s["sv"], s["yt"] * 2 - 1, s["yv"]
gain = ct.mean(0) + 1  # per-input normalisation = fixed synaptic gain
x = ct / gain
xv = sv.sum(1) / gain
rng = np.random.default_rng(0)
eta0 = float(sys.argv[2]) if len(sys.argv) > 2 else 2e-4
w = np.zeros(x.shape[1], np.float32)
b = 0.0
w_avg, b_avg, n_avg = np.zeros_like(w), 0.0, 0
for epoch in range(20):
    eta = eta0 / (1 + epoch / 5)
    for i in rng.permutation(len(x)):
        d = x[i] @ w + b
        e = yt[i] - np.tanh(d)
        w += eta * e * x[i]
        b += eta * e
        # Running average of the weights (a slow synaptic trace); used for the answer.
        w_avg += w
        b_avg += b
        n_avg += 1
    acc = 100 * (((xv @ (w_avg / n_avg) + b_avg / n_avg) > 0) == yv).mean()
    if epoch % 4 == 3:
        print(f"epoch {epoch + 1}: validation (linear on spike counts) {acc:.1f}%", flush=True)
w, b = w_avg / n_avg, b_avg / n_avg
# Spiking check: fold the gain into integer weights of the readout neuron.
wf = np.r_[w / gain, b]
# Rarely active inputs get large weights; clip at the 99.5th percentile of |w| so that 8-bit-like integer
# weights keep the bulk of the synapses (the bias keeps its full value relative to the clip).
cap = np.percentile(np.abs(wf[:-1]), 99.5)
wf[:-1] = np.clip(wf[:-1], -cap, cap)
rq = np.rint(wf / cap * 255).astype(np.int64)
bias = int(np.rint(rq[-1] / sp.ACTIVE))
unit = float(np.median(np.abs(sv[:, :sp.ACTIVE].reshape(-1, sv.shape[-1]) @ rq[:-1] + bias))) + 1
print("spiking readout:", {q: round(100 * (sp.readout_neuron(sv, rq[:-1], bias, q * unit) == yv).mean(), 1) for q in (2, 3, 4, 5, 6)})


def calibrated(sv, rq, bias, k=50):
    """Decision neuron: threshold k × unit, bias raised by thr / ACTIVE per active step, so that reaching the
    threshold over the active steps means total drive > 0 (exact when the potential does not hit the floor
    of 0 on the way). Returns accuracy on the validation labels."""
    unit = float(np.median(np.abs(sv[:, :sp.ACTIVE].reshape(-1, sv.shape[-1]) @ rq + bias))) + 1
    thr = k * unit
    return round(100 * (sp.readout_neuron(sv, rq, bias + int(np.ceil(thr / sp.ACTIVE)), thr) == yv).mean(), 1)


print("dopamine-trained, calibrated decision neuron:", {k: calibrated(sv, rq[:-1], bias, k) for k in (10, 50, 200)})
f = lambda c: np.hstack([c.astype(np.float32), np.ones((len(c), 1), np.float32)])
sol = np.linalg.solve(f(ct).T @ f(ct) + 1e6 * np.eye(ct.shape[1] + 1), f(ct).T @ yt.astype(np.float64))
rr = np.rint(sol / np.abs(sol).max() * 255).astype(np.int64)
print("ridge, calibrated decision neuron:", {k: calibrated(sv, rr[:-1], int(np.rint(rr[-1] / sp.ACTIVE)), k) for k in (10, 50, 200)})