"""Learned local features for small colour images (Coates & Ng 2011 style), prototype for a spiking layer.

Layer 1 is learned without labels: random 6×6×3 patches are contrast-normalised and ZCA-whitened, then
spherical k-means finds K prototypes. This is the same rule as winner-take-all competitive Hebbian learning:
the closest prototype moves towards the patch. Activation of prototype k at a position is
max(0, <patch, d_k> - alpha) (a neuron with threshold alpha). Activations are sum-pooled over a grid of
regions, which gives tolerance to position.

API: learn(images) -> model dict; encode(images, model) -> float32 features [n, K * grid * grid].
images: uint8 [n, 3, 32, 32].
"""
import numpy as np

PATCH = 6


def patches_at(img, ys, xs):
    """img [n,3,32,32] float32 → patches [n, len(ys)*len(xs), 108]."""
    out = np.stack([img[:, :, y:y + PATCH, x:x + PATCH].reshape(len(img), -1) for y in ys for x in xs], axis=1)
    return out


# CONTRAST_DIVIDE=False keeps only the (linear) patch mean subtraction, so mean removal, ZCA and the
# prototypes fold into one weight matrix of a spiking neuron; division by patch contrast is not linear.
CONTRAST_DIVIDE = True


def normalise(p):
    p = p - p.mean(-1, keepdims=True)
    return p / np.sqrt(p.var(-1, keepdims=True) + 10 / 255) if CONTRAST_DIVIDE else p * 4


def learn(images, k=400, n_patches=200000, iters=15, seed=0):
    rng = np.random.default_rng(seed)
    x = images.astype(np.float32) / 255
    idx = rng.integers(0, len(x), n_patches)
    ys, xs = rng.integers(0, 32 - PATCH + 1, n_patches), rng.integers(0, 32 - PATCH + 1, n_patches)
    p = np.stack([x[i, :, y:y + PATCH, xx:xx + PATCH].reshape(-1) for i, y, xx in zip(idx, ys, xs)])
    p = normalise(p)
    mean = p.mean(0)
    cov = np.cov(p - mean, rowvar=False)
    e, v = np.linalg.eigh(cov)
    zca = (v * (1 / np.sqrt(e + 0.1))) @ v.T
    pw = (p - mean) @ zca
    d = rng.normal(size=(k, pw.shape[1])).astype(np.float32)
    d /= np.linalg.norm(d, axis=1, keepdims=True)
    for _ in range(iters):
        s = pw @ d.T
        win = s.argmax(1)
        new = np.zeros_like(d)
        np.add.at(new, win, pw * s[np.arange(len(pw)), win, None])
        empty = np.linalg.norm(new, axis=1) < 1e-6
        new[empty] = rng.normal(size=(empty.sum(), d.shape[1]))
        d = new / np.linalg.norm(new, axis=1, keepdims=True)
    return {"mean": mean.astype(np.float32), "zca": zca.astype(np.float32), "d": d.astype(np.float32)}


def encode(images, model, alpha=0.25, grid=2, batch=250):
    x = images.astype(np.float32) / 255
    pos = list(range(0, 32 - PATCH + 1))
    side = len(pos)
    edges = np.linspace(0, side, grid + 1).astype(int)
    out = []
    for i in range(0, len(x), batch):
        p = normalise(patches_at(x[i:i + batch], pos, pos))
        a = np.maximum(((p - model["mean"]) @ model["zca"]) @ model["d"].T - alpha, 0)
        a = a.reshape(len(p), side, side, -1)
        pooled = [a[:, edges[gy]:edges[gy + 1], edges[gx]:edges[gx + 1]].sum((1, 2))
                  for gy in range(grid) for gx in range(grid)]
        out.append(np.concatenate(pooled, axis=1))
    return np.concatenate(out)