# Trains the spiking CNN on N-MNIST-style simulated saccade events (see events.py) for the "From ReLU to Spikes" post.
# Requires PyTorch + torchvision (CPU, ~30 min). Writes model_events.pt and results_events.json.
# Spiking CNN on N-MNIST-style simulated saccade events (2 polarity channels, 34 x 34, 30 bins).
import json, math, time, torch, torch.nn as nn, torch.nn.functional as F, torchvision
from events import make_events, T as TB
torch.manual_seed(0)
BETA, THETA = 0.9, 1.0
tr = torchvision.datasets.MNIST("data", train=True); te = torchvision.datasets.MNIST("data", train=False)
Xtr = tr.data.float().div(255); Ytr = tr.targets; Xte = te.data.float().div(255); Yte = te.targets
class Spike(torch.autograd.Function):
    @staticmethod
    def forward(ctx, u): ctx.save_for_backward(u); return (u > 0).float()
    @staticmethod
    def backward(ctx, g): u, = ctx.saved_tensors; return g / (1 + (math.pi * u) ** 2)
spike = Spike.apply
class EvNet(nn.Module):
    def __init__(s):
        super().__init__(); s.c1 = nn.Conv2d(2, 12, 5); s.c2 = nn.Conv2d(12, 32, 5); s.fc = nn.Linear(32 * 5 * 5, 10)
    def lif(s, cur, u):
        s_prev = (u > THETA).float().detach(); u = BETA * u + cur - s_prev * THETA; return spike(u - THETA), u
    def forward(s, ev, steps=TB, record=False):   # ev: (T, B, 2, 34, 34)
        B = ev.shape[1]; u1 = torch.zeros(B, 12, 15, 15); u2 = torch.zeros(B, 32, 5, 5); u3 = torch.zeros(B, 10); counts = 0; rec = []
        for t in range(steps):
            s1, u1 = s.lif(F.max_pool2d(s.c1(ev[t]), 2), u1)
            s2, u2 = s.lif(F.max_pool2d(s.c2(s1), 2), u2)
            s3, u3 = s.lif(s.fc(s2.flatten(1)), u3); counts = counts + s3
            if record: rec.append((s1.mean().item(), s2.mean().item()))
        return (counts, rec) if record else counts
def evaluate(net, steps=TB, n=10000):
    net.eval(); c = 0
    with torch.no_grad():
        for i in range(0, n, 500): c += (net(make_events(Xte[i:i + 500]), steps).argmax(1) == Yte[i:i + 500]).sum().item()
    return c / n
net = EvNet(); opt = torch.optim.Adam(net.parameters(), lr=2e-3)
for ep in range(3):
    net.train(); perm = torch.randperm(60000); t0 = time.time()
    for i in range(0, 60000, 128):
        idx = perm[i:i + 128]; out = net(make_events(Xtr[idx])); loss = F.cross_entropy(out, Ytr[idx]); opt.zero_grad(); loss.backward(); opt.step()
    print("epoch", ep, "acc", evaluate(net), round(time.time() - t0), "s", flush=True)
res = {"acc": evaluate(net), "by_T": {t: evaluate(net, t) for t in [3, 5, 10, 15, 20, 25, 30]}}
with torch.no_grad(): _, rec = net(make_events(Xte[:500]), record=True)
res["spike_rate"] = [sum(r[k] for r in rec) / len(rec) for k in range(2)]
ev = make_events(Xte[:2000]); res["events_per_sample"] = ev.sum((0, 2, 3, 4)).mean().item()
print(json.dumps(res)); json.dump(res, open("results_events.json", "w")); torch.save(net.state_dict(), "model_events.pt")