import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms


class SmallResNetBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv1 = nn.Conv2d(channels, channels, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(channels)
        self.conv2 = nn.Conv2d(channels, channels, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(channels)

    def forward(self, x):
        identity = x
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        return F.relu(out + identity)  # residual connection


class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.stem = nn.Conv2d(1, 32, 3, padding=1)
        self.block1 = SmallResNetBlock(32)
        self.pool1 = nn.MaxPool2d(2)   # 28 -> 14
        self.proj = nn.Conv2d(32, 64, 1)
        self.block2 = SmallResNetBlock(64)
        self.pool2 = nn.MaxPool2d(2)   # 14 -> 7
        self.fc = nn.Linear(64 * 7 * 7, num_classes)

    def forward(self, x):
        x = F.relu(self.stem(x))
        x = self.pool1(self.block1(x))
        x = self.proj(x)
        x = self.pool2(self.block2(x))
        x = x.flatten(1)
        return self.fc(x)
