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

class BaselineCNN(nn.Module):
  def __init__(self, num_classes=39):
    super(BaselineCNN, self).__init__()
    

    self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=1)
    self.bn1 = nn.BatchNorm2d(32)
    
    self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
    self.bn2 = nn.BatchNorm2d(64)

    self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
    self.bn3   = nn.BatchNorm2d(128)

    self.pool = nn.MaxPool2d(2, 2)

    self.fc = nn.Linear(128 * 32 * 32, num_classes)

  def forward(self, x):

    x = self.pool(F.relu(self.bn1(self.conv1(x))))
    x = self.pool(F.relu(self.bn2(self.conv2(x))))
    x = self.pool(F.relu(self.bn3(self.conv3(x))))
    x = torch.flatten(x, 1)
    x = self.fc(x)
    return x