import torch

import torch.nn as nn

import pdb

def down_shift(x, pad=None):

    # Pytorch ordering

    xs = [int(y) for y in x.size()]

    # when downshifting, the last row is removed 

    x = x[:, :, :xs[2] - 1, :]

    # padding left, padding right, padding top, padding bottom

    pad = nn.ZeroPad2d((0, 0, 1, 0)) if pad is None else pad

    return pad(x)

class MaskedConv2d(nn.Conv2d):

    def __init__(self, c_in, c_out, k_size, stride, pad, use_down_shift=False):

        super(MaskedConv2d, self).__init__(

            c_in, c_out, k_size, stride, pad, bias=False)

        ch_out, ch_in, height, width = self.weight.size()

                # Mask

        #         -------------------------------------

        #        |  1       1       1       1       1 |

        #        |  1       1       1       1       1 |

        #        |  1       1       1       1       1 |   H // 2

        #        |  0       0       0       0       0 |   H // 2 + 1

        #        |  0       0       0       0       0 |

        #         -------------------------------------

        #  index    0       1     W//2    W//2+1

        mask = torch.ones(ch_out, ch_in, height, width)

        mask[:, :, height // 2 + 1:] = 0

        self.register_buffer('mask', mask)

        self.use_down_shift = use_down_shift

    def forward(self, x):

        self.weight.data *= self.mask

        if self.use_down_shift:

            x = down_shift(x)

        return super(MaskedConv2d, self).forward(x)

def maskConv0(c_in=3, c_out=256, k_size=7, stride=1, pad=3):

    """2D Masked Convolution first layer"""

    return nn.Sequential(

        MaskedConv2d(c_in, c_out * 2, k_size, stride, pad, use_down_shift=True),

        nn.BatchNorm2d(c_out * 2),

        Gate()

        )

class Gate(nn.Module):

    def __init__(self):

        super(Gate, self).__init__()

    def forward(self, x):

        # gated activation 

        xf, xg = torch.chunk(x, 2, dim=1)

        f = torch.tanh(xf)

        g = torch.sigmoid(xg)

        return f * g

class MaskConvBlock(nn.Module):

    def __init__(self, h=128, k_size=3, stride=1, pad=1):

        """1x1 Conv + 2D Masked Convolution (type B) + 1x1 Conv"""

        super(MaskConvBlock, self).__init__()

        self.net = nn.Sequential(

            MaskedConv2d(h, 2 * h, k_size, stride, pad),

            nn.BatchNorm2d(2 * h),

            Gate()

        )

    def forward(self, x):

        """Residual connection"""

        return self.net(x) + x

if __name__ == '__main__':

    def conv(x, kernel):

        return nn.functional.conv2d(x, kernel, padding=1)

    x = torch.ones((1, 1, 5, 5)) * 0.1

    x[:,:,1,0] = 1000

    print("blindspot experiment")

    normal_kernel = torch.ones(1, 1, 3, 3)

    mask_kernel = torch.zeros(1, 1, 3, 3)

    mask_kernel[:,:,0,:] = 1

    mask_b = mask_kernel.clone()

    mask_b[:,:,1,1] = 1

    # mask_kernel[:,:,1,1] = 1

    print("unmasked kernel:", "\n",normal_kernel.squeeze(), "\n")

    print("masked kernel:", "\n", mask_kernel.squeeze(), "\n")

    print("normal conv")

    print("orig image", "\n", x.squeeze(), "\n")

    y = conv(x, normal_kernel)

    print(y[:,0, :,:], "\n")

    y = conv(y, normal_kernel)

    print(y[:,0, :,:], "\n")

    print("with mask")

    print("orig image", "\n", x.squeeze(), "\n")

    y = conv(x, mask_kernel)

    print(y[:,0, :,:], "\n")

    y = conv(y, mask_b)

    print(y[:,0, :,:], "\n")

    y = conv(y, mask_b)

    print(y[:,0, :,:],"\n")

    print("with down_shift")

    print("orig image", x.squeeze(), "\n")

    c_kernel = mask_kernel

    c_kernel[:,:,1,:] = 1

    print("custom kernel:", "\n", c_kernel.squeeze(), "\n")

    y = conv(down_shift(x), c_kernel)

    print(y[:,0, :,:],"\n")

    y = conv(y, c_kernel)

    print(y[:,0, :,:],"\n")

    y = conv(y, c_kernel)

    print(y[:,0, :,:],"\n")

    y = conv(y, c_kernel)

    print(y[:,0, :,:],"\n")