import torch
import torch.nn as nn
import torch.nn.init as init
import math


def pixact(x):
    return (torch.tanh(x) + 1) / 2
    #return x.sigmoid()

# class Net(nn.Module):
#     def __init__(self, upscale_factor):
#         super(Net, self).__init__()
#
#         self.relu = nn.ReLU()
#         self.conv1 = nn.Conv2d(3, 64, (5, 5), (1, 1), (2, 2))
#         self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
#         self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
#         self.conv4 = nn.Conv2d(32, upscale_factor ** 2, (3, 3), (1, 1), (1, 1))
#         self.pixel_shuffle = nn.PixelShuffle(upscale_factor)
#
#         self._initialize_weights()
#
#     def forward(self, x):
#         x = self.relu(self.conv1(x))
#         x = self.relu(self.conv2(x))
#         x = self.relu(self.conv3(x))
#         x = self.pixel_shuffle(self.conv4(x))
#         return x
#
#     def _initialize_weights(self):
#         init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
#         init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
#         init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
#         init.orthogonal_(self.conv4.weight)

# class Net(nn.Module):
#     def __init__(self, upscale_factor=1):
#         super(Net, self).__init__()
#
#         self.relu = nn.ReLU()
#         self.conv1 = nn.Conv2d(3, 64, (5, 5), (1, 1), (2, 2))
#         self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
#         self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
#         self.conv4 = nn.Conv2d(32, 3, (3, 3), (1, 1), (1, 1))
#         #self.pixel_shuffle = nn.PixelShuffle(upscale_factor)
#         self.upsample = nn.Upsample(scale_factor=2, mode='nearest')
#
#         self._initialize_weights()
#
#     def forward(self, x):
#         x = self.relu(self.conv1(x))
#         x = self.relu(self.conv2(x))
#         x = self.upsample(self.relu(self.conv3(x)))
#         x = self.relu(self.conv4(x))
#         return x
#
#     def _initialize_weights(self):
#         init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
#         init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
#         init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
#         init.orthogonal_(self.conv4.weight)

# class Net(nn.Module):
#     def __init__(self, upscale_factor=1):
#         super(Net, self).__init__()
#
#         self.relu = nn.ReLU()
#         self.conv1 = nn.Conv2d(3, 64, (5, 5), (1, 1), (2, 2))
#         self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
#         self.conv3 = nn.Conv2d(64, 128, (3, 3), (1, 1), (1, 1))
#         self.conv4 = nn.Conv2d(128, 64, (3, 3), (1, 1), (1, 1))
#         self.conv5 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
#         self.conv6 = nn.Conv2d(32, 3, (3, 3), (1, 1), (1, 1))
#         #self.pixel_shuffle = nn.PixelShuffle(upscale_factor)
#         self.upsample = nn.Upsample(scale_factor=2, mode='nearest')
#
#         self._initialize_weights()
#
#
#     def forward(self, x):
#         x = self.relu(self.conv1(x))
#         x = self.relu(self.conv2(x))
#         x = self.upsample(self.relu(self.conv3(x)))
#         x = self.relu(self.conv4(x))
#         x = self.upsample(self.relu(self.conv5(x)))
#         x = self.relu(self.conv6(x))
#         return x
#
#     def _initialize_weights(self):
#         init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
#         init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
#         init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
#         init.orthogonal_(self.conv4.weight, init.calculate_gain('relu'))
#         init.orthogonal_(self.conv5.weight, init.calculate_gain('relu'))
#         init.orthogonal_(self.conv6.weight, init.calculate_gain('relu'))

class ResidualBlockG(nn.Module):
  def __init__(self, channels):
    super(ResidualBlockG, self).__init__()
    self.conv1 = dws_block(channels, channels, kernel_size=3, padding=1)
    self.bn1 = nn.InstanceNorm2d(channels)
    self.relu = nn.ReLU()
    self.conv2 = dws_block(channels, channels, kernel_size=3, padding=1)
    self.bn2 = nn.InstanceNorm2d(channels)
  def forward(self, x):
    residual = self.conv1(x)
    residual = self.bn1(residual)
    residual = self.relu(residual)
    residual = self.conv2(residual)
    residual = self.bn2(residual)
    return x + residual

