# -*- coding:utf-8 -*-
import torch.nn as nn
import torch


class ConvUnit(nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size, stride=1, bn=False):
        super(ConvUnit, self).__init__()
        padding = (kernel_size - 1) // 2
        self.features = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding=padding),
            nn.BatchNorm2d(out_channels, eps=0.001, momentum=0, affine=True),
            nn.ReLU(inplace=True),
        ) if bn else nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding=padding),
            nn.ReLU(inplace=True),
        )

    def forward(self, x):
        x = self.features(x)

        return x


class MCNN(nn.Module):
    """Refer from `"MCNN..." <https://www.semanticscholar.org/paper/Single-Image-Crowd-Counting-via-Multi-Column-Neural-Zhang-Zhou/2dc3b3eff8ded8914c8b536d05ee713ff0cdf3cd>`_ paper.

    Args:
        bn (bool): if True, Batch Normalization layer will be used after every convolutional layer

    """
    def __init__(self, bn=False):
        super(MCNN, self).__init__()
        print("*****init MCNN net*****")
        self.column1 = nn.Sequential(
            ConvUnit(3, 16, 9, bn=bn),
            nn.MaxPool2d(2),
            ConvUnit(16, 32, 7, bn=bn),
            nn.MaxPool2d(2),
            ConvUnit(32, 16, 7, bn=bn),
            ConvUnit(16, 8, 7, bn=bn),
        )

        self.column2 = nn.Sequential(
            ConvUnit(3, 20, 7, bn=bn),
            nn.MaxPool2d(2),
            ConvUnit(20, 40, 5, bn=bn),
            nn.MaxPool2d(2),
            ConvUnit(40, 20, 5, bn=bn),
            ConvUnit(20, 10, 5, bn=bn),
        )

        self.column3 = nn.Sequential(
            ConvUnit(3, 24, 5, bn=bn),
            nn.MaxPool2d(2),
            ConvUnit(24, 48, 3, bn=bn),
            nn.MaxPool2d(2),
            ConvUnit(48, 24, 3, bn=bn),
            ConvUnit(24, 12, 3, bn=bn),
        )

        self.merge = nn.Sequential(
            ConvUnit(30, 1, 1, bn=bn)
        )

    def forward(self, x):
        x1 = self.column1(x)
        x2 = self.column2(x)
        x3 = self.column3(x)
        x = torch.cat((x1, x2, x3), dim=1)
        x = self.merge(x)

        return x

