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

class BaselineModel(nn.Module):
    """
    Standard MLP architecture using standard PyTorch dense layers.
    This serves as the benchmark for comparison.
    """
    def __init__(self, input_size=784, hidden_size=256, num_classes=10):
        super(BaselineModel, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

class BaselineCIFARModel(nn.Module):
    """
    Standard CNN architecture for CIFAR-10.
    """
    def __init__(self, num_classes=10):
        super(BaselineCIFARModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        # 32x32 -> 16x16 -> 8x8. 64 channels * 8 * 8 = 4096
        self.fc1 = nn.Linear(64 * 8 * 8, 512)
        self.fc2 = nn.Linear(512, num_classes)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 64 * 8 * 8)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x