"""Design check for the all-on-platform pipeline (numpy prototype, not the final result).

Retina: per colour channel a centre–surround (difference of Gaussians) filter; ON = max(0, DoG), OFF = max(0, −DoG)
→ 6 maps 32×32 (fixed weights, i.e. one fixed layer of neurons). Features: online winner-take-all on 6×6×6 retina
patches (unit-norm weights, winner = largest w·x, winner moves towards the patch), activation max(0, w·x − θ) with
θ a quantile of the winners' responses; sum pooling over quadrants; linear readout (ridge here, delta rule in KNP).

usage: ~/anvenv/bin/python retina_proto.py <data dir with train.npz/val.npz> [--k 200 --stride 2]
"""
import argparse

import numpy as np

PATCH = 6


def gauss(sigma, r=3):
    y, x = np.mgrid[-r:r + 1, -r:r + 1]
    g = np.exp(-(x * x + y * y) / (2 * sigma * sigma))
    return g / g.sum()


DOG = gauss(0.8) - gauss(2.0)


def retina(rgb, gain=4.0):
    """rgb uint8 [n,3,32,32] → float32 [n,6,32,32] in [0,1]: ON and OFF maps per colour."""
    x = rgb.astype(np.float32) / 255
    r = DOG.shape[0] // 2
    xp = np.pad(x, ((0, 0), (0, 0), (r, r), (r, r)), mode="edge")
    d = np.zeros_like(x)
    for dy in range(DOG.shape[0]):
        for dx in range(DOG.shape[1]):
            d += DOG[dy, dx] * xp[:, :, dy:dy + 32, dx:dx + 32]
    return np.clip(np.concatenate([np.maximum(d, 0), np.maximum(-d, 0)], axis=1) * gain, 0, 1)


def patches(maps, stride):
    pos = range(0, 32 - PATCH + 1, stride)
    return np.stack([maps[:, :, y:y + PATCH, x:x + PATCH].reshape(len(maps), -1) for y in pos for x in pos], axis=1)


def learn_wta(maps, k, n=300000, lr0=0.05, seed=0):
    rng = np.random.default_rng(seed)
    i, y, x = rng.integers(0, len(maps), n), rng.integers(0, 27, n), rng.integers(0, 27, n)
    p = np.stack([maps[a, :, b:b + PATCH, c:c + PATCH].reshape(-1) for a, b, c in zip(i, y, x)])
    p = p[p.sum(1) > 0.5]  # skip empty patches (no retina activity)
    w = p[rng.choice(len(p), k, replace=False)].copy()
    w /= np.linalg.norm(w, axis=1, keepdims=True)
    for t, v in enumerate(p):
        win = int(np.argmax(w @ v))
        w[win] += lr0 / (1 + t / 20000) * (v / (np.linalg.norm(v) + 1e-6) - w[win])
        w[win] /= np.linalg.norm(w[win])
    return w, p


def encode(maps, w, theta, stride, batch=500, centre=False):
    """centre=True: subtract the mean prototype (activation max(0, (w_k − mean_j w_j)·x − θ)); still linear."""
    if centre:
        w = w - w.mean(0)
    out = []
    for s in range(0, len(maps), batch):
        p = patches(maps[s:s + batch], stride)
        a = np.maximum(p @ w.T - theta, 0)
        side = int(np.sqrt(p.shape[1]))
        a = a.reshape(len(p), side, side, -1)
        h = side // 2
        out.append(np.concatenate([a[:, :h, :h].sum((1, 2)), a[:, :h, h:].sum((1, 2)),
                                   a[:, h:, :h].sum((1, 2)), a[:, h:, h:].sum((1, 2))], axis=1))
    return np.concatenate(out)


def ridge_acc(ft, yt, fv, yv, lam, classes):
    mu, sd = ft.mean(0), ft.std(0) + 1e-3
    f = lambda z: np.hstack([(z - mu) / sd, np.ones((len(z), 1), np.float32)])
    w = np.linalg.solve(f(ft).T @ f(ft) + lam * np.eye(ft.shape[1] + 1), f(ft).T @ (np.eye(classes)[yt] * 2 - 1))
    return 100 * ((f(fv) @ w).argmax(1) == yv).mean()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("data")
    ap.add_argument("--k", type=int, default=200)
    ap.add_argument("--stride", type=int, default=2)
    ap.add_argument("--flip", action="store_true")
    ap.add_argument("--centre", action="store_true")
    ap.add_argument("--sigmas", default="0.8,2.0")
    ap.add_argument("--gain", type=float, default=4.0)
    a = ap.parse_args()
    global DOG
    s1, s2 = (float(v) for v in a.sigmas.split(","))
    DOG = gauss(s1) - gauss(s2)
    tr, va = np.load(f"{a.data}/train.npz"), np.load(f"{a.data}/val.npz")
    xt, yt, xv, yv = tr["rgb"], tr["labels"].astype(int), va["rgb"], va["labels"].astype(int)
    classes = int(yt.max()) + 1
    mt, mv = retina(xt, a.gain), retina(xv, a.gain)
    w, p = learn_wta(mt, a.k)
    wc = w - w.mean(0) if a.centre else w
    resp = (p[:20000] @ wc.T).max(1)
    for q in ((0, 50) if a.centre else (50, 75)):
        theta = float(np.percentile(resp, q)) if q else 0.0
        ft, fv = encode(mt, w, theta, a.stride, centre=a.centre), encode(mv, w, theta, a.stride, centre=a.centre)
        if a.flip:
            ft = np.vstack([ft, encode(retina(xt[:, :, :, ::-1], a.gain), w, theta, a.stride, centre=a.centre)])
            ytt = np.r_[yt, yt]
        else:
            ytt = yt
        print(f"K {a.k} stride {a.stride} theta=q{q} ({theta:.3f}): "
              + " | ".join(f"lam {lam}: {ridge_acc(ft, ytt, fv, yv, lam, classes):.1f}%" for lam in (100, 1000)), flush=True)


if __name__ == "__main__":
    main()