import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt

def main():
    # get the data
    train_loader, test_loader = download_data()
    print(f"Training examples: {len(train_loader.dataset)}")
    print(f"Test examples: {len(test_loader.dataset)}")

    # a very simple and fast nn
    model = nn.Sequential(
        nn.Flatten(),
        nn.Linear(28*28, 1024),
        nn.Sigmoid(),
        nn.Linear(1024, 1024),
        nn.Sigmoid(),
        nn.Linear(1024, 10)
    )
    
    # define the loss and optimizer ONCE outside run_train
    loss_function = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # see how bad it is before training
    accuracy = run_test(test_loader, model)
    print(f"Untrained accuracy: {accuracy:.2f}%")

    # train the model
    print("Training ...")
    for epoch in range(1, 11):  # fewer epochs just to demo
        run_train(train_loader, model, loss_function, optimizer)
        accuracy = run_test(test_loader, model)
        print(f"Epoch {epoch}, Test Accuracy: {accuracy:.2f}%")

def run_train(train_loader, model, loss_function, optimizer):
    model.train()
    for batch_idx, (image, truth) in enumerate(train_loader):
        # 1. Zero the parameter gradients
        optimizer.zero_grad()
        # 2. Forward pass
        pred = model(image)
        # 3. Calculate the loss
        loss = loss_function(pred, truth)
        # 4. Backprop
        loss.backward()
        # 5. Update params
        optimizer.step()

def run_test(test_loader, model):
    model.eval()
    correct = 0
    total = 0

    # No gradient needed for inference
    with torch.no_grad():
        for image, truth in test_loader:
            # get the probability it is each digit
            predicted_probabilities = model(image)
            # choose the digit with the highest probability
            predicted_label = predicted_probabilities.argmax(dim=1)
            # count how many we got correct
            correct += (predicted_label == truth).sum().item()
            total += truth.size(0)

    return 100.0 * correct / total

def download_data():
    training_data = datasets.MNIST(
        root=".",
        train=True,
        download=True,
        transform=ToTensor()
    )
    # optional: comment this out if you don't want to block on the plot each time
    show_images(training_data)

    test_data = datasets.MNIST(
        root=".",
        train=False,
        download=True,
        transform=ToTensor()
    )

    train_loader = DataLoader(training_data, batch_size=64, shuffle=True)
    test_loader = DataLoader(test_data, batch_size=64, shuffle=False)

    return train_loader, test_loader

def show_images(training_data):
    figure = plt.figure(figsize=(8, 8))
    cols, rows = 5, 5

    for i in range(1, cols * rows + 1):
        sample_idx = torch.randint(len(training_data), size=(1,)).item()
        img, label = training_data[sample_idx]
        figure.add_subplot(rows, cols, i)
        plt.axis("off")
        plt.imshow(img.squeeze(), cmap="gray")
        plt.title(label)
    plt.show()

if __name__ == "__main__":
    main()