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

OUT_CHANNELS = 32

class N902(nn.Module):

    # 32, 144.878
    # 64, 135.952
    # 128, 128.388

    def __init__(self):
        super(N90, self).__init__()
        # 1 input image channel, 6 output channels, 3x3 square convolution
        # kernel
        self.conv1 = nn.Conv2d(1           , OUT_CHANNELS, 3, padding=1)
        self.conv2 = nn.Conv2d(OUT_CHANNELS, 2           , 3, padding=1)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.conv2(x)
        return x


class N903(nn.Module):

    # 32, 79.591
    # 64, 69.663

    def __init__(self):
        super(N90, self).__init__()
        # 1 input image channel, 6 output channels, 3x3 square convolution
        # kernel
        self.conv1 = nn.Conv2d(1           , OUT_CHANNELS, 3, padding=1)
        self.conv2 = nn.Conv2d(OUT_CHANNELS, OUT_CHANNELS, 3, padding=1)
        self.conv3 = nn.Conv2d(OUT_CHANNELS, 2           , 3, padding=1)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = self.conv3(x)
        return x

class N904(nn.Module):

    # 32, 65.503
    # 64, 55.369
    
    def __init__(self):
        super(N90, self).__init__()
        # 1 input image channel, 6 output channels, 3x3 square convolution
        # kernel
        self.conv1 = nn.Conv2d(1           , OUT_CHANNELS, 3, padding=1)
        self.conv2 = nn.Conv2d(OUT_CHANNELS, OUT_CHANNELS, 3, padding=1)
        self.conv3 = nn.Conv2d(OUT_CHANNELS, OUT_CHANNELS, 3, padding=1)
        self.conv4 = nn.Conv2d(OUT_CHANNELS, 2           , 3, padding=1)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = self.conv4(x)
        return x

class N90(nn.Module): 
    # 32, 48.523

    def __init__(self):
        super(N90, self).__init__()
        # 1 input image channel, 6 output channels, 3x3 square convolution
        # kernel
        self.conv1 = nn.Conv2d(1           , OUT_CHANNELS, 3, padding=1)
        self.conv2 = nn.Conv2d(OUT_CHANNELS, OUT_CHANNELS, 3, padding=1)
        self.conv3 = nn.Conv2d(OUT_CHANNELS, OUT_CHANNELS, 3, padding=1)
        self.conv4 = nn.Conv2d(OUT_CHANNELS, OUT_CHANNELS, 3, padding=1)
        self.conv5 = nn.Conv2d(OUT_CHANNELS, 2           , 3, padding=1)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = F.relu(self.conv4(x))
        x = self.conv5(x)
        return x
    

class N906(nn.Module): 
    # 32, 43.330

    def __init__(self):
        super(N90, self).__init__()
        # 1 input image channel, 6 output channels, 3x3 square convolution
        # kernel
        self.conv1 = nn.Conv2d(1           , OUT_CHANNELS, 3, padding=1)
        self.conv2 = nn.Conv2d(OUT_CHANNELS, OUT_CHANNELS, 3, padding=1)
        self.conv3 = nn.Conv2d(OUT_CHANNELS, OUT_CHANNELS, 3, padding=1)
        self.conv4 = nn.Conv2d(OUT_CHANNELS, OUT_CHANNELS, 3, padding=1)
        self.conv5 = nn.Conv2d(OUT_CHANNELS, OUT_CHANNELS, 3, padding=1)
        self.conv6 = nn.Conv2d(OUT_CHANNELS, 2           , 3, padding=1)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = F.relu(self.conv4(x))
        x = F.relu(self.conv5(x))
        x = self.conv6(x)
        return x


# net = N90_100()
# print(net)