import marimo

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


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    # CNN on MNIST

    Train a convolutional neural network on MNIST digits.
    """)
    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()
    X_train = X_train.reshape(-1, 1, 28, 28)
    X_test = X_test.reshape(-1, 1, 28, 28)
    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"""
    ## Build and train
    """)
    return


@app.cell
def _(X_train, nn, torch, train_torch, y_train):
    model = nn.Sequential(
        nn.Conv2d(1, 32, kernel_size=3), nn.ReLU(),  # (N, 1, 28, 28) → (N, 32, 26, 26)
        nn.MaxPool2d(2),                            # → (N, 32, 13, 13)
        nn.Conv2d(32, 64, kernel_size=3), nn.ReLU(), # → (N, 64, 11, 11)
        nn.MaxPool2d(2),                            # → (N, 64, 5, 5)
        nn.Flatten(),                              # → (N, 1600)
        nn.Linear(1600, 128), nn.ReLU(),
        nn.Linear(128, 10),                        # logits
    )
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    history = train_torch(model, X_train, y_train, epochs=5, val_fraction=0.1, optimizer=optimizer)
    return history, model


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


@app.cell
def _(X_test, history, model, plt, run_epoch, y_test):
    _, test_acc = run_epoch(model, X_test, y_test)
    print(f"Test accuracy: {test_acc:.4f}")

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))

    ax1.plot(history["accuracy"], label="Train", linestyle="--", color="#4a6cf7")
    ax1.plot(history["val_accuracy"], label="Validation", color="#4a6cf7")
    ax1.set_title("Accuracy")
    ax1.set_xlabel("Epoch")
    ax1.set_ylabel("Accuracy")
    ax1.legend()
    ax1.grid(True, alpha=0.3)

    ax2.plot(history["loss"], label="Train", linestyle="--", color="#d44")
    ax2.plot(history["val_loss"], label="Validation", color="#d44")
    ax2.set_title("Loss")
    ax2.set_xlabel("Epoch")
    ax2.set_ylabel("Loss")
    ax2.legend()
    ax2.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.show()
    return


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