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

try:
    from huggingface_hub import PyTorchModelHubMixin
except ImportError:  # the model still works without the hub helpers
    class PyTorchModelHubMixin:
        pass


class ResBlock(nn.Module):
    def __init__(self, nf):
        super().__init__()
        self.c1 = nn.Conv2d(nf, nf, 3, 1, 1)
        self.c2 = nn.Conv2d(nf, nf, 3, 1, 1)

    def forward(self, x):
        return x + self.c2(F.relu(self.c1(x), inplace=True))


class SRNet(nn.Module, PyTorchModelHubMixin):
    """x{scale} super-resolution. Input: float RGB (B,3,H,W) in [0,1]. Output: (B,3,H*scale,W*scale), clamp to [0,1]."""

    def __init__(self, scale: int = 4, nf: int = 64, nb: int = 16):
        super().__init__()
        self.scale, self.nf, self.nb = scale, nf, nb
        self.head = nn.Conv2d(3, nf, 3, 1, 1)
        self.body = nn.Sequential(*[ResBlock(nf) for _ in range(nb)])
        self.tail = nn.Conv2d(nf, nf, 3, 1, 1)
        self.up = nn.Conv2d(nf, 3 * scale * scale, 3, 1, 1)
        self.shuffle = nn.PixelShuffle(scale)

    def forward(self, x):
        f = self.head(x)
        f = f + self.tail(self.body(f))
        return self.shuffle(self.up(f)) + F.interpolate(x, scale_factor=self.scale, mode="bicubic", align_corners=False)