class UpsampleBlock(nn.Module):
  def __init__(self, in_channels, up_scale):
    super(UpsampleBlock, self).__init__()
    self.conv = nn.Conv2d(in_channels, in_channels * up_scale ** 2,
                          kernel_size=3, padding=1)
    self.pixel_shuffle = nn.PixelShuffle(up_scale)
    self.prelu = nn.ReLU()
  def forward(self, x):
    x = self.conv(x)
    x = self.pixel_shuffle(x)
    x = self.prelu(x)
    return x


class Generator(nn.Module):
    def __init__(self, scale_factor):
        super(Generator, self).__init__()
        upsample_block_num = int(math.log(scale_factor, 2))

        self.block1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=9, padding=4),
            nn.ReLU()
        )

        self.block2 = ResidualBlockG(64)
        self.block3 = ResidualBlockG(64)
        self.block4 = ResidualBlockG(64)
        self.block5 = ResidualBlockG(64)
        self.block6 = ResidualBlockG(64)
        self.block7 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64)
        )
        block8 = [UpsampleBlock(64, 2) for _ in range(upsample_block_num)]
        block8.append(nn.Conv2d(64, 3, kernel_size=9, padding=4))
        self.block8 = nn.Sequential(*block8)

    def forward(self, x):
        block1 = self.block1(x)
        block2 = self.block2(block1)
        block3 = self.block3(block2)
        block4 = self.block4(block3)
        block5 = self.block5(block4)
        block6 = self.block6(block5)
        block7 = self.block7(block6)
        block8 = self.block8(block1 + block7)
        return block8

class block(nn.Module):
    def __init__(self,channels_in,channels_out,kernel,stride,pad):
        super(block, self).__init__()
        self.conv = nn.Conv2d(channels_in, channels_out, kernel, stride, pad)
        self.act = nn.LeakyReLU(0.2)
        #self.act = nn.ReLU()
        #self.norm = nn.InstanceNorm2d(channels_out)

    def forward(self, x):
        x = self.conv(x)
        #x = self.norm(x)
        x = self.act(x)

        return x

