"""Stroke-detector + pooling features for the AltAI column network (prototype of a spiking feature layer).

Layer 1: 4 oriented line detectors (0°, 45°, 90°, 135°), 5×5 Gabor-like kernels, rectified.
Layer 2: max pooling 3×3 with stride 2 → 4 maps × 14×14 = 784 values, so the network input size is unchanged.
Values are scaled to 0–255 with a fixed global scale (no per-image normalisation, labels are not used).

usage: python stroke_features.py <src.bin> <src.target> <out_prefix>
The output keeps the source layout (training images, then evaluation images); augment first if needed.
"""
import argparse
import shutil
import sys

import numpy as np

SIDE, ORIENT, POOLED = 28, 4, 14
SCALE = 3.5  # fixed scale: 99th percentile of pooled responses on MNIST maps to 255


def kernels(size=5, sigma=1.4, wavelength=4.0):
    half = size // 2
    y, x = np.mgrid[-half:half + 1, -half:half + 1].astype(np.float32)
    out = []
    for k in range(ORIENT):
        t = np.pi * k / ORIENT
        xr = x * np.cos(t) + y * np.sin(t)
        yr = -x * np.sin(t) + y * np.cos(t)
        g = np.exp(-(xr ** 2 + 0.25 * yr ** 2) / (2 * sigma ** 2)) * np.cos(2 * np.pi * xr / wavelength)
        out.append(g - g.mean())
    return np.stack(out)


KERNELS = kernels()


def detect(imgs):
    """imgs: [n, 28, 28] uint8 → rectified responses [n, 4, 28, 28] float32."""
    x = imgs.astype(np.float32) / 255
    pad = KERNELS.shape[-1] // 2
    xp = np.pad(x, ((0, 0), (pad, pad), (pad, pad)))
    out = np.zeros((len(x), ORIENT, SIDE, SIDE), dtype=np.float32)
    for dy in range(KERNELS.shape[1]):
        for dx in range(KERNELS.shape[2]):
            patch = xp[:, dy:dy + SIDE, dx:dx + SIDE]
            out += KERNELS[None, :, dy, dx, None, None] * patch[:, None]
    return np.maximum(out, 0)


def pool(resp):
    """Max pooling 3×3, stride 2, padding 1: [n, 4, 28, 28] → [n, 4, 14, 14]."""
    p = np.pad(resp, ((0, 0), (0, 0), (1, 1), (1, 1)))
    out = np.zeros(resp.shape[:2] + (POOLED, POOLED), dtype=np.float32)
    for dy in range(3):
        for dx in range(3):
            out = np.maximum(out, p[:, :, dy:dy + 2 * POOLED:2, dx:dx + 2 * POOLED:2])
    return out


def features(imgs, batch=5000):
    res = []
    for i in range(0, len(imgs), batch):
        f = pool(detect(imgs[i:i + batch]))
        res.append(np.clip(f * (255 / SCALE), 0, 255).astype(np.uint8).reshape(len(f), -1))
    return np.concatenate(res)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("images")
    ap.add_argument("labels")
    ap.add_argument("out")
    a = ap.parse_args()
    imgs = np.fromfile(a.images, dtype=np.uint8).reshape(-1, SIDE, SIDE)
    f = features(imgs)
    f.tofile(a.out + ".bin")
    shutil.copyfile(a.labels, a.out + ".target")
    print(f"{a.out}: {len(f)} images, mean value {f.mean():.1f}, ~{f.mean() / 255 * 10 * 784:.0f} spikes per image,"
          f" saturated {100 * (f == 255).mean():.2f}%", file=sys.stderr)


if __name__ == "__main__":
    main()