# Trains the ReLU CNN and its spiking (LIF) twin on MNIST for the "From ReLU to Spikes" post, then exports
# weights, 120 test digits and reference spike counts to content/Spiking Neural Networks/demo/snn-mnist.json.
# Requires PyTorch + torchvision (CPU is enough, ~15 min). Usage: python3 train_spiking_cnn.py
# Same small CNN twice: ReLU (ANN) and leaky integrate-and-fire (SNN, surrogate gradient, BPTT over T steps).
import json, math, sys, time, torch, torch.nn as nn, torch.nn.functional as F, torchvision
torch.manual_seed(0)
T, BETA, THETA = 25, 0.9, 1.0
tr = torchvision.datasets.MNIST("data", train=True); te = torchvision.datasets.MNIST("data", train=False)
Xtr = tr.data.float().div(255).unsqueeze(1); Ytr = tr.targets
Xte = te.data.float().div(255).unsqueeze(1); Yte = te.targets

class Spike(torch.autograd.Function):
    """Heaviside forward, arctan surrogate backward: dS/dU ~ 1 / (1 + (pi U)^2)."""
    @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 Net(nn.Module):
    def __init__(s, spiking):
        super().__init__(); s.spiking = spiking
        s.c1 = nn.Conv2d(1, 12, 5); s.c2 = nn.Conv2d(12, 32, 5); s.fc = nn.Linear(32 * 4 * 4, 10)
    def lif(s, cur, u):
        # U[t] = beta U[t-1] + I[t] - S[t-1] theta  (soft reset, reset detached from the graph)
        s_prev = (u > THETA).float().detach()
        u = BETA * u + cur - s_prev * THETA
        return spike(u - THETA), u
    def forward(s, x, steps=T, encoding="direct", record=False):
        if not s.spiking:
            h = F.relu(F.max_pool2d(s.c1(x), 2)); h = F.relu(F.max_pool2d(s.c2(h), 2)); return s.fc(h.flatten(1))
        B = x.shape[0]; u1 = torch.zeros(B, 12, 12, 12); u2 = torch.zeros(B, 32, 4, 4); u3 = torch.zeros(B, 10)
        counts = 0; rec = []
        for t in range(steps):
            xt = x if encoding == "direct" else (torch.rand_like(x) < x).float()
            s1, u1 = s.lif(F.max_pool2d(s.c1(xt), 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(), s3.mean().item()))
        return (counts, rec) if record else counts

def evaluate(net, steps=T, encoding="direct", n=10000):
    net.eval(); correct = 0
    with torch.no_grad():
        for i in range(0, n, 500):
            out = net(Xte[i:i + 500], steps, encoding) if net.spiking else net(Xte[i:i + 500])
            correct += (out.argmax(1) == Yte[i:i + 500]).sum().item()
    return correct / n

def train(net, epochs):
    opt = torch.optim.Adam(net.parameters(), lr=2e-3)
    for ep in range(epochs):
        net.train(); perm = torch.randperm(60000); t0 = time.time()
        for i in range(0, 60000, 128):
            idx = perm[i:i + 128]; out = net(Xtr[idx])
            loss = F.cross_entropy(out, Ytr[idx]); opt.zero_grad(); loss.backward(); opt.step()
        print(("SNN" if net.spiking else "ANN"), "epoch", ep, "acc", evaluate(net), round(time.time() - t0), "s", flush=True)

ann = Net(False); train(ann, 3)
snn = Net(True); train(snn, 3)
res = {"ann": evaluate(ann), "snn": evaluate(snn)}
res["snn_by_T"] = {t: evaluate(snn, t) for t in [1, 2, 3, 4, 5, 6, 8, 10, 12, 15, 20, 25, 30, 40]}
torch.manual_seed(1); res["snn_rate_input"] = evaluate(snn, T, "rate")
with torch.no_grad():
    _, rec = snn(Xte[:1000], T, record=True)
res["spike_rate_per_layer"] = [sum(r[k] for r in rec) / len(rec) for k in range(3)]
print(json.dumps(res, indent=1))
torch.save({"ann": ann.state_dict(), "snn": snn.state_dict()}, "models.pt")
json.dump(res, open("results.json", "w"))

# ---------------- export ----------------
res = json.load(open("results.json")); res["ann_epoch_s"] = 10; res["snn_epoch_s"] = 259
W = lambda m: {k: {"shape": list(v.shape), "data": [float(f"{x:.5g}") for x in v.flatten().tolist()]} for k, v in m.state_dict().items()}
g = torch.Generator().manual_seed(3); idx = torch.randperm(10000, generator=g)[:120].tolist()
samples = [{"px": te.data[i].flatten().tolist(), "y": int(te.targets[i])} for i in idx]
# reference spike counts with the rounded weights, for checking the browser implementation
snn_r = Net(True); snn_r.load_state_dict({k: torch.tensor(v["data"]).reshape(v["shape"]) for k, v in W(snn).items()})
with torch.no_grad():
    x = torch.stack([te.data[i].float().div(255).unsqueeze(0) for i in idx[:5]])
    refs = snn_r(x, 25).tolist()
json.dump({"snn": W(snn), "ann": W(ann), "samples": samples, "results": res, "refs": refs}, open("../../content/Spiking Neural Networks/demo/snn-mnist.json", "w"), separators=(",", ":"))
print(refs)