import torch
import torch.nn as nn
import torch.nn.functional as F


class FSRCNN(nn.Module):
    def __init__(self, scale, num_channels=3, d=56, s=12, m=4):
        super().__init__()
        self.feature_extraction = nn.Sequential(
            nn.Conv2d(num_channels, d, kernel_size=5, padding=2), nn.PReLU(d))
        self.shrink = nn.Sequential(nn.Conv2d(d, s, kernel_size=1), nn.PReLU(s))
        mapping_layers = []
        for _ in range(m):
            mapping_layers += [nn.Conv2d(s, s, kernel_size=3, padding=1), nn.PReLU(s)]
        self.mapping = nn.Sequential(*mapping_layers)
        self.expand = nn.Sequential(nn.Conv2d(s, d, kernel_size=1), nn.PReLU(d))
        self.upsample = nn.Sequential(
            nn.Conv2d(d, num_channels * scale * scale, kernel_size=3, padding=1),
            nn.PixelShuffle(scale))
        self.scale = scale

    def forward(self, x):
        x = self.feature_extraction(x)
        x = self.shrink(x)
        x = self.mapping(x)
        x = self.expand(x)
        x = self.upsample(x)
        return torch.sigmoid(x)