import torch
import torch.nn


class MyCNN(torch.nn.Module):
    def __init__(self,
                 input_channels: int):
        super().__init__()

        self.layers = torch.nn.Sequential(
            torch.nn.Conv2d(in_channels=input_channels, out_channels=32, kernel_size=3, padding='same', bias=False),
            torch.nn.BatchNorm2d(32),
            torch.nn.ReLU(),
            torch.nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding='same', bias=False),
            torch.nn.BatchNorm2d(32),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=3, padding=1),

            torch.nn.Dropout2d(0.1),
            torch.nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding='same', bias=False),
            torch.nn.BatchNorm2d(64),
            torch.nn.ReLU(),
            torch.nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding='same', bias=False),
            torch.nn.BatchNorm2d(64),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=3, padding=1),

            torch.nn.Dropout2d(0.1),
            torch.nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding='same', bias=False),
            torch.nn.BatchNorm2d(128),
            torch.nn.ReLU(),
            torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, padding='same', bias=False),
            torch.nn.BatchNorm2d(128),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=3, padding=1),

            torch.nn.Flatten(),
            torch.nn.Dropout(0.25),
            torch.nn.Linear(in_features=2048, out_features=1024),
            torch.nn.ReLU(),
            torch.nn.Linear(in_features=1024, out_features=512),
            torch.nn.ReLU(),
            torch.nn.Linear(in_features=512, out_features=20, bias=False)
        )

    def forward(self, input_images: torch.Tensor) -> torch.Tensor:
        return self.layers(input_images)

    def __repr__(self):
        return str(self.layers)


model = MyCNN(input_channels=1)