"""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 = 28
# Presets: "fine" — 4 orientations, pool 3×3 stride 2 → 4×14×14;
#          "coarse" — 8 orientations × 2 stroke widths, pool 5×5 stride 4 → 16×7×7. Both give 784 values.
PRESETS = {
    "fine": {"orient": 4, "sigmas": [1.4], "pool": 3, "stride": 2, "pooled": 14, "scale": 3.5,
             "size": 5, "wavelength": 4.0, "norm": False},
    "coarse": {"orient": 8, "sigmas": [1.0, 2.0], "pool": 5, "stride": 4, "pooled": 7, "scale": None,
               "size": 7, "wavelength": None, "norm": True},
}
P = dict(PRESETS["fine"])


def kernels():
    half = P["size"] // 2
    y, x = np.mgrid[-half:half + 1, -half:half + 1].astype(np.float32)
    out = []
    for sigma in P["sigmas"]:
        for k in range(P["orient"]):
            t = np.pi * k / P["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 / (P["wavelength"] or 2.9 * sigma))
            g = g - g.mean()
            out.append(g / np.abs(g).sum() if P["norm"] else g)
    return np.stack(out)


def use_preset(name):
    global KERNELS
    P.clear()
    P.update(PRESETS[name])
    KERNELS = kernels()


KERNELS = None


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), len(KERNELS), 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 P['pool']×P['pool'] with stride P['stride']: [n, maps, 28, 28] → [n, maps, pooled, pooled]."""
    k, st, n = P["pool"], P["stride"], P["pooled"]
    pad = (k - st) // 2 + 1
    p = np.pad(resp, ((0, 0), (0, 0), (pad, pad), (pad, pad)))
    out = np.zeros(resp.shape[:2] + (n, n), dtype=np.float32)
    for dy in range(k):
        for dx in range(k):
            out = np.maximum(out, p[:, :, dy:dy + st * n:st, dx:dx + st * n:st])
    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 / P["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")
    ap.add_argument("--preset", choices=sorted(PRESETS), default="fine")
    ap.add_argument("--scale", type=float, help="response mapped to 255 (default: preset value or 99th percentile)")
    a = ap.parse_args()
    use_preset(a.preset)
    imgs = np.fromfile(a.images, dtype=np.uint8).reshape(-1, SIDE, SIDE)
    if a.scale:
        P["scale"] = a.scale
    elif P["scale"] is None:
        P["scale"] = float(np.percentile(pool(detect(imgs[:2000])), 99))
        print(f"scale (99th percentile of first 2000 images): {P['scale']:.4g}", file=sys.stderr)
    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()