import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import Compose, Lambda, ToTensor

torch.manual_seed(10)

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
print(device)

trainset = datasets.FashionMNIST('data', train=True, download=True,
                          transform=Compose([ToTensor(), Lambda(lambda x: x.flatten())]))
testset = datasets.FashionMNIST('data', train=False, download=True,
                         transform=Compose([ToTensor(), Lambda(lambda x: x.flatten())]))


def train(model, dataset, n_iter=10, batch_size=256):
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    criterion = nn.NLLLoss()
    dl = DataLoader(dataset, batch_size=batch_size)
    model.train()
    for epoch in range(n_iter):
        for images, targets in dl:
            optimizer.zero_grad()
            out = model(images.to(device))
            loss = criterion(out, targets.to(device))
            loss.backward()
            optimizer.step()
        if epoch % 10 == 0:
            print('epoch: %3d loss: %.4f' % (epoch, loss))


def accuracy(model, dataset):
    model.eval()
    correct = sum([(model(images.to(device)).argmax(dim=1) == targets.to(device)).sum()
                   for images, targets in DataLoader(dataset, batch_size=256)])
    print(correct.float() / len(dataset))


neurons = 300

model = nn.Sequential(
    nn.Linear(28 * 28, neurons),
    nn.ReLU(),
    nn.Linear(neurons, 10),
    nn.LogSoftmax(dim=-1)
).to(device)

train(model, trainset)

accuracy(model, testset)