"""Q1: train one sample MNIST model per hardware world and export each to ONNX."""
import json
import os

import onnxruntime as ort
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

ROOT = os.path.dirname(os.path.abspath(__file__))
MODEL_DIR = os.path.join(ROOT, "..", "web", "public", "models")
DATA_DIR = os.path.join(ROOT, "..", "web", "public", "data")
os.makedirs(MODEL_DIR, exist_ok=True)
os.makedirs(DATA_DIR, exist_ok=True)

torch.manual_seed(0)
DEVICE = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
EPOCHS = 4

to_tensor = transforms.ToTensor()  # pixels in [0, 1], white digit on black
train_set = datasets.MNIST(os.path.join(ROOT, "mnist"), train=True, download=True, transform=to_tensor)
test_set = datasets.MNIST(os.path.join(ROOT, "mnist"), train=False, download=True, transform=to_tensor)
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
test_loader = DataLoader(test_set, batch_size=1000)


# ---------- helpers ----------
def binarize(t):
    """Forward: +1 / -1. Backward: pass the gradient straight through."""
    return t + (torch.sign(t) - t).detach()


def spike(v, threshold=1.0):
    """Forward: fire (1) if charge crosses the threshold, else 0. Backward: smooth surrogate."""
    hard = (v > threshold).float()
    soft = torch.sigmoid(5.0 * (v - threshold))
    return hard + (soft - soft.detach())


# ---------- World 0: GPUs (dense multiplication is cheap) ----------
class DenseMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Flatten(), nn.Linear(784, 256), nn.ReLU(),
            nn.Linear(256, 128), nn.ReLU(), nn.Linear(128, 10),
        )

    def forward(self, x):
        return self.net(x)


# ---------- World A: many tiny processors (only neighbor communication) ----------
class NeighborGrid(nn.Module):
    """Every pixel is a tiny processor holding `ch` numbers. Each step, a cell
    updates using only its 3x3 neighborhood. After `steps`, cells vote on the digit."""

    def __init__(self, ch=16, steps=10):
        super().__init__()
        self.ch, self.steps = ch, steps
        self.perceive = nn.Conv2d(ch, 48, 3, padding=1)  # look at 8 neighbors
        self.update = nn.Conv2d(48, ch, 1)                # compute inside the cell
        nn.init.zeros_(self.update.weight)
        nn.init.zeros_(self.update.bias)

    def forward(self, x):
        state = F.pad(x, (0, 0, 0, 0, 0, self.ch - 1))
        for _ in range(self.steps):
            state = state + self.update(F.relu(self.perceive(state)))
            state = torch.clamp(state, -10.0, 10.0)
            state = torch.cat([x, state[:, 1:]], dim=1)  # channel 0 always holds the pixel
        return state[:, -10:].mean(dim=(2, 3))           # every cell votes


# ---------- World B: logic gates (bit operations are nearly free) ----------
class BinaryLinear(nn.Linear):
    def forward(self, x):
        return F.linear(x, binarize(self.weight), None)


class BinaryNet(nn.Module):
    """All weights and activations are +1/-1, so each multiply is one bit comparison.
    BatchNorm + sign acts as a threshold on a bit count."""

    def __init__(self):
        super().__init__()
        self.l1, self.b1 = BinaryLinear(784, 512, bias=False), nn.BatchNorm1d(512)
        self.l2, self.b2 = BinaryLinear(512, 512, bias=False), nn.BatchNorm1d(512)
        self.l3, self.b3 = BinaryLinear(512, 10, bias=False), nn.BatchNorm1d(10)

    def forward(self, x):
        x = binarize(x.flatten(1) - 0.5)
        h = binarize(self.b1(self.l1(x)))
        h = binarize(self.b2(self.l2(h)))
        return self.b3(self.l3(h))

    def clamp_(self):
        for m in self.modules():
            if isinstance(m, BinaryLinear):
                m.weight.data.clamp_(-1, 1)


# ---------- World C: instant lookup (find the closest stored item in one step) ----------
class PrototypeLookup(nn.Module):
    def __init__(self, prototypes, per_class):
        super().__init__()
        self.register_buffer("P", prototypes)  # [10 * per_class, 784], grouped by class
        self.k = per_class

    def forward(self, x):
        x = x.flatten(1)
        d = (x * x).sum(1, keepdim=True) - 2 * x @ self.P.T + (self.P * self.P).sum(1)
        d = d.view(-1, 10, self.k)
        return -d.min(dim=2).values  # closest stored example per class


@torch.no_grad()
def build_prototypes(per_class=30, iters=25):
    X = train_set.data.float().div(255).view(-1, 784).to(DEVICE)
    Y = train_set.targets.to(DEVICE)
    protos = []
    for c in range(10):
        Xc = X[Y == c]
        C = Xc[torch.randperm(len(Xc), device=DEVICE)[:per_class]].clone()
        for _ in range(iters):
            assign = torch.cdist(Xc, C).argmin(1)
            for j in range(per_class):
                members = Xc[assign == j]
                if len(members):
                    C[j] = members.mean(0)
        protos.append(C)
    return torch.cat(protos).cpu()


