# coding:utf-8
import torch.nn as nn
import torch.nn.functional as F


class Conv2D(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, padding='same',
                 stride=1, dilation=1, groups=1):
        super(Conv2D, self).__init__()

        assert type(kernel_size) in [int, tuple], "Allowed kernel type [int or tuple], not {}".format(type(kernel_size))
        assert padding == 'same', "Allowed padding type {}, not {}".format('same', padding)

        self.kernel_size = kernel_size
        if isinstance(kernel_size, tuple):
            self.h_kernel = kernel_size[0]
            self.w_kernel = kernel_size[1]
        else:
            self.h_kernel = kernel_size
            self.w_kernel = kernel_size

        self.padding = padding
        self.stride = stride
        self.dilation = dilation
        self.groups = groups
        self.conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size,
                              stride=self.stride, dilation=self.dilation, groups=self.groups)

    def forward(self, x):

        if self.padding == 'same':

            height, width = x.shape[2:]

            h_pad_need = max(0, (height - 1) * self.stride + self.h_kernel - height)
            w_pad_need = max(0, (width - 1) * self.stride + self.w_kernel - width)

            pad_left = w_pad_need // 2
            pad_right = w_pad_need - pad_left
            pad_top = h_pad_need // 2
            pad_bottom = h_pad_need - pad_top

            padding = (pad_left, pad_right, pad_top, pad_bottom)

            x = F.pad(x, padding, 'constant', 0)

        x = self.conv(x)

        return x


class GCN(nn.Module):
    """
        Large Kernel Matters -- https://arxiv.org/abs/1703.02719
    """
    def __init__(self, in_channels, out_channels, k=3):
        super(GCN, self).__init__()

        self.conv_l1 = Conv2D(in_channels=in_channels, out_channels=out_channels, kernel_size=(k, 1), padding='same')
        self.conv_l2 = Conv2D(in_channels=out_channels, out_channels=out_channels, kernel_size=(1, k), padding='same')

        self.conv_r1 = Conv2D(in_channels=in_channels, out_channels=out_channels, kernel_size=(1, k), padding='same')
        self.conv_r2 = Conv2D(in_channels=out_channels, out_channels=out_channels, kernel_size=(k, 1), padding='same')

    def forward(self, x):
        x1 = self.conv_l1(x)
        x1 = self.conv_l2(x1)

        x2 = self.conv_r1(x)
        x2 = self.conv_r2(x2)

        out = x1 + x2

        return out
