# N-MNIST-style event generation: a static digit on a "screen" is seen by an event camera that makes
# three small saccades (a triangle). Same algorithm as the browser demo (snn-mnist.htm, makeEvents).
import torch
S, PAD = 34, 3                    # sensor 34 x 34 (as N-MNIST), digit placed at offset (3, 3)
T, SUB, C = 30, 4, 0.3            # 30 time bins (10 per saccade), 4 sub-steps per bin, contrast threshold
VERT = [(0.0, 0.0), (2.0, 4.0), (4.0, 0.0), (0.0, 0.0)]   # saccade triangle, pixels (dx, dy)

def path(k):
    """Sensor offset at sub-step k (0 .. T*SUB): linear along the triangle, 3 equal saccades."""
    s = 3 * k / (T * SUB); i = min(2, int(s)); f = s - i
    (x0, y0), (x1, y1) = VERT[i], VERT[i + 1]
    return x0 + f * (x1 - x0), y0 + f * (y1 - y0)

def scene(img, dx, dy):
    """Brightness seen by the sensor when the screen is shifted by (dx, dy): bilinear, black level 0.1."""
    B = img.shape[0]; canvas = torch.zeros(B, S, S); canvas[:, PAD:PAD + 28, PAD:PAD + 28] = img
    ix, iy = int(dx // 1), int(dy // 1); fx, fy = dx - ix, dy - iy
    def sh(a, sx, sy):  # value at (y - sy, x - sx), zero outside
        out = torch.zeros_like(a)
        ys, yd = (slice(0, S - sy), slice(sy, S)) if sy >= 0 else (slice(-sy, S), slice(0, S + sy))
        xs, xd = (slice(0, S - sx), slice(sx, S)) if sx >= 0 else (slice(-sx, S), slice(0, S + sx))
        out[:, yd, xd] = a[:, ys, xs]; return out
    v = ((1 - fx) * (1 - fy) * sh(canvas, ix, iy) + fx * (1 - fy) * sh(canvas, ix + 1, iy)
         + (1 - fx) * fy * sh(canvas, ix, iy + 1) + fx * fy * sh(canvas, ix + 1, iy + 1))
    return 0.1 + 0.9 * v

def make_events(img):
    """img: (B, 28, 28) in [0, 1] -> event counts (T, B, 2, 34, 34): channel 0 = ON, 1 = OFF."""
    B = img.shape[0]; out = torch.zeros(T, B, 2, S, S)
    ref = torch.log(scene(img, *path(0)))
    for k in range(1, T * SUB + 1):
        L = torch.log(scene(img, *path(k))); d = L - ref
        n = torch.floor(d.abs() / C) * torch.sign(d)      # signed number of threshold crossings
        ref = ref + n * C
        b = (k - 1) // SUB
        out[b, :, 0] += n.clamp(min=0); out[b, :, 1] += (-n).clamp(min=0)
    return out