# ---------- World D: brain-like (energy only when a neuron fires) ----------
class SpikingNet(nn.Module):
    def __init__(self, steps=12, beta=0.85):
        super().__init__()
        self.steps, self.beta = steps, beta
        self.fc1, self.fc2 = nn.Linear(784, 256), nn.Linear(256, 10)

    def forward(self, x):
        current = self.fc1(x.flatten(1))
        v1 = current * 0
        v2 = count = total = None
        for _ in range(self.steps):
            v1 = self.beta * v1 + current
            s1 = spike(v1)
            v1 = v1 - s1
            c2 = self.fc2(s1)
            v2 = c2 if v2 is None else self.beta * v2 + c2
            s2 = spike(v2)
            v2 = v2 - s2
            count = s2 if count is None else count + s2
            step = s1.sum(1) + s2.sum(1)
            total = step if total is None else total + step
        self.spike_count = total
        return count + v2  # spike count, leftover charge breaks ties


# ---------- training / evaluation / export ----------
@torch.no_grad()
def evaluate(model, loader=None):
    model.eval().to(DEVICE)
    correct = total = 0
    for x, y in loader or test_loader:
        correct += (model(x.to(DEVICE)).argmax(1).cpu() == y).sum().item()
        total += len(y)
    return correct / total


@torch.no_grad()
def confusion(model):
    model.eval().to(DEVICE)
    M = torch.zeros(10, 10, dtype=torch.long)
    for x, y in test_loader:
        p = model(x.to(DEVICE)).argmax(1).cpu()
        M.index_put_((y, p), torch.ones_like(y), accumulate=True)
    return M.tolist()


@torch.no_grad()
def robustness(model):
    """Accuracy on the test set under Gaussian pixel noise and under translation."""
    model.eval().to(DEVICE)
    g = torch.Generator().manual_seed(0)
    noise, shift = [], []
    for sigma in NOISE_LEVELS:
        correct = 0
        for x, y in test_loader:
            xn = (x + sigma * torch.randn(x.shape, generator=g)).clamp(0, 1)
            correct += (model(xn.to(DEVICE)).argmax(1).cpu() == y).sum().item()
        noise.append(round(correct / len(test_set), 4))
    for k in SHIFTS:
        correct = 0
        for x, y in test_loader:
            xs = torch.roll(x, shifts=(k, k), dims=(2, 3))  # diagonal shift by k pixels
            correct += (model(xs.to(DEVICE)).argmax(1).cpu() == y).sum().item()
        shift.append(round(correct / len(test_set), 4))
    return {"noise": noise, "shift": shift}


def train(model, lr):
    model.to(DEVICE)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, EPOCHS * len(train_loader))
    history = []
    for epoch in range(EPOCHS):
        model.train()
        loss_sum, n = 0.0, 0
        for x, y in train_loader:
            x, y = x.to(DEVICE), y.to(DEVICE)
            loss = F.cross_entropy(model(x), y)
            opt.zero_grad()
            loss.backward()
            opt.step()
            sched.step()
            if hasattr(model, "clamp_"):
                model.clamp_()
            loss_sum += loss.item() * len(y)
            n += len(y)
        acc = evaluate(model)
        history.append({"epoch": epoch + 1, "trainLoss": round(loss_sum / n, 4), "testAccuracy": round(acc, 4)})
        print(f"  epoch {epoch + 1}: train loss {loss_sum / n:.4f}, test accuracy {acc:.4f}")
    return model, history


@torch.no_grad()
def spike_stats(model):
    """Average hidden and output spikes per test digit."""
    model.eval().to(DEVICE)
    hidden = output = 0.0
    for x, _ in test_loader:
        _, s1, s2 = SpikeTrace(model)(x.to(DEVICE))
        hidden += s1.sum().item()
        output += s2.sum().item()
    return hidden / len(test_set), output / len(test_set)


def count_params(model):
    return sum(t.numel() for t in list(model.parameters()) + list(model.buffers()) if t.dtype.is_floating_point)


def export(model, name, outputs=("logits",), check=True):
    model.eval().cpu()
    path = os.path.join(MODEL_DIR, f"{name}.onnx")
    torch.onnx.export(
        model, torch.zeros(1, 1, 28, 28), path,
        input_names=["input"], output_names=list(outputs),
        opset_version=17, dynamo=False,
    )
    size_kb = os.path.getsize(path) / 1024
    agree = None
    if check:
        sess = ort.InferenceSession(path)
        agree = 0
        with torch.no_grad():
            for i in range(100):
                x = test_set[i][0].unsqueeze(0)
                a = model(x).argmax(1).item()
                b = sess.run(None, {"input": x.numpy()})[0].argmax(1).item()
                agree += a == b
        print(f"  ONNX matches PyTorch on {agree}/100 images, {size_kb:.0f} KB")
    return round(size_kb), agree


