import marimo

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


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    # Number of Epochs: Overfitting Check

    Train lr=0.1, bs=32 for 50 epochs to see when validation
    loss stops improving and starts increasing (overfitting).
    """)
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## Load 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
def _(nn, torch):
    def make_model(hidden_sizes=(128,), activation=nn.ReLU):
        torch.manual_seed(42)
        layers = []
        n_inputs = 784
        for size in hidden_sizes:
            layers.extend([nn.Linear(n_inputs, size), activation()])
            n_inputs = size
        layers.append(nn.Linear(n_inputs, 10))
        return nn.Sequential(*layers)
    return (make_model,)


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


@app.cell
def _(X_train, make_model, nn, run_epoch, train_torch, y_train):
    model = make_model()

    # Capture initial loss before any training
    initial = run_epoch(model, X_train[:48000], y_train[:48000])
    initial_val = run_epoch(model, X_train[48000:], y_train[48000:])

    print(f"Before training (random weights) — accuracy: {initial[1]:.4f} — loss: {initial[0]:.4f} — val_accuracy: {initial_val[1]:.4f} — val_loss: {initial_val[0]:.4f}")

    history = train_torch(model, X_train, y_train, epochs=50, batch_size=32,
                        lr=0.1)

    hist = history
    for key, val in [("loss", initial[0]), ("accuracy", initial[1]),
                     ("val_loss", initial_val[0]), ("val_accuracy", initial_val[1])]:
        hist[key] = [val] + hist[key]

    return (hist,)


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## Plot: train vs validation over 50 epochs
    """)
    return


@app.cell
def _(hist, plt):
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

    epochs = range(0, len(hist["loss"]))

    ax1.plot(epochs, hist["accuracy"], label="Train")
    ax1.plot(epochs, hist["val_accuracy"], label="Validation")
    ax1.set_title("Accuracy (lr=0.1, 50 epochs)")
    ax1.set_xlabel("Epoch")
    ax1.set_ylabel("Accuracy")
    ax1.legend()
    ax1.grid(True, alpha=0.3)

    ax2.plot(epochs, hist["loss"], label="Train")
    ax2.plot(epochs, hist["val_loss"], label="Validation")
    ax2.set_title("Loss (lr=0.1, 50 epochs)")
    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()