import torch 
import torchvision
import torchvision.transforms as transforms



transform = transforms.Compose([transforms.ToTensor()])
trainset = torchvision.datasets.MNIST(root='.data/', train=True, download=True, transform=transform)
train_data = torch.utils.data.DataLoader(trainset, shuffle=True, batch_size=4,num_workers=2)
testset = torchvision.datasets.MNIST(root='.data/', train=False,download=True, transform=transform  )
test_data = torch.utils.data.DataLoader(testset, shuffle=True, batch_size=4)



class SimpleConv(torch.nn.Module):
    def __init__(self):
    super(SimpleConv, self).__init__()
    self.conv1 = torch.nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5)
    self.pool = torch.nn.MaxPool2d(kernel_size=2, stride=2)


    #return 28/2 = 14*14
    self.layer1 = torch.nn.Sequential(
    # 1 - count of input map 
    # 32 - count of output map
    torch.nn.Conv2d(1,32, kernel_size=5, stride=1),
    torch.nn.ReLU(),
    torch.nn.MaxPool2d(kernel_size=2, stride=2)
    )

    #return 14/2 = 7*7
    self.layer2 = torch.nn.Sequential(
    torch.nn.Conv2d(32,64,kernel_size=5,stride=1,padding=2),
    torch.nn.ReLU(),
    torch.nn.MaxPool2d(kernel_size=2,stride=2)
    )

    self.drop_out = torch.nn.Dropout() 


    self.fc1 = torch.nn.Linear(7 * 7 * 64, 1000) 
    self.fc2 = torch.nn.Linear(1000, 10)


def forward(self, x):
    out = self.layer1(x) 
    out = self.layer2(out) 
    print(out.shape)
    out = x.view(-1, 7 * 7 * 64)
    print(out.shape)
    out = self.drop_out(out) 
    out = self.fc1(out) 
    out = self.fc2(out) 
    return out



 model      = SimpleConv()
 loss_fun = torch.nn.CrossEntropyLoss()
 optim    = torch.optim.Adam(model.parameters(),lr= 0.001)



num_epochs = 5 
num_classes = 10 



for epoch in range(num_epochs):
   for i, batch in enumerate(train_data): 
          X_batch, y_batch = batch

          print(X_batch.shape)
          print(y_batch.shape)

          optim.zero_grad()
          output = model(X_batch)



          loss = loss_fun(output, y_batch)