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'))