import torch
import torch.nn as nn

# O = (I-K+2P)/S + 1
class Net(torch.nn.Module):
    def __init__(self, num_channels, upscale_factor, d=56, s=12, m=4):
        super(Net, self).__init__()

        self.first_part = nn.Sequential(nn.Conv2d(in_channels=num_channels, out_channels=d, kernel_size=5, stride=1, padding=2),
                                        nn.PReLU())
        self.layers = []
        self.layers.append(nn.Sequential(nn.Conv2d(in_channels=d, out_channels=s, kernel_size=1, stride=1, padding=0),
                                         nn.PReLU()))
        for _ in range(m):
            self.layers.append(nn.Conv2d(in_channels=s, out_channels=s, kernel_size=3, stride=1, padding=1))
        self.layers.append(nn.PReLU())
        self.layers.append(nn.Sequential(nn.Conv2d(in_channels=s, out_channels=d, kernel_size=1, stride=1, padding=0),
                                         nn.PReLU()))

        self.mid_part = torch.nn.Sequential(*self.layers)

        # Deconvolution
        # Change this when dealing with different upscale factor
        if upscale_factor == 2:
            self.last_part = nn.ConvTranspose2d(in_channels=d, out_channels=num_channels, kernel_size=8, stride=upscale_factor, padding=3, output_padding=0)
        elif upscale_factor == 3:
            self.last_part = nn.ConvTranspose2d(in_channels=d, out_channels=num_channels, kernel_size=7, stride=upscale_factor, padding=2, output_padding=0)
        elif upscale_factor == 4:
            self.last_part = nn.ConvTranspose2d(in_channels=d, out_channels=num_channels, kernel_size=8, stride=upscale_factor, padding=2, output_padding=0)

    def forward(self, x):
        out = self.first_part(x)
        out = self.mid_part(out)
        out = self.last_part(out)
        return out

    def weight_init(self, mean=0.0, std=0.02):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                m.weight.data.normal_(mean, std)
                if m.bias is not None:
                    m.bias.data.zero_()
            if isinstance(m, nn.ConvTranspose2d):
                m.weight.data.normal_(0.0, 0.0001)
                if m.bias is not None:
                    m.bias.data.zero_()