import torch.nn as nn

class MainNet(nn.Module):
    def __init__(self, z_size=100, out_size=3, ngf=128):
        super(MainNet, self).__init__()
        self.z_size = z_size
        self.ngf = ngf
        self.out_size = out_size

        self.layer0 = nn.Sequential(
            # input size is z_size
            nn.ConvTranspose2d(self.z_size, self.ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(self.ngf * 8),
            nn.ReLU(inplace=True)
            # state size: (ngf * 8) x 4 x 4
            )

        self.layer1 = nn.Sequential(
            # state size: (ngf * 8) x 4 x 4
            nn.ConvTranspose2d(self.ngf * 8, self.ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(self.ngf * 4),
            nn.ReLU(inplace=True),
            # state size: (ngf * 4) x 8 x 8
            )
        self.layer2 = nn.Sequential(
            # state size: (ngf * 4) x 8 x 8
            nn.ConvTranspose2d(self.ngf * 4, self.ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(self.ngf * 2),
            nn.ReLU(inplace=True),
            # state size: (ngf * 2) x 16 x 16
            )
        self.layer3 = nn.Sequential(
            # state size: (ngf * 2) x 16 x 16
            nn.ConvTranspose2d(self.ngf * 2, self.ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(self.ngf),
            nn.ReLU(inplace=True),
            # state size: ngf x 32 x 32
            )
        self.layer4 = nn.Sequential(
            # state size: ngf x 32 x 32
            nn.ConvTranspose2d(self.ngf, self.out_size, 4, 2, 1, bias=False),
            nn.Tanh()
            # state size: out_size x 64 x 64
            )
        for m in self.modules():
            if isinstance(m, nn.ConvTranspose2d):
                m.weight.data.normal_(0.0, 0.02)
                if m.bias is not None:
                    m.bias.data.zero_()

    def forward(self, input):
        self.l1out = output = self.layer0(input)
        self.l2out = output = self.layer1(output)
        self.l3out = output = self.layer2(output)
        self.l4out = output = self.layer3(output)
        output = self.layer4(output)
        return output