import torch
import torch.nn as nn
import torch.nn.functional as F 

# class MNISTNetwork(nn.Module): 
#     # achieved 97 percent accuracy 
#     def __init__(self): 
#         super().__init__()
#         self.layer1 = nn.Linear(784, 400)
#         self.layer2 = nn.Linear(400, 256)
#         self.layer3 = nn.Linear(256, 64)
#         self.layer4 = nn.Linear(64, 32)
#         self.layer5 = nn.Linear(32, 10)
    
#     def forward(self, x):
#         x = x.view(-1, 28*28) 
#         x = torch.relu(self.layer1(x))
#         x = torch.relu(self.layer2(x))
#         x = torch.relu(self.layer3(x))
#         x = torch.relu(self.layer4(x))
#         x = torch.relu(self.layer5(x))
#         return F.log_softmax(x, dim=1)

class MNISTNetwork(nn.Module):
    # achieved 98.783 percent accuracy 
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(64*7*7, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2)
        x = x.view(-1, 64*7*7)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)