import marimo

__generated_with = "0.21.1"
app = marimo.App(width="medium")


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    # Training a Real Network on MNIST

    This notebook follows the [blog article](/posts/1008/en/training-and-experimenting-with-mnist/) —
    training a neural network with PyTorch to classify all 10 handwritten digits (0–9).
    """)
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## Prepare data
    """)
    return


@app.cell
def _():
    import marimo as mo
    import matplotlib.pyplot as plt
    import torch
    from torch import nn
    from torch.utils.data import DataLoader, TensorDataset
    from torchvision.datasets import MNIST
    import numpy as np

    train_data = MNIST("./data", train=True, download=True)
    test_data = MNIST("./data", train=False, download=True)
    # NumPy arrays also work with the from-scratch implementation.
    X_train = train_data.data.numpy().reshape(-1, 784).astype("float32") / 255.0
    X_test = test_data.data.numpy().reshape(-1, 784).astype("float32") / 255.0
    y_train = train_data.targets.numpy()
    y_test = test_data.targets.numpy()
    return DataLoader, TensorDataset, X_test, X_train, mo, nn, np, plt, torch, y_test, y_train


@app.cell
def _(DataLoader, TensorDataset, nn, torch):
    def run_epoch(model, X, y, batch_size=32, optimizer=None):
        training = optimizer is not None
        model.train(training)
        device = next(model.parameters()).device
        loader = DataLoader(
            TensorDataset(torch.as_tensor(X, dtype=torch.float32),
                          torch.as_tensor(y, dtype=torch.long)),
            batch_size=batch_size, shuffle=training,
        )
        loss_fn = nn.CrossEntropyLoss()
        total_loss, correct = 0.0, 0
        with torch.enable_grad() if training else torch.no_grad():
            for xb, yb in loader:
                xb, yb = xb.to(device), yb.to(device)
                if training:
                    optimizer.zero_grad()
                logits = model(xb)
                loss = loss_fn(logits, yb)
                if training:
                    loss.backward()
                    optimizer.step()
                total_loss += loss.item() * len(yb)
                correct += (logits.argmax(dim=1) == yb).sum().item()
        return total_loss / len(y), correct / len(y)


    def train_torch(model, X, y, epochs=5, lr=0.1, batch_size=32,
                    val_fraction=0.2, optimizer=None, include_initial=False):
        split = int(len(X) * (1 - val_fraction))
        X_fit, y_fit = X[:split], y[:split]
        X_val, y_val = X[split:], y[split:]
        if optimizer is None:
            optimizer = torch.optim.SGD(model.parameters(), lr=lr)
        history = {key: [] for key in ("loss", "accuracy", "val_loss", "val_accuracy")}
        for epoch in range(0 if include_initial else 1, epochs + 1):
            loss, accuracy = run_epoch(
                model, X_fit, y_fit, batch_size, optimizer if epoch else None)
            val_loss, val_accuracy = run_epoch(model, X_val, y_val, batch_size)
            for key, value in zip(history, (loss, accuracy, val_loss, val_accuracy)):
                history[key].append(value)
            print(f"Epoch {epoch}/{epochs} — loss: {loss:.4f}, accuracy: {accuracy:.4f}, "
                  f"val_loss: {val_loss:.4f}, val_accuracy: {val_accuracy:.4f}")
        return history
    return run_epoch, train_torch


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## Define model
    """)
    return


@app.cell
def _(nn):
    model = nn.Sequential(
        nn.Linear(784, 128),
        nn.ReLU(),
        nn.Linear(128, 10),  # raw logits; CrossEntropyLoss applies log-softmax
    )
    return (model,)


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## Train model
    """)
    return