class SubPixelConvolutionalBlock(nn.Module):
    """
    A subpixel convolutional block, comprising convolutional, pixel-shuffle, and PReLU activation layers.
    """

    def __init__(self, kernel_size=3, n_channels=64, scaling_factor=2):
        """
        :param kernel_size: kernel size of the convolution
        :param n_channels: number of input and output channels
        :param scaling_factor: factor to scale input images by (along both dimensions)
        """
        super(SubPixelConvolutionalBlock, self).__init__()

        # A convolutional layer that increases the number of channels by scaling factor^2, followed by pixel shuffle and PReLU
        self.conv = nn.Conv2d(in_channels=n_channels, out_channels=n_channels * (scaling_factor ** 2),
                              kernel_size=kernel_size, padding=kernel_size // 2)
        # These additional channels are shuffled to form additional pixels, upscaling each dimension by the scaling factor
        self.pixel_shuffle = nn.PixelShuffle(upscale_factor=scaling_factor)
        #self.lrelu = nn.ReLU()
        self.lrelu = nn.LeakyReLU(0.2)

    def forward(self, input):
        """
        Forward propagation.

        :param input: input images, a tensor of size (N, n_channels, w, h)
        :return: scaled output images, a tensor of size (N, n_channels, w * scaling factor, h * scaling factor)
        """
        output = self.conv(input)  # (N, n_channels * scaling factor^2, w, h)
        output = self.pixel_shuffle(output)  # (N, n_channels, w * scaling factor, h * scaling factor)
        output = self.lrelu(output)  # (N, n_channels, w * scaling factor, h * scaling factor)

        return output

class Net(nn.Module):
    def __init__(self, upscale_factor=1):
        super(Net, self).__init__()

        self.conv1 = block(3, 64, (7, 7), (1, 1), (3, 3))
        self.conv2 = block(64, 64, (5, 5), (1, 1), (2, 2))
        self.conv3 = block(64, 128, (5, 5), (1, 1), (2, 2))
        self.conv4 = block(128, 256, (5, 5), (1, 1), (2, 2))
        self.conv4_1 = block(256, 256, (5, 5), (1, 1), (2, 2))
        self.conv5 = block(256, 128, (5, 5), (1, 1), (2, 2))
        self.conv6 = block(128, 64, (5, 5), (1, 1), (2, 2))
        self.conv7 = block(64, 32, (3, 3), (1, 1), (1, 1))
        self.conv7_1 = block(32, 16, (3, 3), (1, 1), (1, 1))
        self.conv8 = nn.Conv2d(16, 3, (3, 3), (1, 1), (1, 1))
        #self.pixel_shuffle = nn.PixelShuffle(upscale_factor)
        #self.relu = nn.ReLU(inplace=True)
        #self.upsample = nn.Upsample(scale_factor=2, mode='bicubic')
        self.spc1 = SubPixelConvolutionalBlock( kernel_size=3, n_channels=128, scaling_factor=2)
        #self.spc1 = nn.Upsample(scale_factor=2, mode='bicubic')
        self.spc2 = SubPixelConvolutionalBlock( kernel_size=3, n_channels=128, scaling_factor=2)
        #self.spc2 = nn.Upsample(scale_factor=2, mode='bicubic')
        self.spc3 = SubPixelConvolutionalBlock(kernel_size=3, n_channels=128, scaling_factor=4)

        #self._initialize_weights()

    def forward(self, x):
        x =self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.spc1(x)
        x = self.conv4(x)
        x = self.conv4_1(x)
        x = self.conv5(x)
        x = self.spc2(x)
        x = self.conv6(x)
        x = self.conv7(x)
        x = self.conv7_1(x)
        x = self.conv8(x)
        return x

    # def _initialize_weights(self):
    #     init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
    #     init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
    #     init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
    #     init.orthogonal_(self.conv4.weight, init.calculate_gain('relu'))
    #     init.orthogonal_(self.conv4_1.weight, init.calculate_gain('relu'))
    #     init.orthogonal_(self.conv5.weight, init.calculate_gain('relu'))
    #     init.orthogonal_(self.conv6.weight, init.calculate_gain('relu'))
    #     init.orthogonal_(self.conv7.weight, init.calculate_gain('relu'))
    #     init.orthogonal_(self.conv7_1.weight, init.calculate_gain('relu'))
    #     init.orthogonal_(self.conv8.weight, init.calculate_gain('relu'))

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                # n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                # m.weight.data.normal_(0, math.sqrt(2. / n))
                # g = nn.init.calculate_gain('leaky_relu', 0.1)
                m.weight.data.normal_(0, 0.01)
                # nn.init.xavier_normal_(m.weight,gai=g)
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.InstanceNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                m.weight.data.normal_(0, 0.01)
                m.bias.data.zero_()
        # self.sp.weight.data.normal_(0,0.01)


class dws_block(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, padding=1,stride = 1):
        super(dws_block, self).__init__()
        self.dc = nn.Conv2d(in_channels, in_channels, kernel_size=kernel_size, padding=padding, groups=in_channels,stride=stride)
        self.pc = nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0)

    def forward(self, x):
        # Encoder
        x = self.dc(x)
        x = self.pc(x)

        return x


class _Residual_Block(nn.Module):
    def __init__(self):
        super(_Residual_Block, self).__init__()

        self.conv1 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn1 = nn.InstanceNorm2d(64)
        self.relu = nn.LeakyReLU(0.2, inplace=True)
        self.pad2 = nn.ReflectionPad2d(1)
        self.conv2 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.InstanceNorm2d(64)

    def forward(self, x):
        identity_data = x
        output = self.relu(self.bn1(self.conv1(x)))
        output = self.bn2(self.conv2(x))
        # add output of the ResBlock with its input with a skip connection
        output = torch.add(output, identity_data)
        return output


