import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms


class DNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(1 * 28 * 28, 128),
            nn.ReLU(),
            nn.Linear(128, 10),
        )

    def forward(self, x):
        x = x.view(x.shape[0], -1)
        x = self.model(x)
        return x


def train(device, model, train_loader, optimizer):
    sum_loss = 0

    model.train()
    for x, y in train_loader:
        x = x.to(device)
        y = y.to(device)

        optimizer.zero_grad()
        yh = model(x)
        loss = F.cross_entropy(yh, y)
        loss.backward()
        optimizer.step()

        sum_loss += loss.item()

    avg_loss = sum_loss / len(train_loader)
    return avg_loss


def test(device, model, test_loader):
    sum_loss = 0
    total = 0
    correct = 0

    with torch.no_grad():
        model.eval()
        for x, y in test_loader:
            x = x.to(device)
            y = y.to(device)

            yh = model(x)
            loss = F.cross_entropy(yh, y)

            sum_loss += loss.item()
            total += x.shape[0]
            correct += (torch.argmax(yh, 1) == y).int().sum().item()

    avg_loss = sum_loss / len(test_loader)
    accuracy = correct / total
    return avg_loss, accuracy


def main():
    use_cuda = True
    max_epoch = 10
    learning_rate = 0.0002
    train_kwargs = {
        "batch_size": 32,
        "shuffle": True,
        "drop_last": True,
    }
    test_kwargs = {
        "batch_size": 32,
        "shuffle": True,
        "drop_last": True,
    }

    if use_cuda and torch.cuda.is_available():
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")

    transform = transforms.ToTensor()
    train_dataset = datasets.MNIST(
        "./data", train=True, download=True, transform=transform
    )
    train_loader = torch.utils.data.DataLoader(train_dataset, **train_kwargs)

    test_dataset = datasets.MNIST(
        "./data", train=False, download=True, transform=transform
    )
    test_loader = torch.utils.data.DataLoader(test_dataset, **test_kwargs)

    model = DNN().to(device)
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    for epoch in range(max_epoch):
        train_loss = train(device, model, train_loader, optimizer)
        test_loss, test_accuracy = test(device, model, test_loader)
        print(
            "epoch",
            epoch,
            "train_loss",
            train_loss,
            "test_loss",
            test_loss,
            "test_accuracy",
            test_accuracy,
        )


if __name__ == "__main__":
    main()
