import cv2 as cv
import os
import torch.nn as nn
import torch
import torchvision
import tarfile
import torch.functional as f
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torch.optim
from torchvision.datasets.utils import download_url
from torch.utils.data import DataLoader as DataLoader
from google.colab.patches import cv2_imshow
from zipfile import ZipFile
import glob
from google.colab import files
import imageio
from albumentations import HorizontalFlip, VerticalFlip, Rotate
import matplotlib.pyplot as plt

batch_size=64
input_size  = 1
output_classes  = 10
lr  = 0.001
num_epochs=5

class   cnn(nn.Module):
  def __init__(self,input_size, output_classes):
    super().__init__()
    self.conv1  = nn.Sequential(
       nn.Conv2d(input_size,4,kernel_size=(3,3), padding=1, stride=1),
       nn.ReLU(),
       nn.MaxPool2d(2),
       nn.Conv2d(4,8,kernel_size=(3,3), padding=1, stride=1),
       nn.ReLU(),
       nn.MaxPool2d(2))
    self.linear2  = nn.Linear(392,output_classes)   
  def forward(self, input):    
    x   = self.conv1(input)
    x = x.reshape(x.shape[0],-1)
    x=self.linear2(x)
    return x 

device  = 'cuda' if torch.cuda.is_available() else 'cpu'
print(device)
model = cnn(1,10).to(device)

train_ds  = datasets.MNIST(root='/dataset', train=True, transform = transforms.ToTensor(), download = True)
test_ds  = datasets.MNIST(root='/dataset', train=False, transform = transforms.ToTensor(), download = True)
train_dl  = DataLoader(train_ds,batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)
test_dl   = DataLoader(test_ds, batch_size=batch_size, shuffle = False, num_workers=4,pin_memory=True)