class UpscaleNet(nn.Module):
    def __init__(self):
        super(UpscaleNet, self).__init__()
        # Input is 3 Channels => to 64 channels (kernel size 9 stride 1)

        self.conv_input = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=9, stride=1, padding=4, bias=False)
        self.relu = nn.LeakyReLU(0.2, inplace=True)

        # Running 16 resBlocks
        self.residual = self.make_layer(_Residual_Block, 16)

        # Position of the Network wise skip connection(transfering spatial information from the low dimension image)
        self.conv_mid = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn_mid = nn.InstanceNorm2d(64)

        # Upscale module with PixelShuffle (Fractional Convolution)
        self.upscale4x = nn.Sequential(
            # using x2 two times to upscale by 4 times
            nn.Conv2d(in_channels=64, out_channels=256, kernel_size=3, stride=1, padding=1, bias=False),
            nn.PixelShuffle(2),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(in_channels=64, out_channels=256, kernel_size=3, stride=1, padding=1, bias=False),
            nn.PixelShuffle(2),
            nn.LeakyReLU(0.2, inplace=True),
        )

        # Final convolution to aleviate artifacts from the upsampling and colapsing channels from 64 to 3

        self.conv_output = nn.Conv2d(in_channels=64, out_channels=3, kernel_size=9, stride=1, padding=4, bias=False)

        for m in self.modules():
            # Initialisation of conv layers
            if isinstance(m, nn.Conv2d):
                # init.orthogonal(m.weight, math.sqrt(2))
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
                if m.bias is not None:
                    m.bias.data.zero_()

    def make_layer(self, block, num_of_layer):
        # Generate the ResBlocks
        layers = []
        for _ in range(num_of_layer):
            layers.append(block())
        return nn.Sequential(*layers)

    def forward(self, x):
        # FeedForward
        out = self.relu(self.conv_input(x))
        # save input conv for skip connection
        residual = out
        # Apply ResBlocks
        out = self.residual(out)
        out = self.bn_mid(self.conv_mid(out))
        # Apply skip connection
        out = torch.add(out, residual)
        # Upsclae
        out = self.upscale4x(out)
        # Generate output
        out = self.conv_output(out)
        return out

