import torch.nn as nn

from layers_custom import maskConv0, MaskConvBlock

import torch

class MaskCNN(nn.Module):

    def __init__(self, n_channel=1024, h=128):

        """PixelCNN Model"""

        super(MaskCNN, self).__init__()

        self.MaskConv0 = maskConv0(n_channel, h, k_size=7, stride=1, pad=3)

        # large 7 x 7 masked filter with image downshift to ensure that each output neuron's receptive field only sees what is above it in the image 

        MaskConv = []

        # stack of 10 gated residual masked conv blocks

        for i in range(10):

            MaskConv.append(MaskConvBlock(h, k_size=3, stride=1, pad=1))

        self.MaskConv = nn.Sequential(*MaskConv)

        # 1x1 conv to upsample to required feature (channel) length

        self.out = nn.Sequential(

            nn.ReLU(),

            nn.Conv2d(h, n_channel, kernel_size=1, stride=1, padding=0),

            nn.BatchNorm2d(n_channel),

            nn.ReLU()

            )

    def forward(self, x):

        """

        Args:

            x: [batch_size, channel, height, width]

        Return:

            out [batch_size, channel, height, width]

        """

        # fully convolutional, feature map dimension maintained constant throughout

        x = self.MaskConv0(x)

        x = self.MaskConv(x)

        x = self.out(x)

        return x

if __name__ == '__main__':

    from torchsummary import summary

    model = PixelCNN(1024, 128)

    summary(model, (1024, 7,7))

    x = torch.rand(2, 1024, 7, 7)

    x = model(x)

    print(x.shape)