@app.cell
def _(X_train, model, train_torch, y_train):
    history = train_torch(model, X_train, y_train, epochs=5)
    return (history,)


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## Evaluate model
    """)
    return


@app.cell
def _(X_test, model, run_epoch, y_test):
    # Final evaluation on the test set (images the model never saw during training or validation)
    torch_test_loss, torch_test_acc = run_epoch(model, X_test, y_test)
    print(f"Final test loss:     {torch_test_loss:.4f}")
    print(f"Final test accuracy: {torch_test_acc:.2%}")
    return



@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## Same network from scratch in NumPy
    """)
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ### Define layers
    """)
    return


@app.cell
def _(np):
    class HiddenLayer:
        def __init__(self, n_inputs, n_neurons):
            self.W = np.random.randn(n_neurons, n_inputs) * np.sqrt(2.0 / n_inputs)
            self.b = np.zeros(n_neurons)

        def forward(self, x):
            self.x = x
            self.z = self.W @ x + self.b
            self.out = np.maximum(0, self.z)
            return self.out

        def backward(self, grad_output):
            grad_z = grad_output * (self.z > 0)
            self.grad_W = np.outer(grad_z, self.x)
            self.grad_b = grad_z
            return self.W.T @ grad_z

        def update(self, lr):
            self.W -= lr * self.grad_W
            self.b -= lr * self.grad_b

    class OutputLayer:
        def __init__(self, n_inputs, n_classes):
            self.W = np.random.randn(n_classes, n_inputs) * np.sqrt(2.0 / n_inputs)
            self.b = np.zeros(n_classes)

        def forward(self, x):
            self.x = x
            self.z = self.W @ x + self.b
            exp = np.exp(self.z - np.max(self.z))
            self.probs = exp / np.sum(exp)
            return self.probs

        def backward(self, label):
            grad_z = self.probs.copy()
            grad_z[label] -= 1
            self.grad_W = np.outer(grad_z, self.x)
            self.grad_b = grad_z
            return self.W.T @ grad_z

        def update(self, lr):
            self.W -= lr * self.grad_W
            self.b -= lr * self.grad_b

    return HiddenLayer, OutputLayer


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ### Define model
    """)
    return


@app.cell
def _(HiddenLayer, OutputLayer, np):
    def evaluate(layers, output_layer, X, y):
        correct = 0
        for x, label in zip(X, y):
            h = x
            for layer in layers:
                h = layer.forward(h)
            probs = output_layer.forward(h)
            correct += (np.argmax(probs) == label)
        return correct / len(y)

    layer1 = HiddenLayer(784, 128)
    output = OutputLayer(128, 10)

    lr = 0.1
    batch_size = 32
    all_layers = [layer1, output]

    return all_layers, batch_size, evaluate, layer1, lr, output


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ### Train model
    """)
    return


@app.cell
def _(X_train, all_layers, batch_size, evaluate, layer1, lr, np, output, y_train):
    # Split: first 80% train, last 20% validation, shared by both implementations
    n_val = int(0.2 * len(X_train))
    X_val, y_val = X_train[-n_val:], y_train[-n_val:]
    X_train_sub, y_train_sub = X_train[:-n_val], y_train[:-n_val]

    n = X_train_sub.shape[0]
    for epoch in range(5):
        indices = np.random.permutation(n)
        X_shuffled = X_train_sub[indices]
        y_shuffled = y_train_sub[indices]

        epoch_loss = 0.0
        correct = 0

        for i in range(0, n, batch_size):
            X_batch = X_shuffled[i:i+batch_size]
            y_batch = y_shuffled[i:i+batch_size]
            bs = len(X_batch)

            # Accumulate gradients over the batch
            accumulated = {id(l): (np.zeros_like(l.W), np.zeros_like(l.b))
                           for l in all_layers}

            for x, label in zip(X_batch, y_batch):
                h = layer1.forward(x)
                probs = output.forward(h)

                epoch_loss += -np.log(probs[label] + 1e-10)
                correct += (np.argmax(probs) == label)

                grad = output.backward(label)
                layer1.backward(grad)

                for l in all_layers:
                    accumulated[id(l)][0][:] += l.grad_W
                    accumulated[id(l)][1][:] += l.grad_b

            # Average gradients and update
            for l in all_layers:
                l.grad_W = accumulated[id(l)][0] / bs
                l.grad_b = accumulated[id(l)][1] / bs
                l.update(lr)

        train_acc = correct / n
        val_acc = evaluate([layer1], output, X_val, y_val)
        print(f"Epoch {epoch+1}/5 — loss: {epoch_loss/n:.4f}, "
              f"train acc: {train_acc:.3%}, val acc: {val_acc:.3%}")
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ### Evaluate model
    """)
    return


@app.cell
def _(X_test, evaluate, layer1, output, y_test):
    # Final evaluation on the test set (images the model never saw during training or validation)
    numpy_test_acc = evaluate([layer1], output, X_test, y_test)
    print(f"Final test accuracy: {numpy_test_acc:.2%}")
    return


if __name__ == "__main__":
    app.run()