"""Build ON/OFF-coded MNIST datasets for mnist_learn_example (KNP_MNIST_INPUT_SIZE=1568).

Each image becomes 1568 bytes: 784 ON pixels (unchanged) followed by 784 OFF pixels.
OFF modes:
  full  - 255 - p everywhere;
  half  - (255 - p) / 2 everywhere;
  local - 255 - p only within `--radius` pixels of the digit (p > 0), zero elsewhere.

usage: python prepare_onoff.py <src.bin> <src.target> <out_prefix> --mode local [--radius 2]
Writes <out_prefix>.bin and <out_prefix>.target (labels copied as is).
"""
import argparse
import shutil

import numpy as np

SIDE = 28


def dilate(mask, radius):
    out = mask.copy()
    for dy in range(-radius, radius + 1):
        for dx in range(-radius, radius + 1):
            if dy * dy + dx * dx > radius * radius:
                continue
            shifted = np.zeros_like(mask)
            ys, yd = (slice(dy, None), slice(None, -dy or None)) if dy >= 0 else (slice(None, dy), slice(-dy, None))
            xs, xd = (slice(dx, None), slice(None, -dx or None)) if dx >= 0 else (slice(None, dx), slice(-dx, None))
            shifted[:, ys, xs] = mask[:, yd, xd]
            out |= shifted
    return out


def off_channel(imgs, mode, radius):
    inv = 255 - imgs.astype(np.int32)
    if mode == "full":
        off = inv
    elif mode == "half":
        off = inv // 2
    elif mode == "local":
        off = inv * dilate(imgs > 0, radius)
    else:
        raise ValueError(mode)
    return off.astype(np.uint8)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("images")
    ap.add_argument("labels")
    ap.add_argument("out")
    ap.add_argument("--mode", choices=["full", "half", "local"], required=True)
    ap.add_argument("--radius", type=int, default=2)
    a = ap.parse_args()
    imgs = np.fromfile(a.images, dtype=np.uint8).reshape(-1, SIDE, SIDE)
    off = off_channel(imgs, a.mode, a.radius)
    out = np.concatenate([imgs.reshape(len(imgs), -1), off.reshape(len(imgs), -1)], axis=1)
    out.tofile(a.out + ".bin")
    shutil.copyfile(a.labels, a.out + ".target")
    on_spk, off_spk = imgs.mean() / 255 * 10 * SIDE * SIDE, off.mean() / 255 * 10 * SIDE * SIDE
    print(f"{a.out}: {len(out)} images, ~{on_spk:.0f} ON / ~{off_spk:.0f} OFF spikes per image")


if __name__ == "__main__":
    main()