import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary


class DigitalLetterNet(nn.Module):
    def __init__(self, in_chs, out_chs, fcn=False):
        super(DigitalLetterNet, self).__init__()
        self.in_chs = in_chs
        self.out_chs = out_chs
        self.fcn = fcn

        self.conv = nn.Sequential(
            nn.Conv2d(self.in_chs, 16, 5, padding=2),
            nn.ReLU(True),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1),
            nn.ReLU(True),
            nn.BatchNorm2d(32),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(True),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 3, padding=1),
            nn.ReLU(True),
            nn.MaxPool2d(2),
            nn.Conv2d(128, 256, 3, padding=1),
            nn.ReLU(True),
            nn.MaxPool2d(2),
        )
        if self.fcn:
            self.classifier = nn.Conv2d(256, self.out_chs, 1)
        else:
            self.classifier = nn.Linear(256, self.out_chs)

    def forward(self, x):
        x = self.conv(x)
        if not self.fcn:
            x.squeeze_()
        x = self.classifier(x)
        if self.fcn:
            x.squeeze_()
        x = F.log_softmax(x, dim=-1)
        return x


# just for test
if __name__ == '__main__':
    dl_net = DigitalLetterNet(3, 36, True)
    print(summary(dl_net, (3, 32, 32), device='cpu'))