# ---------- extra exports used by the interactive figures ----------
class CASteps(nn.Module):
    """Same computation as NeighborGrid, but returns every cell's vote at every step."""

    def __init__(self, m):
        super().__init__()
        self.m = m

    def forward(self, x):
        m = self.m
        state = F.pad(x, (0, 0, 0, 0, 0, m.ch - 1))
        votes = []
        for _ in range(m.steps):
            state = state + m.update(F.relu(m.perceive(state)))
            state = torch.clamp(state, -10.0, 10.0)
            state = torch.cat([x, state[:, 1:]], dim=1)
            votes.append(state[:, -10:])
        return torch.stack(votes, 1)  # [B, steps, 10, 28, 28]


class SpikeTrace(nn.Module):
    """Same computation as SpikingNet, but returns the spike trains."""

    def __init__(self, m):
        super().__init__()
        self.m = m

    def forward(self, x):
        m = self.m
        current = m.fc1(x.flatten(1))
        v1 = current * 0
        v2 = count = None
        h, o = [], []
        for _ in range(m.steps):
            v1 = m.beta * v1 + current
            s1 = spike(v1)
            v1 = v1 - s1
            c2 = m.fc2(s1)
            v2 = c2 if v2 is None else m.beta * v2 + c2
            s2 = spike(v2)
            v2 = v2 - s2
            count = s2 if count is None else count + s2
            h.append(s1)
            o.append(s2)
        return count + v2, torch.stack(h, 1), torch.stack(o, 1)  # logits, [B,T,256], [B,T,10]


def main():
    results = []

    def record(id_, m, history, ops, extra=None):
        size_kb, agree = export(m, id_)
        r = {"id": id_, "accuracy": round(evaluate(m), 4), "sizeKB": size_kb, "onnxAgreement": agree,
             "params": count_params(m), "history": history, "ops": ops,
             "confusion": confusion(m), "robustness": robustness(m)}
        r.update(extra or {})
        results.append(r)
        return r

    print("World 0: GPUs -> dense network")
    m, h = train(DenseMLP(), 1e-3)
    macs = 784 * 256 + 256 * 128 + 128 * 10
    record("gpu_mlp", m, h, {"multiplies": macs, "bitOps": 0, "lookups": 0, "events": 0,
                             "note": "Dense multiply-accumulates (MACs)"},
           {"statLabel": "Multiplications per digit", "statValue": f"{macs:,}"})

    print("World A: many tiny processors -> neighbor grid")
    m, h = train(NeighborGrid(), 2e-3)
    per_cell = (16 * 48 * 9 + 48 * 16) * 10
    record("ca_grid", m, h, {"multiplies": per_cell * 784, "perCell": per_cell, "bitOps": 0, "lookups": 0, "events": 0,
                             "note": "MACs, all local: each cell reads only its 8 neighbors"},
           {"statLabel": "Multiplications per cell (all local)", "statValue": f"{per_cell:,}"})
    export(CASteps(m), "ca_grid_steps", outputs=("votes",), check=False)

    print("World B: logic gates -> binary network")
    m, h = train(BinaryNet(), 1e-3)
    bits = 784 * 512 + 512 * 512 + 512 * 10
    record("binary_net", m, h, {"multiplies": 0, "bitOps": bits, "lookups": 0, "events": 0,
                                "weightBits": bits, "note": "XNOR + popcount; BatchNorm folds into thresholds"},
           {"statLabel": "Bit operations per digit", "statValue": f"{bits:,}",
            "weightKB": round(bits / 8 / 1024, 1)})

    print("World C: instant lookup -> stored prototypes")
    protos = build_prototypes(30)
    m = PrototypeLookup(protos, 30)
    print(f"  test accuracy {evaluate(m):.4f}")
    record("lookup_protos", m, [], {"multiplies": 0, "bitOps": 0, "lookups": 1, "events": 0,
                                    "comparisons": 300, "note": "One nearest-match search over 300 stored digits"},
           {"statLabel": "Stored examples compared per digit", "statValue": "300"})
    with open(os.path.join(MODEL_DIR, "prototypes.bin"), "wb") as f:
        f.write((protos.clamp(0, 1) * 255).round().to(torch.uint8).numpy().tobytes())  # [300, 784], 30 per class

    print("World D: brain-like -> spiking network")
    m, h = train(SpikingNet(), 1e-3)
    hid, out = spike_stats(m)
    record("spiking_net", m, h, {"multiplies": 784 * 256, "bitOps": 0, "lookups": 0,
                                 "events": round(hid * 10 + out), "hiddenSpikes": round(hid, 1),
                                 "outputSpikes": round(out, 1), "steps": 12,
                                 "note": "Input layer is computed once; after that, work happens only when a neuron spikes"},
           {"statLabel": "Average spikes per digit (energy is spent only on these)", "statValue": f"{hid + out:,.0f}"})
    export(SpikeTrace(m), "spiking_trace", outputs=("logits", "hidden", "output"), check=False)

    out = {"noiseLevels": NOISE_LEVELS, "shifts": SHIFTS, "epochs": EPOCHS, "trainSize": len(train_set),
           "testSize": len(test_set), "models": results}
    with open(os.path.join(DATA_DIR, "q1_results.json"), "w") as f:
        json.dump(out, f, indent=2)
    print("Wrote q1_results.json")


NOISE_LEVELS = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5]
SHIFTS = [0, 1, 2, 3, 4]

if __name__ == "__main__":
    main()