class UNet(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(UNet, self).__init__()

        # Encoder (contracting path)
        self.encoder1 = self.contracting_block(in_channels, 64)
        self.encoder2 = self.contracting_block(64, 128)
        self.encoder3 = self.contracting_block(128, 256)
        self.encoder4 = self.contracting_block(256, 512)

        # Bottleneck
        self.bottleneck = nn.Sequential(
            dws_block(512, 1024, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            dws_block(1024, 1024, kernel_size=3, padding=1),
            nn.ReLU(inplace=True)
        )

        # Decoder (expansive path)
        self.decoder1 = self.expansive_block(1024, 512)
        self.decoder2 = self.expansive_block(512, 256)
        self.decoder3 = self.expansive_block(256, 128)
        self.decoder4 = self.expansive_block(128, 64)
        self.decoder5 = self.expansive_block(32, 32)
        self.decoder6 = self.expansive_block(16, 32)

        # Output layer
        self.final_conv = nn.Conv2d(16, out_channels, kernel_size=1)

    def contracting_block(self, in_channels, out_channels):
        return nn.Sequential(
            dws_block(in_channels, out_channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            dws_block(out_channels, out_channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

    def expansive_block(self, in_channels, out_channels):
        return nn.Sequential(
            dws_block(in_channels, out_channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            dws_block(out_channels, out_channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(out_channels, out_channels // 2, kernel_size=2, stride=2)
        )

    def forward(self, x):
        # Encoder
        enc1 = self.encoder1(x)
        enc2 = self.encoder2(enc1)
        enc3 = self.encoder3(enc2)
        enc4 = self.encoder4(enc3)

        # Bottleneck
        bottleneck = self.bottleneck(enc4)

        # Decoder
        dec1 = self.decoder1(bottleneck)
        dec2 = self.decoder2(torch.cat([dec1, enc3], dim=1))
        dec3 = self.decoder3(torch.cat([dec2, enc2], dim=1))
        dec4 = self.decoder4(torch.cat([dec3, enc1], dim=1))
        dec5 = self.decoder5(dec4)
        dec6 = self.decoder6(dec5)
        # Output layer
        output = self.final_conv(dec6)
        return output

class ConvolutionalBlock(nn.Module):
    """
    A convolutional block, comprising convolutional, BN, activation layers.
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride=1, batch_norm=False, activation=None):
        """
        :param in_channels: number of input channels
        :param out_channels: number of output channe;s
        :param kernel_size: kernel size
        :param stride: stride
        :param batch_norm: include a BN layer?
        :param activation: Type of activation; None if none
        """
        super(ConvolutionalBlock, self).__init__()

        if activation is not None:
            activation = activation.lower()
            assert activation in {'prelu', 'leakyrelu', 'tanh'}

        # A container that will hold the layers in this convolutional block
        layers = list()

        # A convolutional layer
        layers.append(
            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride,
                      padding=kernel_size // 2))

        # A batch normalization (BN) layer, if wanted
        if batch_norm is True:
            layers.append(nn.BatchNorm2d(num_features=out_channels))

        # An activation layer, if wanted
        if activation == 'prelu':
            layers.append(nn.PReLU())
        elif activation == 'leakyrelu':
            layers.append(nn.LeakyReLU(0.2))
        elif activation == 'tanh':
            layers.append(nn.Tanh())

        # Put together the convolutional block as a sequence of the layers in this container
        self.conv_block = nn.Sequential(*layers)

    def forward(self, input):
        """
        Forward propagation.

        :param input: input images, a tensor of size (N, in_channels, w, h)
        :return: output images, a tensor of size (N, out_channels, w, h)
        """
        output = self.conv_block(input)  # (N, out_channels, w, h)

        return output


class ResidualBlock(nn.Module):
    """
    A residual block, comprising two convolutional blocks with a residual connection across them.
    """

    def __init__(self, kernel_size=3, n_channels=64):
        """
        :param kernel_size: kernel size
        :param n_channels: number of input and output channels (same because the input must be added to the output)
        """
        super(ResidualBlock, self).__init__()

        # The first convolutional block
        self.conv_block1 = ConvolutionalBlock(in_channels=n_channels, out_channels=n_channels, kernel_size=kernel_size,
                                              batch_norm=True, activation='PReLu')

        # The second convolutional block
        self.conv_block2 = ConvolutionalBlock(in_channels=n_channels, out_channels=n_channels, kernel_size=kernel_size,
                                              batch_norm=True, activation=None)

    def forward(self, input):
        """
        Forward propagation.

        :param input: input images, a tensor of size (N, n_channels, w, h)
        :return: output images, a tensor of size (N, n_channels, w, h)
        """
        residual = input  # (N, n_channels, w, h)
        output = self.conv_block1(input)  # (N, n_channels, w, h)
        output = self.conv_block2(output)  # (N, n_channels, w, h)
        output = output + residual  # (N, n_channels, w, h)

        return output


class SRResNet(nn.Module):
    """
    The SRResNet, as defined in the paper.
    """

    def __init__(self, large_kernel_size=9, small_kernel_size=3, n_channels=64, n_blocks=16, scaling_factor=4):
        """
        :param large_kernel_size: kernel size of the first and last convolutions which transform the inputs and outputs
        :param small_kernel_size: kernel size of all convolutions in-between, i.e. those in the residual and subpixel convolutional blocks
        :param n_channels: number of channels in-between, i.e. the input and output channels for the residual and subpixel convolutional blocks
        :param n_blocks: number of residual blocks
        :param scaling_factor: factor to scale input images by (along both dimensions) in the subpixel convolutional block
        """
        super(SRResNet, self).__init__()

        # Scaling factor must be 2, 4, or 8
        scaling_factor = int(scaling_factor)
        assert scaling_factor in {2, 4, 8}, "The scaling factor must be 2, 4, or 8!"

        # The first convolutional block
        self.conv_block1 = ConvolutionalBlock(in_channels=3, out_channels=n_channels, kernel_size=large_kernel_size,
                                              batch_norm=False, activation='PReLu')

        # A sequence of n_blocks residual blocks, each containing a skip-connection across the block
        self.residual_blocks = nn.Sequential(
            *[ResidualBlock(kernel_size=small_kernel_size, n_channels=n_channels) for i in range(n_blocks)])

        # Another convolutional block
        self.conv_block2 = ConvolutionalBlock(in_channels=n_channels, out_channels=n_channels,
                                              kernel_size=small_kernel_size,
                                              batch_norm=True, activation=None)

        # Upscaling is done by sub-pixel convolution, with each such block upscaling by a factor of 2
        n_subpixel_convolution_blocks = int(math.log2(scaling_factor))
        self.subpixel_convolutional_blocks = nn.Sequential(
            *[SubPixelConvolutionalBlock(kernel_size=small_kernel_size, n_channels=n_channels, scaling_factor=2) for i
              in range(n_subpixel_convolution_blocks)])

        # The last convolutional block
        self.conv_block3 = ConvolutionalBlock(in_channels=n_channels, out_channels=3, kernel_size=large_kernel_size,
                                              batch_norm=False, activation='Tanh')

    def forward(self, lr_imgs):
        """
        Forward prop.

        :param lr_imgs: low-resolution input images, a tensor of size (N, 3, w, h)
        :return: super-resolution output images, a tensor of size (N, 3, w * scaling factor, h * scaling factor)
        """
        output = self.conv_block1(lr_imgs)  # (N, 3, w, h)
        residual = output  # (N, n_channels, w, h)
        output = self.residual_blocks(output)  # (N, n_channels, w, h)
        output = self.conv_block2(output)  # (N, n_channels, w, h)
        output = output + residual  # (N, n_channels, w, h)
        output = self.subpixel_convolutional_blocks(output)  # (N, n_channels, w * scaling factor, h * scaling factor)
        sr_imgs = self.conv_block3(output)  # (N, 3, w * scaling factor, h * scaling factor)

        return sr_imgs

class Discriminator(nn.Module):
  def __init__(self):
    super(Discriminator, self).__init__()
    self.net = nn.Sequential(
        dws_block(3, 64, kernel_size=3, padding=1),
        nn.LeakyReLU(0.2),

        dws_block(64, 64, kernel_size=3, stride=2, padding=1),
        nn.BatchNorm2d(64),
        nn.LeakyReLU(0.2),

        dws_block(64, 128, kernel_size=3, padding=1),
        nn.BatchNorm2d(128),
        nn.LeakyReLU(0.2),

        dws_block(128, 256, kernel_size=3, padding=1),
        nn.BatchNorm2d(256),
        nn.LeakyReLU(0.2),

        dws_block(256, 256, kernel_size=3, stride=2, padding=1),
        nn.BatchNorm2d(256),
        nn.LeakyReLU(0.2),

        dws_block(256, 512, kernel_size=3, padding=1),
        nn.BatchNorm2d(512),
        nn.LeakyReLU(0.2),

        dws_block(512, 512, kernel_size=3, stride=2, padding=1),
        nn.BatchNorm2d(512),
        nn.LeakyReLU(0.2),

        nn.AdaptiveAvgPool2d(1),
        nn.Conv2d(512, 1024, kernel_size=1),
        nn.LeakyReLU(0.2),
        nn.Conv2d(1024, 1, kernel_size=1)
    )
  def forward(self, x):
    batch_size=x.size()[0]
    return torch.sigmoid(self.net(x).view(batch_size))